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.datamodel.AlignmentI;
import jalview.datamodel.DBRefEntry;
import jalview.datamodel.GeneLoci;
import jalview.datamodel.Mapping;
import jalview.datamodel.Sequence;
import jalview.datamodel.SequenceFeature;
import jalview.datamodel.SequenceI;
import jalview.ext.htsjdk.VCFReader;
import jalview.io.gff.SequenceOntologyI;
import jalview.util.MapList;

import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

public class VCFLoader
{
  AlignmentI al;

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

  /**
   * Loads VCF on to an alignment - provided it can be related to one or more
   * sequence's chromosomal coordinates
   * 
   * @param filePath
   */
  public void loadVCF(String filePath)
  {
    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;
          computePeptideVariants(seq);
        }
      }
      long elapsed = System.currentTimeMillis() - start;
      System.out.println(String.format(
              "Added %d VCF variants to %d sequence(s) (%dms)", varCount,
              seqCount, elapsed));

      reader.close();
    } catch (Throwable e)
    {
      System.err.println("Error processing VCF: " + e.getMessage());
      e.printStackTrace();
    }
  }

  /**
   * (Re-)computes peptide variants from dna variants, for any protein sequence
   * to which the dna sequence has a mapping. Note that although duplicate
   * features may get computed, they will not be added, since duplicate sequence
   * features are ignored in Sequence.addSequenceFeature.
   * 
   * @param dnaSeq
   */
  protected void computePeptideVariants(SequenceI dnaSeq)
  {
    DBRefEntry[] dbrefs = dnaSeq.getDBRefs();
    if (dbrefs == null)
    {
      return;
    }
    for (DBRefEntry dbref : dbrefs)
    {
      Mapping mapping = dbref.getMap();
      if (mapping == null || mapping.getTo() == null
              || mapping.getMap().getFromRatio() != 3)
      {
        continue;
      }
      AlignmentUtils.computeProteinFeatures(dnaSeq, mapping.getTo(),
              mapping.getMap());
    }
  }

  /**
   * 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;
    GeneLoci seqCoords = ((Sequence) seq).getGeneLoci();
    if (seqCoords == null)
    {
      return 0;
    }

    MapList mapping = seqCoords.getMapping();
    List<int[]> seqChromosomalContigs = mapping.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)
  {
    GeneLoci seqCoords = ((Sequence) seq).getGeneLoci();

    String chromosome = seqCoords.getChromosome();
    String seqRef = seqCoords.getReference();
    String species = seqCoords.getSpecies();
    
    /*
     * map chromosomal coordinates from GRCh38 (sequence) to
     * GRCh37 (VCF) if necessary
     */
    int offset = 0;
    if ("GRCh38".equalsIgnoreCase(seqRef) && isVcfRefGrch37)
    {
      int[] newRange = mapReferenceRange(range, chromosome, species,
              "GRCh38",
              "GRCh37");
      offset = newRange[0] - range[0];
      range = newRange;
    }

    /*
     * query the VCF for overlaps
     */
    int count = 0;
    MapList mapping = seqCoords.getMapping();

    CloseableIterator<VariantContext> variants = reader.query(chromosome,
            range[0], range[1]);
    while (variants.hasNext())
    {
      /*
       * get variant location in sequence chromosomal coordinates
       */
      VariantContext variant = variants.next();
      count++;
      int start = variant.getStart() - offset;
      int end = variant.getEnd() - offset;

      /*
       * get location in sequence coordinates
       */
      int[] seqLocation = mapping.locateInFrom(start, end);
      int featureStart = seqLocation[0];
      int featureEnd = seqLocation[1];
      addVariantFeatures(seq, variant, featureStart, featureEnd);
    }

    variants.close();

    return count;
  }

  /**
   * Inspects the VCF variant record, and adds variant features to the sequence
   * 
   * @param seq
   * @param variant
   * @param featureStart
   * @param featureEnd
   */
  protected void addVariantFeatures(SequenceI seq, VariantContext variant,
          int featureStart, int featureEnd)
  {
    StringBuilder sb = new StringBuilder();
    sb.append(variant.getReference().getBaseString());

    int alleleCount = 0;
    for (Allele allele : variant.getAlleles())
    {
      if (!allele.isReference())
      {
        sb.append(",").append(allele.getBaseString());
        alleleCount++;
      }
    }
    String alleles = sb.toString(); // e.g. G,A,C

    String type = SequenceOntologyI.SEQUENCE_VARIANT;
    float score = 0f;
    if (alleleCount == 1)
    {
      try
      {
        score = (float) variant.getAttributeAsDouble("AF", 0d);
      } catch (NumberFormatException e)
      {
        // leave score as 0
      }
    }
    SequenceFeature sf = new SequenceFeature(type, alleles, featureStart,
            featureEnd, score, "VCF");

    /*
     * only add 'alleles' property if a SNP, as we can
     * only handle SNPs when computing peptide variants
     */
    if (variant.isSNP())
    {
      sf.setValue("alleles", alleles);
    }

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

  /**
   * Call the Ensembl REST service that maps from one assembly reference's
   * coordinates to another's
   * 
   * @param range
   * @param chromosome
   * @param species
   * @param fromRef
   * @param toRef
   * @return
   */
  protected int[] mapReferenceRange(int[] range, String chromosome,
          String species, String fromRef, String toRef)
  {
    // TODO call
    // http://rest.ensembl.org/map/species/fromRef/chromosome:range[0]..range[1]:1/toRef?content-type=application/json
    // and parse the JSON response

    // 1922632 is the difference between 37 and 38 for chromosome 17
    return new int[] { range[0] - 1922632, range[1] - 1922632 };
  }
}