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 Polar-Charged-Hydrophobic classification
 * Symbol groupls: STQN, KRH, DE, LAMIVFWY, C, P, G
 *
 * polar uncharged:             STQN
 * polar positively chared:     KRH
 * polar negatively charged:    DE
 * hydrophobic:                 LAMIVFWY
 * non-polar small:             P
 * non-polar tiny:              G
 * ambiguous:                   C
 *
 */
public class ChargedProteinLogoStyle implements SymbolStyle {
    private ArrayList<Color> colorList;

    public ChargedProteinLogoStyle() {
        AminoAcidGrouping aaGroup = new AminoAcidGrouping();
        aaGroup.setColoring(AminoAcidGrouping.COLORING_POLAR_CHARGED);
        this.colorList = aaGroup.getColorList();
    }

    public ChargedProteinLogoStyle(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;

        //polar uncharged (S,T,Q,N)
        if (s.equals(ProteinTools.s()) || s.equals(ProteinTools.t()) ||
                s.equals(ProteinTools.q()) || s.equals(ProteinTools.n())) {
            fillColor = colorList.get(0);

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

        //polar negative (D,E)
        } else if (s.equals(ProteinTools.d()) || s.equals(ProteinTools.e())) {
            fillColor = colorList.get(2);

        //non-polar (PGC)
        } else if (s.equals(ProteinTools.p()) || s.equals(ProteinTools.g()) || s.equals(ProteinTools.c())) {
            fillColor = colorList.get(3);

        //hydrophobic (LAMIVFWY)
        } else {
            fillColor = colorList.get(4);
        }
        return fillColor;
    }
}
