package org.baderlab.brain;

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

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

/**
 * Copyright (c) 2005 Memorial Sloan-Kettering Cancer Center
 * *
 * * Code written by: Gary Bader
 * * Authors: Gary Bader, Chris Sander
 * *
 * * This library is free software; you can redistribute it and/or modify it
 * * under the terms of the GNU Lesser General Public License as published
 * * by the Free Software Foundation; either version 2.1 of the License, or
 * * any later version.
 * *
 * * This library is distributed in the hope that it will be useful, but
 * * WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF
 * * MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.  The software and
 * * documentation provided hereunder is on an "as is" basis, and
 * * Memorial Sloan-Kettering Cancer Center
 * * has no obligations to provide maintenance, support,
 * * updates, enhancements or modifications.  In no event shall the
 * * Memorial Sloan-Kettering Cancer Center
 * * be liable to any party for direct, indirect, special,
 * * incidental or consequential damages, including lost profits, arising
 * * out of the use of this software and its documentation, even if
 * * Memorial Sloan-Kettering Cancer Center
 * * has been advised of the possibility of such damage.  See
 * * the GNU Lesser General Public License for more details.
 * *
 * * You should have received a copy of the GNU Lesser General Public License
 * * along with this library; if not, write to the Free Software Foundation,
 * * Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA.
 * *
 * * User: GaryBader
 * * Date: Feb 20, 2006
 * * Time: 7:23:44 PM
 */

/**
 * Defines a number of amino acid groupings
 */
public class AminoAcidGrouping {
    public final static String GROUPING_NO_GROUPING =   "No Grouping";
    public final static String GROUPING_POLAR_CHARGED = "Polar-Hydrophobic";
    public final static String GROUPING_HYDROPHOBIC =   "Hydrophobic-Size";
    public final static String GROUPING_HYDROXYL =      "Hydroxyl+";
    public final static String GROUPING_PDZ =           "PDZ";
    public final static String GROUPING_WEBLOGO =       "WebLogo";
    public final static String GROUPING_DEFAULT = GROUPING_NO_GROUPING;

    public final static String COLORING_POLAR_CHARGED = "Polar-Hydrophobic";
    public final static String COLORING_HYDROPHOBIC =   "Hydrophobic-Size";
    public final static String COLORING_HYDROXYL =      "Hydroxyl+";
    public final static String COLORING_PDZ =           "PDZ";
    public final static String COLORING_WEBLOGO =       "WebLogo";
    public final static String COLORING_DEFAULT = COLORING_WEBLOGO;

    public final static String DESCRIPTION_NO_GROUPING =    "No amino acid grouping applied";
    public final static String DESCRIPTION_POLAR_CHARGED =  "Grouped by polarity and charge; hydrophobic residues are grouped separately";
    public final static String DESCRIPTION_HYDROPHOBIC =    "Grouped by hydrophobicity and size";
    public final static String DESCRIPTION_HYDROXYL =       "Grouped by common physico-chemical properties";
    public final static String DESCRIPTION_WEBLOGO =        "WebLogo default grouping scheme";
    public final static String DESCRIPTION_PDZ =            "Grouping based on PDZ 2007 meeting";

    /*
 * polar uncharged:             STQN
 * polar positively chared:     KRH
 * polar negatively charged:    DE
 * hydrophobic:                 LAMIVFWY
 * non-polar small:             P
 * non-polar tiny:              G
 * ambiguous:                   C
     */

    //amino acid grouping methods
    public final static String[] GROUPING_METHODS = {
            GROUPING_NO_GROUPING,
            GROUPING_POLAR_CHARGED,
            GROUPING_HYDROPHOBIC,
            GROUPING_HYDROXYL,
            GROUPING_PDZ,
            GROUPING_WEBLOGO
    };

    //amino acid coloring methods
    public final static String[] COLORING_METHODS = {
            COLORING_POLAR_CHARGED,
            COLORING_HYDROPHOBIC,
            COLORING_HYDROXYL,
            COLORING_PDZ,
            COLORING_WEBLOGO
    };

    private final static int[][] RGB_POLAR_CHARGED =  { {255, 4, 255}, {204, 2, 4}, {4, 2, 204}, {91, 91, 91}, {102, 0, 102}, {255, 170, 0}, {102, 51, 0}, {0, 0, 0} };
    private final static int[][] RGB_HYDROPHOBIC =    { {200, 200, 200}, {4, 2, 255}, {4, 2, 204}, {50, 50, 255}, {60, 80, 255}, {150, 150, 150}, {91, 91, 91}, {0, 0, 0} };
    private final static int[][] RGB_HYDROXYL =       { {255, 4, 255}, {4, 2, 204}, {204, 2, 4}, {91, 91, 91}, {0, 0, 0} };
    private final static int[][] RGB_PDZ =            { {4, 206, 4}, {204, 2, 4}, {4, 2, 204}, {102, 0, 102}, {255, 170, 0}, {102, 51, 0}, {0, 0, 0} };
    private final static int[][] RGB_WEBLOGO =        { {4, 206, 4}, {204, 2, 4}, {4, 2, 204}, {0, 0, 0} };

    private final static String[] GROUPS_POLAR_CHARGED = { "STQN", "KRH", "DE", "C", "P", "G", "LAMIVFWY" };
    private final static String[] GROUPS_HYDROPHOBIC =   { "ST", "KRH", "DEQN", "G", "P", "VAC", "ILM", "FWY" };
    private final static String[] GROUPS_HYDROXYL =      { "ST", "KRH", "DEQN", "G", "FLAMPWIVCY" };
    private final static String[] GROUPS_PDZ =           { "STQN", "DE", "KRH", "G", "C", "P", "AVLIMFYW" };
    private final static String[] GROUPS_WEBLOGO =       { "GSTYCQN", "DE", "KRH", "AVLIPWFM" };

    private final static String[] COLORS_POLAR_CHARGED = GROUPS_POLAR_CHARGED;
    private final static String[] COLORS_HYDROPHOBIC =   GROUPS_HYDROPHOBIC;
    private final static String[] COLORS_HYDROXYL =      GROUPS_HYDROXYL;
    private final static String[] COLORS_PDZ =           GROUPS_PDZ;
    private final static String[] COLORS_WEBLOGO =       GROUPS_WEBLOGO;

    private final static String[] COLOR_CATEGORIES_POLAR_CHARGED =    { "Polar Neutral", "Polar Positive", "Polar Negative", "Non-Polar Tiny", "Non-Polar Small", "Non-Polar Small", "Hydrophobic" };
    private final static String[] COLOR_CATEGORIES_HYDROPHOBIC =      { "Hydrophobic-Small", "Hydrophilic-Large", "Hydrophilic-Small", "Hydrophilic-Tiny", "Small", "Hydrophobic-Small (Polar)", "Hydrophobic-Medium (Non-polar)", "Hydrophobic-Large" };
    private final static String[] COLOR_CATEGORIES_HYDROXYL =         { "Hydroxyl-Polar", "Basic", "Polar-Acidic", "Hydrophilic-Nonpolar", "Hydrophobic" };
    private final static String[] COLOR_CATEGORIES_PDZ =              { "Polar", "Acidic", "Basic", "Tiny", "Small", "Small", "Hydrophobic"};
    private final static String[] COLOR_CATEGORIES_WEBLOGO =          { "Polar", "Acidic", "Basic", "Hydrophobic"};

    private final static String[] GROUP_CATEGORIES_POLAR_CHARGED =    COLOR_CATEGORIES_POLAR_CHARGED;
    private final static String[] GROUP_CATEGORIES_HYDROPHOBIC =      COLOR_CATEGORIES_HYDROPHOBIC;
    private final static String[] GROUP_CATEGORIES_HYDROXYL =         COLOR_CATEGORIES_HYDROXYL;
    private final static String[] GROUP_CATEGORIES_PDZ =              COLOR_CATEGORIES_PDZ;
    private final static String[] GROUP_CATEGORIES_WEBLOGO =          COLOR_CATEGORIES_WEBLOGO;

    private String selectedGrouping;
    private String selectedColoring;
    private String selectedDescription;
    private String selectedGroups;
    private ArrayList<Color> colorList;
    private ArrayList<String> colorSymbolList;
    private ArrayList<String> colorCategoryList;
    private ArrayList<String> groupList;
    private ArrayList<String> groupCategoryList;
    private SymbolStyle proteinLogoStyle;
    private int numberOfGroups = 0;
    private int numberOfColors = 0;


    //constructor
    public AminoAcidGrouping() {
        colorList = new ArrayList<Color>();
        colorSymbolList = new ArrayList<String>();
        colorCategoryList = new ArrayList<String>();
        groupList = new ArrayList<String>();
        groupCategoryList = new ArrayList<String>();
        this.setGrouping(GROUPING_DEFAULT);
        this.setColoring(COLORING_DEFAULT);
    }

    public void setColoring(String colorType) {
        selectedColoring = colorType;
        colorList.clear();
        colorSymbolList.clear();
        colorCategoryList.clear();

        if (colorType == COLORING_POLAR_CHARGED) {
            proteinLogoStyle = new ChargedProteinLogoStyle(colorList);
            numberOfColors = COLORS_POLAR_CHARGED.length;
            selectedDescription = DESCRIPTION_POLAR_CHARGED;
            for (int i = 0; i < numberOfColors; i++) {
                colorList.add(new Color(RGB_POLAR_CHARGED[i][0], RGB_POLAR_CHARGED[i][1], RGB_POLAR_CHARGED[i][2]));
                colorSymbolList.add(COLORS_POLAR_CHARGED[i]);
                colorCategoryList.add(COLOR_CATEGORIES_POLAR_CHARGED[i]);
            }
        }
        else if (colorType == COLORING_HYDROPHOBIC) {
            proteinLogoStyle = new HydrophobicBySizeProteinLogoStyle(colorList);
            numberOfColors = COLORS_HYDROPHOBIC.length;
            selectedDescription = DESCRIPTION_HYDROPHOBIC;
            for (int i = 0; i < numberOfColors; i++) {
                colorList.add(new Color(RGB_HYDROPHOBIC[i][0], RGB_HYDROPHOBIC[i][1], RGB_HYDROPHOBIC[i][2]));
                colorSymbolList.add(COLORS_HYDROPHOBIC[i]);
                colorCategoryList.add(COLOR_CATEGORIES_HYDROPHOBIC[i]);
            }
        }
        else if (colorType == COLORING_HYDROXYL) {
            proteinLogoStyle = new HydroxylProteinLogoStyle(colorList);
            numberOfColors = COLORS_HYDROXYL.length;
            selectedDescription = DESCRIPTION_HYDROXYL;
            for (int i = 0; i < numberOfColors; i++) {
                colorList.add(new Color(RGB_HYDROXYL[i][0], RGB_HYDROXYL[i][1], RGB_HYDROXYL[i][2]));
                colorSymbolList.add(COLORS_HYDROXYL[i]);
                colorCategoryList.add(COLOR_CATEGORIES_HYDROXYL[i]);
            }
        }
        else if (colorType == COLORING_PDZ) {
            proteinLogoStyle = new PDZProteinLogoStyle(colorList);
            numberOfColors = COLORS_PDZ.length;
            selectedDescription = DESCRIPTION_PDZ;
            for (int i = 0; i < numberOfColors; i++) {
                colorList.add(new Color(RGB_PDZ[i][0], RGB_PDZ[i][1], RGB_PDZ[i][2]));
                colorSymbolList.add(COLORS_PDZ[i]);
                colorCategoryList.add(COLOR_CATEGORIES_PDZ[i]);
            }
        }
        else if (colorType == COLORING_WEBLOGO) {
            proteinLogoStyle = new WebLogoProteinStyle(this);
            numberOfColors = COLORS_WEBLOGO.length;
            selectedDescription = DESCRIPTION_WEBLOGO;
            for (int i = 0; i < numberOfColors; i++) {
                colorList.add(new Color(RGB_WEBLOGO[i][0], RGB_WEBLOGO[i][1], RGB_WEBLOGO[i][2]));
                colorSymbolList.add(COLORS_WEBLOGO[i]);
                colorCategoryList.add(COLOR_CATEGORIES_WEBLOGO[i]);
            }
        }
        else {
            throw new IllegalStateException("'" + colorType + "' is not a recognized color scheme.");
        }
    }

    public void setGrouping(String groupType) {
        //set selected grouping, logo style, colour list, and group list
        selectedGrouping = groupType;
        groupList.clear();
        groupCategoryList.clear();
        if (groupType == GROUPING_NO_GROUPING) {
            selectedDescription = DESCRIPTION_NO_GROUPING;
            selectedGroups = null;
            numberOfGroups = 0;
        }
        else if (groupType == GROUPING_POLAR_CHARGED) {
            selectedDescription = DESCRIPTION_POLAR_CHARGED;
            selectedGroups = Arrays.toString(GROUPS_POLAR_CHARGED);
            numberOfGroups = GROUPS_POLAR_CHARGED.length;
            for (int i = 0; i < numberOfGroups; i++) {
                groupList.add(GROUPS_POLAR_CHARGED[i]);
                groupCategoryList.add(GROUP_CATEGORIES_POLAR_CHARGED[i]);
            }
        }
        else if (groupType == GROUPING_HYDROPHOBIC) {
            selectedDescription = DESCRIPTION_HYDROPHOBIC;
            selectedGroups = Arrays.toString(GROUPS_HYDROPHOBIC);
            numberOfGroups = GROUPS_HYDROPHOBIC.length;
            for (int i = 0; i < numberOfGroups; i++) {
                groupList.add(GROUPS_HYDROPHOBIC[i]);
                groupCategoryList.add(GROUP_CATEGORIES_HYDROPHOBIC[i]);
            }
        }
        else if (groupType == GROUPING_HYDROXYL) {
            selectedDescription = DESCRIPTION_HYDROXYL;
            selectedGroups = Arrays.toString(GROUPS_HYDROXYL);
            numberOfGroups = GROUPS_HYDROXYL.length;
            for (int i = 0; i < numberOfGroups; i++) {
                groupList.add(GROUPS_HYDROXYL[i]);
                groupCategoryList.add(GROUP_CATEGORIES_HYDROXYL[i]);
            }
        }
        else if (groupType == GROUPING_PDZ) {
            selectedDescription = DESCRIPTION_PDZ;
            selectedGroups = Arrays.toString(GROUPS_PDZ);
            numberOfGroups = GROUPS_PDZ.length;
            for (int i = 0; i < numberOfGroups; i++) {
                groupList.add(GROUPS_PDZ[i]);
                groupCategoryList.add(GROUP_CATEGORIES_PDZ[i]);
            }
        }
        else if (groupType == GROUPING_WEBLOGO) {
            selectedDescription = DESCRIPTION_WEBLOGO;
            selectedGroups = Arrays.toString(GROUPS_WEBLOGO);
            numberOfGroups = GROUPS_WEBLOGO.length;
            for (int i = 0; i < numberOfGroups; i++) {
                groupList.add(GROUPS_WEBLOGO[i]);
                groupCategoryList.add(GROUP_CATEGORIES_WEBLOGO[i]);
            }
        }
    }

    public ArrayList getColorList() {
        return colorList;
    }

    public ArrayList getColorSymbolList() {
        return colorSymbolList;
    }

    public ArrayList getGroupList() {
        return groupList;
    }

    public ArrayList getColorCategoryList() {
        return colorCategoryList;
    }

    public ArrayList getGroupCategoryList() {
        return groupCategoryList;
    }

    public SymbolStyle getProteinLogoStyle() {
        //return proteinLogoStyle;
        return new WebLogoProteinStyle(this);
    }

    public int getNumberOfGroups() {
        return numberOfGroups;
    }

    public int getNumberOfColors() {
        return numberOfColors;
    }

    public String getSelectedColoring() {
        return selectedColoring;
    }

    public String getSelectedGrouping() {
        return selectedGrouping;
    }

    public String getSelectedDescription() {
        return selectedDescription;
    }

    public String getSelectedGroups() {
        return selectedGroups;
    }

    public void addColor(Color color, String symbols, String category) {
        colorList.add(color);
        colorSymbolList.add(symbols);
        colorCategoryList.add(category);
        ++numberOfColors;
    }

    public void removeColor(int index) {
        colorList.remove(index);
        colorSymbolList.remove(index);
        colorCategoryList.remove(index);
        --numberOfColors;
    }

    public void setColorSymbols(int index, String symbols) {
        if (symbols == null) {
            throw new IllegalArgumentException("Symbol parameter must be non-null.");
        }
        if (index > (colorSymbolList.size() - 1)) {
            throw new IllegalArgumentException("Index parameter is out of range.");
        }
        colorSymbolList.set(index, symbols);
    }

    public void setColor(int index, Color color) {
        if (color == null) {
            throw new IllegalArgumentException("Color parameter must be non-null.");
        }
        if (index > (colorList.size() - 1)) {
            throw new IllegalArgumentException("Index parameter is out of range.");
        }
        colorList.set(index, color);
    }

    public Color getColor(Symbol symbol) {


        String residue = getSymbolString(symbol);
        Color color = null;

        for (int i = 0; i < colorSymbolList.size(); i++ ) {
            String symbols = colorSymbolList.get(i);
            if (symbols == null) {
                System.out.println("This should not happen!!");
            }
            if (residue == null) {
                System.out.println("This should also not happen!!!");
            }
            if (symbols.contains(residue)) {
                color = colorList.get(i);
            }
        }
        return color;
    }

    private String getSymbolString(Symbol s) {
        
        if (s.equals(ProteinTools.a()))
            return "A";
        else if (s.equals(ProteinTools.c()))
            return "C";
        else if (s.equals(ProteinTools.d()))
            return "D";
        else if (s.equals(ProteinTools.e()))
            return "E";
        else if (s.equals(ProteinTools.f()))
            return "F";
        else if (s.equals(ProteinTools.g()))
            return "G";
        else if (s.equals(ProteinTools.h()))
            return "h";
        else if (s.equals(ProteinTools.i()))
            return "I";
        else if (s.equals(ProteinTools.k()))
            return "K";
        else if (s.equals(ProteinTools.l()))
            return "L";
        else if (s.equals(ProteinTools.m()))
            return "M";
        else if (s.equals(ProteinTools.n()))
            return "N";
        else if (s.equals(ProteinTools.p()))
            return "P";
        else if (s.equals(ProteinTools.q()))
            return "Q";
        else if (s.equals(ProteinTools.r()))
            return "R";
        else if (s.equals(ProteinTools.s()))
            return "S";
        else if (s.equals(ProteinTools.t()))
            return "T";
        else if (s.equals(ProteinTools.u()))
            return "U";
        else if (s.equals(ProteinTools.v()))
            return "V";
        else if (s.equals(ProteinTools.w()))
            return "W";
        else if (s.equals(ProteinTools.y()))
            return "Y";
        else if (s.equals(ProteinTools.ter()))
            return "*";



        return null;
    }

    //The following are 'legacy' methods from a previous version of this class.
    //All (but one) return a comma-delimited string of the amino acid groups defined above in the GROUPS_* arrays.
    //The exception is getPositionSpecificPDZGrouping which is not yet defined in these arrays.

    public static String getHydroxylBasicAcidicpolarHydrophobeGrouping() {
        return Arrays.toString(GROUPS_HYDROXYL);
    }

    public static String getPolarChargedHydrophobeGrouping() {
        return Arrays.toString(GROUPS_POLAR_CHARGED);
    }

    public static String getGroupingHydrophobeBySize() {
        return Arrays.toString(GROUPS_HYDROPHOBIC);
    }

    public static String getWebLogo() {
        return Arrays.toString(GROUPS_WEBLOGO);
    }


    public static String getPDZ() {
        return Arrays.toString(GROUPS_PDZ);
    }

    public static String[] getPositionSpecificPDZGrouping() {
        String[] groupingByPosition = {
                "ST, KRH, DEQN, FLAMPWIVCY, G", //position -3
                "ST, KRH, DEQN, MPVCYFWLIA, G", //position -2 - ST, large hydrophobes, G
                "ST, KRH, DEQN, FLAMPIVCY, G, W", //position -1 - W is separate because of backbone binding capability
                "ST, KRH, DEQN, G, P, VAC, ILM, FWY" //position 0 - hydrophobe by size
        };
        return groupingByPosition;
    }


}
