package org.baderlab.brain;

import org.biojava.bio.gui.SymbolStyle;
import org.biojava.bio.symbol.Symbol;
import org.biojava.bio.symbol.IllegalSymbolException;
import org.biojava.bio.seq.ProteinTools;

import java.awt.*;
import java.util.ArrayList;

/**
 * Paints colors for the sequence logo according to the Bydrophobic-by-Size classification
 * Symbol groups: ST, KRH, DEQN, G, P, VAC, ILM, FWY
 *
 * ST       hydrophobic-polar-small     light gray
 * KRH      hydrophilic-polar+-large    dark blue
 * DEQN     hydrophilic-polar-small     blue
 * G        hydrophilic-nonpolar-tiny   sky blue
 * P        ?-nonpolar-small            light blue
 * VAC      hydrophobic-nonpolar-small  gray
 * ILM      hydrophobic-nonpolar-medium dark gray
 * FWY      hydrophobic-large           black
 *
 */
public class HydrophobicBySizeProteinLogoStyle implements SymbolStyle {
    private ArrayList<Color> colorList;

    public HydrophobicBySizeProteinLogoStyle() {
        AminoAcidGrouping aaGroup = new AminoAcidGrouping();
        aaGroup.setColoring(AminoAcidGrouping.COLORING_HYDROPHOBIC);
        this.colorList = aaGroup.getColorList();
    }

    public HydrophobicBySizeProteinLogoStyle(ArrayList<Color> colorList) {
        this.colorList = colorList;
    }

    public Paint outlinePaint(Symbol s) throws IllegalSymbolException {
        return fillPaint(s);
    }

    public Paint fillPaint(Symbol s) throws IllegalSymbolException {
        Paint fillColor = null;

        //(S,T)
        if (s.equals(ProteinTools.s()) || s.equals(ProteinTools.t())) {
            fillColor = colorList.get(0);

        //(K,R,H)
        } else if (s.equals(ProteinTools.k()) || s.equals(ProteinTools.r()) || s.equals(ProteinTools.h())) {
            fillColor = colorList.get(1);

        //(D,E,Q,N)
        } else if (s.equals(ProteinTools.d()) || s.equals(ProteinTools.e())
                || s.equals(ProteinTools.q()) || s.equals(ProteinTools.n())) {
            fillColor = colorList.get(2);

        //(G)
        } else if (s.equals(ProteinTools.g())) {
            fillColor = colorList.get(3);

        //(P)
        } else if (s.equals(ProteinTools.p())) {
            fillColor = colorList.get(4);

        //(VAC)
        } else if (s.equals(ProteinTools.v()) || s.equals(ProteinTools.a())
                || s.equals(ProteinTools.c())) {
            fillColor = colorList.get(5);

        //(ILM)
        } else if (s.equals(ProteinTools.i()) || s.equals(ProteinTools.l())
                || s.equals(ProteinTools.m())) {
            fillColor = colorList.get(6);

        //(FWY)
        } else if (s.equals(ProteinTools.f()) || s.equals(ProteinTools.w())
                || s.equals(ProteinTools.y())) {
            fillColor = colorList.get(7);
        }

        return fillColor;
    }
}
