package jalview.io.vcf;

import htsjdk.samtools.util.CloseableIterator;
import htsjdk.variant.variantcontext.Allele;
import htsjdk.variant.variantcontext.VariantContext;
import htsjdk.variant.vcf.VCFHeader;
import htsjdk.variant.vcf.VCFHeaderLine;

import jalview.analysis.AlignmentUtils;
import jalview.analysis.Dna;
import jalview.api.AlignViewControllerGuiI;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.DBRefEntry;
import jalview.datamodel.GeneLociI;
import jalview.datamodel.Mapping;
import jalview.datamodel.SequenceFeature;
import jalview.datamodel.SequenceI;
import jalview.ext.ensembl.EnsemblMap;
import jalview.ext.htsjdk.VCFReader;
import jalview.io.gff.Gff3Helper;
import jalview.io.gff.SequenceOntologyI;
import jalview.util.MapList;
import jalview.util.MappingUtils;
import jalview.util.MessageManager;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/**
 * A class to read VCF data (using the htsjdk) and add variants as sequence
 * features on dna and any related protein product sequences
 * 
 * @author gmcarstairs
 */
public class VCFLoader
{
  private static final String FEATURE_GROUP_VCF = "VCF";

  private static final String EXCL = "!";

  /*
   * the alignment we are associated VCF data with
   */
  private AlignmentI al;

  /*
   * mappings between VCF and sequence reference assembly regions, as 
   * key = "species!chromosome!fromAssembly!toAssembly
   * value = Map{fromRange, toRange}
   */
  private Map<String, Map<int[], int[]>> assemblyMappings;

  /**
   * Constructor given an alignment context
   * 
   * @param alignment
   */
  public VCFLoader(AlignmentI alignment)
  {
    al = alignment;

    // map of species!chromosome!fromAssembly!toAssembly to {fromRange, toRange}
    assemblyMappings = new HashMap<String, Map<int[], int[]>>();
  }

  /**
   * Starts a new thread to query and load VCF variant data on to the alignment
   * <p>
   * This method is not thread safe - concurrent threads should use separate
   * instances of this class.
   * 
   * @param filePath
   * @param gui
   */
  public void loadVCF(final String filePath,
          final AlignViewControllerGuiI gui)
  {
    if (gui != null)
    {
      gui.setStatus(MessageManager.getString("label.searching_vcf"));
    }

    new Thread()
    {

      @Override
      public void run()
      {
        VCFLoader.this.doLoad(filePath, gui);
      }

    }.start();
  }

  /**
   * Loads VCF on to an alignment - provided it can be related to one or more
   * sequence's chromosomal coordinates.
   * 
   * @param filePath
   * @param gui
   */
  protected void doLoad(String filePath, AlignViewControllerGuiI gui)
  {
    VCFReader reader = null;
    try
    {
      // long start = System.currentTimeMillis();
      reader = new VCFReader(filePath);

      VCFHeader header = reader.getFileHeader();
      VCFHeaderLine ref = header
              .getOtherHeaderLine(VCFHeader.REFERENCE_KEY);
      // check if reference is wrt assembly19 (GRCh37)
      // todo may need to allow user to specify reference assembly?
      boolean isRefGrch37 = (ref != null && ref.getValue().contains(
              "assembly19"));

      int varCount = 0;
      int seqCount = 0;

      /*
       * query for VCF overlapping each sequence in turn
       */
      for (SequenceI seq : al.getSequences())
      {
        int added = loadVCF(seq, reader, isRefGrch37);
        if (added > 0)
        {
          seqCount++;
          varCount += added;
          transferAddedFeatures(seq);
        }
      }
      if (gui != null)
      {
        // long elapsed = System.currentTimeMillis() - start;
        String msg = MessageManager.formatMessage("label.added_vcf",
                varCount, seqCount);
        gui.setStatus(msg);
        gui.getFeatureSettingsUI().discoverAllFeatureData();
      }
    } catch (Throwable e)
    {
      System.err.println("Error processing VCF: " + e.getMessage());
      e.printStackTrace();
      if (gui != null)
      {
        gui.setStatus("Error occurred - see console for details");
      }
    } finally
    {
      if (reader != null)
      {
        try
        {
          reader.close();
        } catch (IOException e)
        {
          // ignore
        }
      }
    }
  }

  /**
   * Transfers VCF features to sequences to which this sequence has a mapping.
   * If the mapping is 1:3, computes peptide variants from nucleotide variants.
   * 
   * @param seq
   */
  protected void transferAddedFeatures(SequenceI seq)
  {
    DBRefEntry[] dbrefs = seq.getDBRefs();
    if (dbrefs == null)
    {
      return;
    }
    for (DBRefEntry dbref : dbrefs)
    {
      Mapping mapping = dbref.getMap();
      if (mapping == null || mapping.getTo() == null)
      {
        continue;
      }

      SequenceI mapTo = mapping.getTo();
      MapList map = mapping.getMap();
      if (map.getFromRatio() == 3)
      {
        /*
         * dna-to-peptide product mapping
         */
        AlignmentUtils.computeProteinFeatures(seq, mapTo, map);
      }
      else
      {
        /*
         * nucleotide-to-nucleotide mapping e.g. transcript to CDS
         */
        // TODO no DBRef to CDS is added to transcripts
        List<SequenceFeature> features = seq.getFeatures()
                .getPositionalFeatures(SequenceOntologyI.SEQUENCE_VARIANT);
        for (SequenceFeature sf : features)
        {
          if (FEATURE_GROUP_VCF.equals(sf.getFeatureGroup()))
          {
            transferFeature(sf, mapTo, map);
          }
        }
      }
    }
  }

  /**
   * Tries to add overlapping variants read from a VCF file to the given
   * sequence, and returns the number of overlapping variants found. Note that
   * this requires the sequence to hold information as to its chromosomal
   * positions and reference, in order to be able to map the VCF variants to the
   * sequence.
   * 
   * @param seq
   * @param reader
   * @param isVcfRefGrch37
   * @return
   */
  protected int loadVCF(SequenceI seq, VCFReader reader,
          boolean isVcfRefGrch37)
  {
    int count = 0;
    GeneLociI seqCoords = seq.getGeneLoci();
    if (seqCoords == null)
    {
      return 0;
    }

    List<int[]> seqChromosomalContigs = seqCoords.getMap().getToRanges();
    for (int[] range : seqChromosomalContigs)
    {
      count += addVcfVariants(seq, reader, range, isVcfRefGrch37);
    }

    return count;
  }

  /**
   * Queries the VCF reader for any variants that overlap the given chromosome
   * region of the sequence, and adds as variant features. Returns the number of
   * overlapping variants found.
   * 
   * @param seq
   * @param reader
   * @param range
   *          start-end range of a sequence region in its chromosomal
   *          coordinates
   * @param isVcfRefGrch37
   *          true if the VCF is with reference to GRCh37
   * @return
   */
  protected int addVcfVariants(SequenceI seq, VCFReader reader,
          int[] range, boolean isVcfRefGrch37)
  {
    GeneLociI seqCoords = seq.getGeneLoci();

    String chromosome = seqCoords.getChromosomeId();
    String seqRef = seqCoords.getAssemblyId();
    String species = seqCoords.getSpeciesId();

    // TODO handle species properly
    if ("".equals(species))
    {
      species = "human";
    }

    /*
     * map chromosomal coordinates from GRCh38 (sequence) to
     * GRCh37 (VCF) if necessary
     */
    // TODO generalise for other assemblies and species
    int offset = 0;
    String fromRef = "GRCh38";
    if (fromRef.equalsIgnoreCase(seqRef) && isVcfRefGrch37)
    {
      String toRef = "GRCh37";
      int[] newRange = mapReferenceRange(range, chromosome, species,
              fromRef, toRef);
      if (newRange == null)
      {
        System.err.println(String.format(
                "Failed to map %s:%s:%s:%d:%d to %s", species, chromosome,
                fromRef, range[0], range[1], toRef));
        return 0;
      }
      offset = newRange[0] - range[0];
      range = newRange;
    }

    boolean forwardStrand = range[0] <= range[1];

    /*
     * query the VCF for overlaps
     * (convert a reverse strand range to forwards)
     */
    int count = 0;
    MapList mapping = seqCoords.getMap();

    int fromLocus = Math.min(range[0], range[1]);
    int toLocus = Math.max(range[0], range[1]);
    CloseableIterator<VariantContext> variants = reader.query(chromosome,
            fromLocus, toLocus);
    while (variants.hasNext())
    {
      /*
       * get variant location in sequence chromosomal coordinates
       */
      VariantContext variant = variants.next();

      /*
       * we can only process SNP variants (which can be reported
       * as part of a MIXED variant record
       */
      if (!variant.isSNP() && !variant.isMixed())
      {
        continue;
      }

      count++;
      int start = variant.getStart() - offset;
      int end = variant.getEnd() - offset;

      /*
       * convert chromosomal location to sequence coordinates
       * - null if a partially overlapping feature
       */
      int[] seqLocation = mapping.locateInFrom(start, end);
      if (seqLocation != null)
      {
        addVariantFeatures(seq, variant, seqLocation[0], seqLocation[1],
                forwardStrand);
      }
    }

    variants.close();

    return count;
  }

  /**
   * Inspects the VCF variant record, and adds variant features to the sequence.
   * Only SNP variants are added, not INDELs.
   * <p>
   * If the sequence maps to the reverse strand of the chromosome, reference and
   * variant bases are recorded as their complements (C/G, A/T).
   * 
   * @param seq
   * @param variant
   * @param featureStart
   * @param featureEnd
   * @param forwardStrand
   */
  protected void addVariantFeatures(SequenceI seq, VariantContext variant,
          int featureStart, int featureEnd, boolean forwardStrand)
  {
    byte[] reference = variant.getReference().getBases();
    if (reference.length != 1)
    {
      /*
       * sorry, we don't handle INDEL variants
       */
      return;
    }

    /*
     * for now we extract allele frequency as feature score; note
     * this attribute is String for a simple SNP, but List<String> if
     * multiple alleles at the locus; we extract for the simple case only
     */
    Object af = variant.getAttribute("AF");
    float score = 0f;
    if (af instanceof String)
    {
      try
      {
        score = Float.parseFloat((String) af);
      } catch (NumberFormatException e)
      {
        // leave as 0
      }
    }

    StringBuilder sb = new StringBuilder();
    sb.append(forwardStrand ? (char) reference[0] : complement(reference));

    /*
     * inspect alleles and record SNP variants (as the variant
     * record could be MIXED and include INDEL and SNP alleles)
     * warning: getAlleles gives no guarantee as to the order 
     * in which they are returned
     */
    for (Allele allele : variant.getAlleles())
    {
      if (!allele.isReference())
      {
        byte[] alleleBase = allele.getBases();
        if (alleleBase.length == 1)
        {
          sb.append(",").append(
                  forwardStrand ? (char) alleleBase[0]
                          : complement(alleleBase));
        }
      }
    }
    String alleles = sb.toString(); // e.g. G,A,C

    String type = SequenceOntologyI.SEQUENCE_VARIANT;

    SequenceFeature sf = new SequenceFeature(type, alleles, featureStart,
            featureEnd, score, FEATURE_GROUP_VCF);

    sf.setValue(Gff3Helper.ALLELES, alleles);

    Map<String, Object> atts = variant.getAttributes();
    for (Entry<String, Object> att : atts.entrySet())
    {
      sf.setValue(att.getKey(), att.getValue());
    }
    seq.addSequenceFeature(sf);
  }

  /**
   * A convenience method to complement a dna base and return the string value
   * of its complement
   * 
   * @param reference
   * @return
   */
  protected String complement(byte[] reference)
  {
    return String.valueOf(Dna.getComplement((char) reference[0]));
  }

  /**
   * Determines the location of the query range (chromosome positions) in a
   * different reference assembly.
   * <p>
   * If the range is just a subregion of one for which we already have a mapping
   * (for example, an exon sub-region of a gene), then the mapping is just
   * computed arithmetically.
   * <p>
   * Otherwise, calls the Ensembl REST service that maps from one assembly
   * reference's coordinates to another's
   * 
   * @param queryRange
   *          start-end chromosomal range in 'fromRef' coordinates
   * @param chromosome
   * @param species
   * @param fromRef
   *          assembly reference for the query coordinates
   * @param toRef
   *          assembly reference we wish to translate to
   * @return the start-end range in 'toRef' coordinates
   */
  protected int[] mapReferenceRange(int[] queryRange, String chromosome,
          String species, String fromRef, String toRef)
  {
    /*
     * first try shorcut of computing the mapping as a subregion of one
     * we already have (e.g. for an exon, if we have the gene mapping)
     */
    int[] mappedRange = findSubsumedRangeMapping(queryRange, chromosome,
            species, fromRef, toRef);
    if (mappedRange != null)
    {
      return mappedRange;
    }

    /*
     * call (e.g.) http://rest.ensembl.org/map/human/GRCh38/17:45051610..45109016:1/GRCh37
     */
    EnsemblMap mapper = new EnsemblMap();
    int[] mapping = mapper.getMapping(species, chromosome, fromRef, toRef,
            queryRange);

    if (mapping == null)
    {
      // mapping service failure
      return null;
    }

    /*
     * save mapping for possible future re-use
     */
    String key = makeRangesKey(chromosome, species, fromRef, toRef);
    if (!assemblyMappings.containsKey(key))
    {
      assemblyMappings.put(key, new HashMap<int[], int[]>());
    }

    assemblyMappings.get(key).put(queryRange, mapping);

    return mapping;
  }

  /**
   * If we already have a 1:1 contiguous mapping which subsumes the given query
   * range, this method just calculates and returns the subset of that mapping,
   * else it returns null. In practical terms, if a gene has a contiguous
   * mapping between (for example) GRCh37 and GRCh38, then we assume that its
   * subsidiary exons occupy unchanged relative positions, and just compute
   * these as offsets, rather than do another lookup of the mapping.
   * <p>
   * If in future these assumptions prove invalid (e.g. for bacterial dna?!),
   * simply remove this method or let it always return null.
   * <p>
   * Warning: many rapid calls to the /map service map result in a 429 overload
   * error response
   * 
   * @param queryRange
   * @param chromosome
   * @param species
   * @param fromRef
   * @param toRef
   * @return
   */
  protected int[] findSubsumedRangeMapping(int[] queryRange, String chromosome,
          String species, String fromRef, String toRef)
  {
    String key = makeRangesKey(chromosome, species, fromRef, toRef);
    if (assemblyMappings.containsKey(key))
    {
      Map<int[], int[]> mappedRanges = assemblyMappings.get(key);
      for (Entry<int[], int[]> mappedRange : mappedRanges.entrySet())
      {
        int[] fromRange = mappedRange.getKey();
        int[] toRange = mappedRange.getValue();
        if (fromRange[1] - fromRange[0] == toRange[1] - toRange[0])
        {
          /*
           * mapping is 1:1 in length, so we trust it to have no discontinuities
           */
          if (MappingUtils.rangeContains(fromRange, queryRange))
          {
            /*
             * fromRange subsumes our query range
             */
            int offset = queryRange[0] - fromRange[0];
            int mappedRangeFrom = toRange[0] + offset;
            int mappedRangeTo = mappedRangeFrom + (queryRange[1] - queryRange[0]);
            return new int[] { mappedRangeFrom, mappedRangeTo };
          }
        }
      }
    }
    return null;
  }

  /**
   * Transfers the sequence feature to the target sequence, locating its start
   * and end range based on the mapping. Features which do not overlap the
   * target sequence are ignored.
   * 
   * @param sf
   * @param targetSequence
   * @param mapping
   *          mapping from the feature's coordinates to the target sequence
   */
  protected void transferFeature(SequenceFeature sf,
          SequenceI targetSequence, MapList mapping)
  {
    int[] mappedRange = mapping.locateInTo(sf.getBegin(), sf.getEnd());
  
    if (mappedRange != null)
    {
      String group = sf.getFeatureGroup();
      int newBegin = Math.min(mappedRange[0], mappedRange[1]);
      int newEnd = Math.max(mappedRange[0], mappedRange[1]);
      SequenceFeature copy = new SequenceFeature(sf, newBegin, newEnd,
              group, sf.getScore());
      targetSequence.addSequenceFeature(copy);
    }
  }

  /**
   * Formats a ranges map lookup key
   * 
   * @param chromosome
   * @param species
   * @param fromRef
   * @param toRef
   * @return
   */
  protected static String makeRangesKey(String chromosome, String species,
          String fromRef, String toRef)
  {
    return species + EXCL + chromosome + EXCL + fromRef + EXCL
            + toRef;
  }
}