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.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;




public class HMMMatchScoreColourScheme extends ResidueColourScheme
{
  
  private Map<Character, Map<Integer, Map<String, Double>>> probabilities;

  private List<Integer> ranges;

  private static double binSize;

  public class MatchProbReader
  {
    private BufferedReader reader;
    
    MatchProbReader() throws FileNotFoundException
    {
      reader = new BufferedReader(
              new FileReader("resources/ProbabilityOfMatch"));
    }
    
    /*
    public Map<Character, Map<String, Double>>  getProbabilities() throws IOException
    {
      Map<Character, Map<String, Double>> probabilities = new HashMap<>();
      
      String[] alphabet = reader.readLine().split("\\,");
      for (int i = 1; i < alphabet.length - 1; i++)
      {
        probabilities.put(alphabet[i].replaceAll("\\ ", "").charAt(0),
                new HashMap<>());
      }
      
      String line = reader.readLine();
      while(line != null)
      {
        String[] contents = line.split("\\,");
        
        for(int i = 1; i < contents.length; i++)
        {
          probabilities.get(alphabet[i].replaceAll("\\ ", "").charAt(0))
                  .put(contents[0], Double
                          .valueOf(contents[i].replaceAll("\\ ", "")));
        }
        line = reader.readLine();
    
      }
      reader.close();
      return probabilities;
    }
    */
    
    public Map<Character, Map<Integer, Map<String, Double>>> getProbabilities()
            throws IOException
    {

      Map<Character, Map<Integer, Map<String, Double>>> probabilities = new HashMap<>();

      ranges = new ArrayList<>();
      ranges.add(0);

      binSize = Double.valueOf((reader.readLine().replaceAll("\\ ", "")));
      String line = reader.readLine();
      char c = line.charAt(0);

      while (line != null)
      {
        line = reader.readLine();
        while (line != null && line.split("\\,").length != 1)
        {
          String[] llrs = line.split("\\,");
          String[] counts = reader.readLine().split("\\,");
          int range = Integer.valueOf(llrs[0]);

          if (!ranges.contains(range))
          {
            ranges.add(range);
          }
          if (!probabilities.containsKey(c))
          {
            probabilities.put(c, new HashMap<>());
          }
          probabilities.get(c).put(range, new HashMap<>());

          for (int i = 1; i < llrs.length; i++)
          {
            probabilities.get(c).get(range).put(
                    llrs[i].replaceAll("\\ ", ""),
                    Double.valueOf(counts[i].replaceAll("\\ ", "")));
          }

          line = reader.readLine();
        }
        if (line != null)
        {
          c = line.charAt(0);
        }
      }

      return probabilities;
    }
    
  }

  private static final Color INSERTION_COLOUR = Color.white;

  /*
   * the aligned HMM consensus sequence to use as reference for colouring
   */
  private SequenceI hmmSeq;

  private HiddenMarkovModel hmm;

  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
   * @throws IOException
   */
  public HMMMatchScoreColourScheme(List<SequenceI> hmmSeqs)
  {
    hmmSeq = hmmSeqs.isEmpty() ? null : hmmSeqs.get(0);
    hmm = hmmSeq == null ? null : hmmSeq.getHMM();

    try
    {
      MatchProbReader probabilityReader = new MatchProbReader();
      probabilities = probabilityReader.getProbabilities();
    } catch (IOException e)
    {
      System.out.println(e.getStackTrace());
    }
  }

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

  @Override
  public Color findColour(char symbol, int column, SequenceI seq,
          String consensusResidue, float pid)
  {
    if (seq == null)
    {
      return null;
    }
    return findColour(symbol, column, seq.gapMap().length);
  }

  // TODO change documentation
  /**
   * 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, int length)
  {
    if (getHmm() == null || Comparison.isGap(symbol))
    {
      return Color.white;
    }
    if (Comparison.isGap(hmmSeq.getCharAt(column)))
    {
      return INSERTION_COLOUR;
    }
    if (Character.isLowerCase(symbol))
    {
      symbol = Character.toUpperCase(symbol);
    }

    double prob = 0;
    if (hmm.getBackgroundFrequencies().containsKey(symbol))
    {
      int lengthBin = getLengthBin(length);

      double llr = Math
              .log(getHmm().getMatchEmissionProbability(column, symbol)
                      / hmm.getBackgroundFrequencies().get(symbol));

      if (!probabilities.containsKey(symbol)
              || !probabilities.get(symbol).get(lengthBin)
              .containsKey(format(llr)))
      {
        return new Color(140, 140, 140);
      }


      prob = probabilities.get(symbol).get(lengthBin).get(format(llr));
    }
    else
    {
      return new Color(140, 140, 140);
    }

    Color colour = Color.ORANGE;
    if (prob >= 0.5)
    {

      colour = ColorUtils.getGraduatedColour((float) prob, 0.5f,
              Color.WHITE, 1f,
              Color.blue);
    }
    else
    {
      colour = ColorUtils.getGraduatedColour((float) prob, 0f, Color.red,
              0.5f, Color.WHITE);
    }

    return colour;
  }

  public static String format(Double d)
  {
    String formatArg = String.valueOf(binSize);

    // if bin size, need format "%.n" where n is number of decimal places
    if (binSize < 1)
    {
      formatArg = "." + formatArg.split("\\.")[1].length();
    }

    Double rounded = Math.round(d / binSize) * binSize;
    String formatted = String.format("%" + formatArg + "f", rounded);

    // format sometimes returns a number rounded to 0 as -0
    // this ensures output will always be 0
    if (Double.valueOf(formatted) == 0)
    {
      formatted = "0";
    }
    return formatted;

  }

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

  /**
   * 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);
  }

  /**
   * Constructor given a sequence collection
   * 
   * @param ac
   */
  public HMMMatchScoreColourScheme(AnnotatedCollectionI ac)
  {
    this(ac.getHmmSequences());
  }

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

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

  /**
   * Answers true if the sequence collection has an HMM consensus sequence and
   * that the first HMM sequence contains background frequencies, else false
   */
  @Override
  public boolean isApplicableTo(AnnotatedCollectionI ac)
  {
    return !ac.getHmmSequences().isEmpty() && ac.getHmmSequences().get(0)
            .getHMM().getBackgroundFrequencies() != null;
  }

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

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

  protected HiddenMarkovModel getHmm()
  {
    return hmm;
  }

  protected SequenceI getHmmSequence()
  {
    return hmmSeq;
  }

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

  private int getLengthBin(int l)
  {
    for (int i = 1; i < ranges.size(); i++)
    {
      if (l >= ranges.get(i - 1) && l < ranges.get(i))
      {
        return ranges.get(i);
      }
    }
    return -1;
  }
}