/*
 * Jalview - A Sequence Alignment Editor and Viewer (Version 2.8.2)
 * Copyright (C) 2014 The Jalview Authors
 * 
 * This file is part of Jalview.
 * 
 * Jalview is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License 
 * as published by the Free Software Foundation, either version 3
 * of the License, or (at your option) any later version.
 *  
 * Jalview 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.  See the GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with Jalview.  If not, see <http://www.gnu.org/licenses/>.
 * The Jalview Authors are detailed in the 'AUTHORS' file.
 */
package jalview.analysis;

import jalview.datamodel.AlignedCodonFrame;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.SequenceI;
import jalview.schemes.ResidueProperties;
import jalview.util.MapList;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * grab bag of useful alignment manipulation operations Expect these to be
 * refactored elsewhere at some point.
 * 
 * @author jimp
 * 
 */
public class AlignmentUtils
{

  /**
   * Represents the 3 possible results of trying to map one alignment to
   * another.
   */
  public enum MappingResult
  {
    Mapped, NotMapped, AlreadyMapped
  }

  /**
   * given an existing alignment, create a new alignment including all, or up to
   * flankSize additional symbols from each sequence's dataset sequence
   * 
   * @param core
   * @param flankSize
   * @return AlignmentI
   */
  public static AlignmentI expandContext(AlignmentI core, int flankSize)
  {
    List<SequenceI> sq = new ArrayList<SequenceI>();
    int maxoffset = 0;
    for (SequenceI s : core.getSequences())
    {
      SequenceI newSeq = s.deriveSequence();
      if (newSeq.getStart() > maxoffset
              && newSeq.getDatasetSequence().getStart() < s.getStart())
      {
        maxoffset = newSeq.getStart();
      }
      sq.add(newSeq);
    }
    if (flankSize > -1)
    {
      maxoffset = flankSize;
    }
    // now add offset to create a new expanded alignment
    for (SequenceI s : sq)
    {
      SequenceI ds = s;
      while (ds.getDatasetSequence() != null)
      {
        ds = ds.getDatasetSequence();
      }
      int s_end = s.findPosition(s.getStart() + s.getLength());
      // find available flanking residues for sequence
      int ustream_ds = s.getStart() - ds.getStart(), dstream_ds = ds
              .getEnd() - s_end;

      // build new flanked sequence

      // compute gap padding to start of flanking sequence
      int offset = maxoffset - ustream_ds;

      // padding is gapChar x ( maxoffset - min(ustream_ds, flank)
      if (flankSize >= 0)
      {
        if (flankSize < ustream_ds)
        {
          // take up to flankSize residues
          offset = maxoffset - flankSize;
          ustream_ds = flankSize;
        }
        if (flankSize < dstream_ds)
        {
          dstream_ds = flankSize;
        }
      }
      char[] upstream = new String(ds.getSequence(s.getStart() - 1
              - ustream_ds, s.getStart() - 1)).toLowerCase().toCharArray();
      char[] downstream = new String(ds.getSequence(s_end - 1, s_end + 1
              + dstream_ds)).toLowerCase().toCharArray();
      char[] coreseq = s.getSequence();
      char[] nseq = new char[offset + upstream.length + downstream.length
              + coreseq.length];
      char c = core.getGapCharacter();
      // TODO could lowercase the flanking regions
      int p = 0;
      for (; p < offset; p++)
      {
        nseq[p] = c;
      }
      // s.setSequence(new String(upstream).toLowerCase()+new String(coreseq) +
      // new String(downstream).toLowerCase());
      System.arraycopy(upstream, 0, nseq, p, upstream.length);
      System.arraycopy(coreseq, 0, nseq, p + upstream.length,
              coreseq.length);
      System.arraycopy(downstream, 0, nseq, p + coreseq.length
              + upstream.length, downstream.length);
      s.setSequence(new String(nseq));
      s.setStart(s.getStart() - ustream_ds);
      s.setEnd(s_end + downstream.length);
    }
    AlignmentI newAl = new jalview.datamodel.Alignment(
            sq.toArray(new SequenceI[0]));
    for (SequenceI s : sq)
    {
      if (s.getAnnotation() != null)
      {
        for (AlignmentAnnotation aa : s.getAnnotation())
        {
          newAl.addAnnotation(aa);
        }
      }
    }
    newAl.setDataset(core.getDataset());
    return newAl;
  }

  /**
   * Returns the index (zero-based position) of a sequence in an alignment, or
   * -1 if not found.
   * 
   * @param al
   * @param seq
   * @return
   */
  public static int getSequenceIndex(AlignmentI al, SequenceI seq)
  {
    int result = -1;
    int pos = 0;
    for (SequenceI alSeq : al.getSequences())
    {
      if (alSeq == seq)
      {
        result = pos;
        break;
      }
      pos++;
    }
    return result;
  }

  /**
   * Returns a map of lists of sequences in the alignment, keyed by sequence
   * name. For use in mapping between different alignment views of the same
   * sequences.
   * 
   * @see jalview.datamodel.AlignmentI#getSequencesByName()
   */
  public static Map<String, List<SequenceI>> getSequencesByName(
          AlignmentI al)
  {
    Map<String, List<SequenceI>> theMap = new LinkedHashMap<String, List<SequenceI>>();
    for (SequenceI seq : al.getSequences())
    {
      String name = seq.getName();
      if (name != null)
      {
        List<SequenceI> seqs = theMap.get(name);
        if (seqs == null)
        {
          seqs = new ArrayList<SequenceI>();
          theMap.put(name, seqs);
        }
        seqs.add(seq);
      }
    }
    return theMap;
  }

  /**
   * Build mapping of protein to cDNA alignment. Mappings are made between
   * sequences which have the same name and compatible lengths. Has a 3-valued
   * result: either Mapped (at least one sequence mapping was created),
   * AlreadyMapped (all possible sequence mappings already exist), or NotMapped
   * (no possible sequence mappings exist).
   * 
   * @param proteinAlignment
   * @param cdnaAlignment
   * @return
   */
  public static MappingResult mapProteinToCdna(
          final AlignmentI proteinAlignment,
          final AlignmentI cdnaAlignment)
  {
    if (proteinAlignment == null || cdnaAlignment == null)
    {
      return MappingResult.NotMapped;
    }

    boolean mappingPossible = false;
    boolean mappingPerformed = false;

    List<SequenceI> thisSeqs = proteinAlignment.getSequences();
  
    /*
     * Build a look-up of cDNA sequences by name, for matching purposes.
     */
    Map<String, List<SequenceI>> cdnaSeqs = cdnaAlignment
            .getSequencesByName();
  
    for (SequenceI aaSeq : thisSeqs)
    {
      AlignedCodonFrame acf = new AlignedCodonFrame();
      List<SequenceI> candidates = cdnaSeqs.get(aaSeq.getName());
      if (candidates == null)
      {
        /*
         * No cDNA sequence with matching name, so no mapping possible for this
         * protein sequence
         */
        continue;
      }
      mappingPossible = true;
      for (SequenceI cdnaSeq : candidates)
      {
        if (!mappingExists(proteinAlignment.getCodonFrames(),
                aaSeq.getDatasetSequence(), cdnaSeq.getDatasetSequence()))
        {
          MapList map = mapProteinToCdna(aaSeq, cdnaSeq);
          if (map != null)
          {
            acf.addMap(cdnaSeq, aaSeq, map);
            mappingPerformed = true;
          }
        }
      }
      proteinAlignment.addCodonFrame(acf);
    }

    /*
     * If at least one mapping was possible but none was done, then the
     * alignments are already as mapped as they can be.
     */
    if (mappingPossible && !mappingPerformed)
    {
      return MappingResult.AlreadyMapped;
    }
    else
    {
      return mappingPerformed ? MappingResult.Mapped
              : MappingResult.NotMapped;
    }
  }

  /**
   * Answers true if the mappings include one between the given (dataset)
   * sequences.
   */
  public static boolean mappingExists(Set<AlignedCodonFrame> set,
          SequenceI aaSeq, SequenceI cdnaSeq)
  {
    if (set != null)
    {
      for (AlignedCodonFrame acf : set)
      {
        if (cdnaSeq == acf.getDnaForAaSeq(aaSeq))
        {
          return true;
        }
      }
    }
    return false;
  }

  /**
   * Build a mapping (if possible) of a protein to a cDNA sequence. The cDNA
   * must be three times the length of the protein, possibly after ignoring
   * start and/or stop codons. Returns null if no mapping is determined.
   * 
   * @param proteinSeqs
   * @param cdnaSeq
   * @return
   */
  public static MapList mapProteinToCdna(SequenceI proteinSeq,
          SequenceI cdnaSeq)
  {
    /*
     * Here we handle either dataset sequence set (desktop) or absent (applet)
     */
    final SequenceI proteinDataset = proteinSeq.getDatasetSequence();
    String aaSeqString = proteinDataset != null ? proteinDataset
            .getSequenceAsString() : proteinSeq.getSequenceAsString();
    final SequenceI cdnaDataset = cdnaSeq.getDatasetSequence();
    String cdnaSeqString = cdnaDataset != null ? cdnaDataset
            .getSequenceAsString() : cdnaSeq.getSequenceAsString();
    if (aaSeqString == null || cdnaSeqString == null)
    {
      return null;
    }

    final int mappedLength = 3 * aaSeqString.length();
    int cdnaLength = cdnaSeqString.length();
    int cdnaStart = 1;
    int cdnaEnd = cdnaLength;
    final int proteinStart = 1;
    final int proteinEnd = aaSeqString.length();

    /*
     * If lengths don't match, try ignoring stop codon.
     */
    if (cdnaLength != mappedLength)
    {
      for (Object stop : ResidueProperties.STOP)
      {
        if (cdnaSeqString.toUpperCase().endsWith((String) stop))
        {
          cdnaEnd -= 3;
          cdnaLength -= 3;
          break;
        }
      }
    }

    /*
     * If lengths still don't match, try ignoring start codon.
     */
    if (cdnaLength != mappedLength
            && cdnaSeqString.toUpperCase().startsWith(
                    ResidueProperties.START))
    {
      cdnaStart += 3;
      cdnaLength -= 3;
    }

    if (cdnaLength == mappedLength)
    {
      MapList map = new MapList(new int[]
      { cdnaStart, cdnaEnd }, new int[]
      { proteinStart, proteinEnd }, 3, 1);
      return map;
    }
    else
    {
      return null;
    }
  }

  /**
   * Align sequence 'seq' to match the alignment of a mapped sequence. Note this
   * currently assumes that we are aligning cDNA to match protein.
   * 
   * @param seq
   *          the sequence to be realigned
   * @param al
   *          the alignment whose sequence alignment is to be 'copied'
   * @param gap
   *          character string represent a gap in the realigned sequence
   * @param preserveUnmappedGaps
   * @param preserveMappedGaps
   * @return true if the sequence was realigned, false if it could not be
   */
  public static boolean alignSequenceAs(SequenceI seq, AlignmentI al,
          String gap, boolean preserveMappedGaps,
          boolean preserveUnmappedGaps)
  {
    /*
     * Get any mappings from the source alignment to the target (dataset) sequence.
     */
    // TODO there may be one AlignedCodonFrame per dataset sequence, or one with
    // all mappings. Would it help to constrain this?
    List<AlignedCodonFrame> mappings = al.getCodonFrame(seq);
    if (mappings == null)
    {
      return false;
    }
  
    /*
     * Locate the aligned source sequence whose dataset sequence is mapped. We
     * just take the first match here (as we can't align cDNA like more than one
     * protein sequence).
     */
    SequenceI alignFrom = null;
    AlignedCodonFrame mapping = null;
    for (AlignedCodonFrame mp : mappings)
    {
      alignFrom = mp.findAlignedSequence(seq.getDatasetSequence(), al);
      if (alignFrom != null)
      {
        mapping = mp;
        break;
      }
    }
  
    if (alignFrom == null)
    {
      return false;
    }
    alignSequenceAs(seq, alignFrom, mapping, gap, al.getGapCharacter(),
            preserveMappedGaps, preserveUnmappedGaps);
    return true;
  }

  /**
   * Align sequence 'alignTo' the same way as 'alignFrom', using the mapping to
   * match residues and codons. Flags control whether existing gaps in unmapped
   * (intron) and mapped (exon) regions are preserved or not. Gaps linking intro
   * and exon are only retained if both flags are set.
   * 
   * @param alignTo
   * @param alignFrom
   * @param mapping
   * @param myGap
   * @param sourceGap
   * @param preserveUnmappedGaps
   * @param preserveMappedGaps
   */
  public static void alignSequenceAs(SequenceI alignTo,
          SequenceI alignFrom,
          AlignedCodonFrame mapping, String myGap, char sourceGap,
          boolean preserveMappedGaps, boolean preserveUnmappedGaps)
  {
    // TODO generalise to work for Protein-Protein, dna-dna, dna-protein
    final char[] thisSeq = alignTo.getSequence();
    final char[] thatAligned = alignFrom.getSequence();
    StringBuilder thisAligned = new StringBuilder(2 * thisSeq.length);
  
    // aligned and dataset sequence positions, all base zero
    int thisSeqPos = 0;
    int sourceDsPos = 0;

    int basesWritten = 0;
    char myGapChar = myGap.charAt(0);
    int ratio = myGap.length();

    /*
     * Traverse the aligned protein sequence.
     */
    int sourceGapMappedLength = 0;
    boolean inExon = false;
    for (char sourceChar : thatAligned)
    {
      if (sourceChar == sourceGap)
      {
        sourceGapMappedLength += ratio;
        continue;
      }

      /*
       * Found a residue. Locate its mapped codon (start) position.
       */
      sourceDsPos++;
      // Note mapping positions are base 1, our sequence positions base 0
      int[] mappedPos = mapping.getMappedRegion(alignTo, alignFrom,
              sourceDsPos);
      if (mappedPos == null)
      {
        /*
         * Abort realignment if unmapped protein. Or could ignore it??
         */
        System.err.println("Can't align: no codon mapping to residue "
                + sourceDsPos + "(" + sourceChar + ")");
        return;
      }

      int mappedCodonStart = mappedPos[0]; // position (1...) of codon start
      int mappedCodonEnd = mappedPos[mappedPos.length - 1]; // codon end pos
      StringBuilder trailingCopiedGap = new StringBuilder();

      /*
       * Copy dna sequence up to and including this codon. Optionally, include
       * gaps before the codon starts (in introns) and/or after the codon starts
       * (in exons).
       * 
       * Note this only works for 'linear' splicing, not reverse or interleaved.
       * But then 'align dna as protein' doesn't make much sense otherwise.
       */
      int intronLength = 0;
      while (basesWritten < mappedCodonEnd && thisSeqPos < thisSeq.length)
      {
        final char c = thisSeq[thisSeqPos++];
        if (c != myGapChar)
        {
          basesWritten++;

          if (basesWritten < mappedCodonStart)
          {
            /*
             * Found an unmapped (intron) base. First add in any preceding gaps
             * (if wanted).
             */
            if (preserveUnmappedGaps && trailingCopiedGap.length() > 0)
            {
              thisAligned.append(trailingCopiedGap.toString());
              intronLength += trailingCopiedGap.length();
              trailingCopiedGap = new StringBuilder();
            }
            intronLength++;
            inExon = false;
          }
          else
          {
            final boolean startOfCodon = basesWritten == mappedCodonStart;
            int gapsToAdd = calculateGapsToInsert(preserveMappedGaps,
                    preserveUnmappedGaps, sourceGapMappedLength, inExon,
                    trailingCopiedGap.length(), intronLength, startOfCodon);
            for (int i = 0; i < gapsToAdd; i++)
            {
              thisAligned.append(myGapChar);
            }
            sourceGapMappedLength = 0;
            inExon = true;
          }
          thisAligned.append(c);
          trailingCopiedGap = new StringBuilder();
        }
        else
        {
          if (inExon && preserveMappedGaps)
          {
            trailingCopiedGap.append(myGapChar);
          }
          else if (!inExon && preserveUnmappedGaps)
          {
            trailingCopiedGap.append(myGapChar);
          }
        }
      }
    }

    /*
     * At end of protein sequence. Copy any remaining dna sequence, optionally
     * including (intron) gaps. We do not copy trailing gaps in protein.
     */
    while (thisSeqPos < thisSeq.length)
    {
      final char c = thisSeq[thisSeqPos++];
      if (c != myGapChar || preserveUnmappedGaps)
      {
        thisAligned.append(c);
      }
    }

    /*
     * All done aligning, set the aligned sequence.
     */
    alignTo.setSequence(new String(thisAligned));
  }

  /**
   * Helper method to work out how many gaps to insert when realigning.
   * 
   * @param preserveMappedGaps
   * @param preserveUnmappedGaps
   * @param sourceGapMappedLength
   * @param inExon
   * @param trailingCopiedGap
   * @param intronLength
   * @param startOfCodon
   * @return
   */
  protected static int calculateGapsToInsert(boolean preserveMappedGaps,
          boolean preserveUnmappedGaps, int sourceGapMappedLength,
          boolean inExon, int trailingGapLength,
          int intronLength, final boolean startOfCodon)
  {
    int gapsToAdd = 0;
    if (startOfCodon)
    {
      /*
       * Reached start of codon. Ignore trailing gaps in intron unless we are
       * preserving gaps in both exon and intron. Ignore them anyway if the
       * protein alignment introduces a gap at least as large as the intronic
       * region.
       */
      if (inExon && !preserveMappedGaps)
      {
        trailingGapLength = 0;
      }
      if (!inExon && !(preserveMappedGaps && preserveUnmappedGaps))
      {
        trailingGapLength = 0;
      }
      if (inExon)
      {
        gapsToAdd = Math.max(sourceGapMappedLength, trailingGapLength);
      }
      else
      {
        if (intronLength + trailingGapLength <= sourceGapMappedLength)
        {
          gapsToAdd = sourceGapMappedLength - intronLength;
        }
        else
        {
          gapsToAdd = Math.min(intronLength + trailingGapLength
                  - sourceGapMappedLength, trailingGapLength);
        }
      }
    }
    else
    {
      /*
       * second or third base of codon; check for any gaps in dna
       */
      if (!preserveMappedGaps)
      {
        trailingGapLength = 0;
      }
      gapsToAdd = Math.max(sourceGapMappedLength, trailingGapLength);
    }
    return gapsToAdd;
  }
}