package jalview.schemes;

import jalview.api.AlignViewportI;
import jalview.datamodel.AnnotatedCollectionI;
import jalview.datamodel.HiddenMarkovModel;
import jalview.datamodel.SequenceI;
import jalview.util.ColorUtils;
import jalview.util.Comparison;

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

/**
 * Base class for colour schemes based on a selected Hidden Markov Model. The
 * colour is with reference to an HMM consensus sequence and HMM profile
 * <ul>
 * <li>white for a gap</li>
 * <li>red for an insertion (position is gapped in the HMM consensus)</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>
 * 
 * Sub-class implementations use either global ('Uniprot') or local
 * ('alignment') background frequencies.
 * 
 * @author tzvanaalten
 * @author gmcarstairs
 */
public abstract class HmmerColourScheme extends ResidueColourScheme
{
  private static final Color INSERTION_COLOUR = new Color(230, 0, 0); // reddish

  /*
   * the aligned HMM consensus sequence to use as reference for colouring
   */
  private List<SequenceI> hmmSeqs=null;

  
  private Map<Character, Float> frequencies;

  /**
   * Constructor given a list of Hidden Markov Model consensus sequences. The
   * first sequence provides the HMM profile from which we can read the emission
   * probabilities that determine the colour.
   * 
   * @param hmmSeqs
   */
  public HmmerColourScheme(List<SequenceI> hmmSeqs)
  {
    this.hmmSeqs = hmmSeqs.isEmpty() ? null : hmmSeqs;
  }

  protected   String getAlphabetType()
  {
    if (hmmSeqs!=null) {
      for (SequenceI hmmSeq : hmmSeqs)
    {
      if (hmmSeq != null)
      {
        HiddenMarkovModel hmm = hmmSeq.getHMM();
        if (hmm != null)
        {
          return hmm.getAlphabetType();
        }
      }
    }}
    return ResidueProperties.ALPHABET_AMINO;
  }

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

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

  /**
   * 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 (hmmSeqs==null || Comparison.isGap(symbol))
    {
      // todo: could return probability of seeing a gap ?
      return Color.white;
    }
    // locate first hmm with a non-gap at this position
    for (SequenceI hmmSeq:hmmSeqs)
    {
      if (hmmSeq==null)
      {
        continue;
      }
      if (!Comparison.isGap(hmmSeq.getCharAt(column)))
      {
        if (symbol >= 'a')
        {
          symbol += 'A' - 'a';
        }

        final double prob = hmmSeq.getHMM()
                .getMatchEmissionProbability(column, symbol);

        Float freq = 0f;

        if (!frequencies.containsKey(symbol))
        {
          return Color.WHITE;
        }
        else
        {
          freq = frequencies.get(symbol);
        }

        /*
         * Orange if match emission probability is less than background probability
         */
        double infoRatio = prob / freq.floatValue();
        Color colour = Color.ORANGE;
        if (infoRatio >= 1)
        {
          /*
           * log-scale graduated shade of blue if prob is greater than background  
           */
          float infoLog = (float) Math.log(infoRatio);
          colour = ColorUtils.getGraduatedColour(infoLog, 0, Color.WHITE,
                  getMaxInformationScore(), Color.blue);
        }

        return colour;
      }
    }
    // no more seqs. so
    return INSERTION_COLOUR;
  }

  /**
   * Answers the maximum possible value of information score (log ratio), for
   * use in scaling a graduated colour range
   * 
   * @return
   */
  abstract float getMaxInformationScore();

  /**
   * Answers a new colour scheme instance based on the HMM of the first sequence
   * in ac that has an HMM
   */
  @Override
  public ColourSchemeI getInstance(AlignViewportI viewport,
          AnnotatedCollectionI ac)
  {
    return newInstance(ac);
  }

  /**
   * Answers a new instance of the colour scheme for the given HMM
   * 
   * @param ac
   * @return
   */
  protected abstract HmmerColourScheme newInstance(AnnotatedCollectionI ac);

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

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

  protected Map<Character, Float> getFrequencies()
  {
    return frequencies;
  }

  protected void setFrequencies(Map<Character, Float> frequencies)
  {
    this.frequencies = frequencies;
  }
}