package jalview.schemes;

import jalview.datamodel.AnnotatedCollectionI;
import jalview.datamodel.HiddenMarkovModel;
import jalview.datamodel.ResidueCount;
import jalview.datamodel.SequenceCollectionI;
import jalview.datamodel.SequenceI;
import jalview.util.ColorUtils;
import jalview.util.Comparison;

import java.awt.Color;
import java.util.HashMap;
import java.util.Map;

/**
 * A colour scheme based on a selected Hidden Markov Model. The colour is
 * <ul>
 * <li>white for a gap</li>
 * <li>red for an insertion</li>
 * <li>orange for negative information content</li>
 * <li>white to blue for increasing information content</li>
 * </ul>
 * where information content is the log ratio
 * 
 * <pre>
 *   log(profile match emission probability / residue background probability>
 * </pre>
 * 
 * using the alignment's background frequencies for residues.
 * 
 * @author tzvanaalten
 *
 */
public class HMMERAlignmentColourScheme extends ResidueColourScheme
{
  /*
   * the ratio, for each symbol, of its frequency to total symbol count
   */
  Map<Character, Double> frequency = new HashMap<>();

  float logTotalCount;

  HiddenMarkovModel hmm;

  /**
   * Constructor given a Hidden Markov Model
   * 
   * @param sg
   * 
   * @param markov
   */
  public HMMERAlignmentColourScheme(AnnotatedCollectionI sg,
          HiddenMarkovModel markov)
  {
    hmm = markov;
    countFrequencies(sg);
  }

  /**
   * Default constructor (required by ColourSchemes.loadColourSchemes)
   */
  public HMMERAlignmentColourScheme()
  {
  }

  @Override
  public Color findColour(char symbol, int position, SequenceI seq,
          String consensusResidue, float pid)
  {
    return findColour(symbol, position);
  }

  /**
   * Returns the colour at a particular symbol at a column in the alignment:
   * <ul>
   * <li>white for a gap</li>
   * <li>red for an insertion</li>
   * <li>orange for negative information content</li>
   * <li>white to blue for increasing information content</li>
   * </ul>
   * 
   * @param symbol
   * @param column
   * @return
   */
  private Color findColour(char symbol, int column)
  {
    if (hmm == null || Comparison.isGap(symbol))
    {
      return Color.white;
    }
    if (Character.isLowerCase(symbol))
    {
      symbol = Character.toUpperCase(symbol);
    }
    double prob = hmm.getMatchEmissionProbability(column, symbol);
    Double freq = frequency.get(symbol);
    if (freq == null)
    {
      return Color.white;
    }
    if (prob == 0)
    {
      return new Color(230, 0, 0);
    }
    double value = Math.log(prob / freq.doubleValue());
    Color colour = null;
    if (value > 0)
    {
      colour = ColorUtils.getGraduatedColour((float) value, 0,
              Color.WHITE, logTotalCount, Color.blue);
    }
    else if (value < 0)
    {
      return Color.ORANGE;
    }
    return colour;
  }

  @Override
  public void alignmentChanged(AnnotatedCollectionI collection,
          Map<SequenceI, SequenceCollectionI> hiddenReps)
  {
    /*
     * ? no need to do anything if alignment is adjusted
     * since findColour() handles everything 
     */
  }

  @Override
  public ColourSchemeI getInstance(AnnotatedCollectionI sg,
          Map<SequenceI, SequenceCollectionI> hiddenRepSequences)
  {
    SequenceI hmmSeq = sg.getHmmConsensus();
    HiddenMarkovModel model = hmmSeq == null ? null : hmmSeq.getHMM();
    return new HMMERAlignmentColourScheme(sg, model);
  }

  /**
   * Answers true if the sequence collection has an HMM consensus sequence, else
   * false
   */
  @Override
  public boolean isApplicableTo(AnnotatedCollectionI ac)
  {
    return ac.getHmmConsensus() != null;
  }

  @Override
  public String getSchemeName()
  {
    return JalviewColourScheme.HMMERA.toString();
  }

  @Override
  public boolean isSimple()
  {
    return false;
  }

  /**
   * Counts and stores the relatively frequency of every residue in the
   * alignment
   * 
   * @param sg
   */
  public void countFrequencies(AnnotatedCollectionI sg)
  {
    ResidueCount counts = new ResidueCount(sg.getSequences());
    int total = counts.getTotalResidueCount(); // excludes gaps

    for (char symbol : counts.getSymbolCounts().symbols)
    {
      double freq = counts.getCount(symbol) / (double) total;
      frequency.put(symbol, freq);
    }
    logTotalCount = (float) Math.log(total);
  }
}