package org.broad.igv.sam;


import org.broad.igv.feature.genome.Genome;
import org.broad.igv.ui.util.MessageUtils;
import org.broad.igv.util.Pair;

import java.util.*;


/**
 * Experimental class for phasing alignments for high ploidy regions.
 *
 * @author Jim Robinson
 */

public class HaplotypeUtils {


    private final AlignmentInterval alignmentInterval;
    Genome genome;

    public HaplotypeUtils(AlignmentInterval alignmentInterval, Genome genome) {
        this.alignmentInterval = alignmentInterval;
        this.genome = genome;
    }

    public void clusterAlignments(String chr, int start, int end, int nClasses) {


        AlignmentCounts counts = this.alignmentInterval.getCounts();

        final byte[] reference = genome.getSequence(chr, start, end);

        // Find snp positions
        List<Integer> snpPos = findVariantPositions(start, end, counts, reference);

        if(snpPos.size() == 0) {
            MessageUtils.showMessage("No variants in selected range.");
            return;
        }

        if(snpPos.size() < nClasses - 1) {
            nClasses = snpPos.size() + 1;
            MessageUtils.showMessage("Not enough variants, reducing # of clusters: " + nClasses);
        }


        // Adjust start and end to min and max snp positions, there is no information outside these bounds
        start = snpPos.get(0) - 1;
        end = snpPos.get(snpPos.size() - 1) + 1;

        // Label alignments
        Map<String, List<Alignment>> labelAlignmentMap = labelAlignments(start, end, snpPos, reference, this.alignmentInterval.getAlignmentIterator());

        // Sort labels (entries) by # of associated alignments
        List<String> labels = new ArrayList(labelAlignmentMap.keySet());
        labels.sort((o1, o2) -> {
            return labelAlignmentMap.get(o2).size() - labelAlignmentMap.get(o1).size();
        });

        // Create initial cluster centroids
        List<V> clusters = new ArrayList<>();
        for (int i = 0; i < nClasses; i++) {
            String label = labels.get(i);
            V v = new V(i + 1, label);
            clusters.add(v);
        }

        // Now assign all labels to a cluster

        int n = 0;
        int max = 50;
        while (true) {
            for (String label : labels) {

                double min = Double.MAX_VALUE;
                V centroid = null;

                for (V c : clusters) {
                    double dist = c.distance(label);
                    if (dist < min) {
                        centroid = c;
                        min = dist;
                    }
                }

                if (centroid != null) {
                    centroid.add(label);
                }
            }

            boolean movement = false;
            for (V c : clusters) {
                if (c.movement()) {
                    movement = true;
                    break;
                }
            }

            if (movement && n++ < max) {
                for (V c : clusters) {
                    c.reset();
                }
            } else {
                break;
            }
        }

        System.out.println("Converged in: " + n);

        // Now label alignments
        for (int i = 0; i < clusters.size(); i++) {

            V c = clusters.get(i);
            String label = "" + c.id;
            for (String l : c.allLabels) {

                List<Alignment> alignments = labelAlignmentMap.get(l);
                for (Alignment a : alignments) {
                    a.setHaplotypeName(label);
                }
            }
        }

    }

    private List<Integer> findVariantPositions(int start, int end, AlignmentCounts counts, byte[] reference) {

        List<Integer> snpPos = new ArrayList<>();

        for (int i = start; i < end; i++) {

            byte ref = reference[i - start];

            float mismatchCount = getMismatchCount(counts, i, ref);

            if (mismatchCount > 0.2f) {
                snpPos.add(i);
            }
        }
        return snpPos;
    }

    public Map<String, List<Alignment>> labelAlignments(int start, int end, List<Integer> positions, byte[] reference, Iterator<Alignment> iter) {

        Map<String, List<Alignment>> alignmentMap = new HashMap<>();

        while (iter.hasNext()) {

            Alignment alignment = iter.next();

            if (start >= alignment.getStart() && end <= alignment.getEnd()) {

                String hapName = "";
                int dist = 0;

                for (Integer pos : positions) {

                    byte ref = reference[pos - start];
                    boolean found = false;
                    for (AlignmentBlock block : alignment.getAlignmentBlocks()) {

                        if (block.isSoftClipped()) continue;
                        if (block.contains(pos)) {
                            int blockOffset = pos - block.getStart();
                            hapName += (char) block.getBase(blockOffset);
                            found = true;

                            if (ref != block.getBase(blockOffset)) {
                                dist++;
                            }

                            break;
                        }
                    }
                    if (!found) {
                        hapName += "_";
                    }
                }

                hapName = hapName.toLowerCase();

                List<Alignment> alignments = alignmentMap.get(hapName);
                if (alignments == null) {
                    alignments = new ArrayList<>();
                    alignmentMap.put(hapName, alignments);
                }
                alignments.add(alignment);

            }
        }

        return alignmentMap;
    }


    public float getMismatchCount(AlignmentCounts counts, int pos, byte ref) {


        float mismatchQualitySum = 0;


        if (ref < 96) ref += 32;  // a fast "toLowercase"
        for (char c : BaseAlignmentCounts.nucleotides) {
            if (c != ref && c != 'n') {
                mismatchQualitySum += counts.getCount(pos, (byte) c);
            }
        }
        return mismatchQualitySum / counts.getTotalCount(pos);


    }

//    List<V> combineClusters(List<V> clusters) {
//
//        List<Pair<Integer, Integer>> combine = new ArrayList<>();
//        for (int i = 0; i < clusters.size(); i++) {
//            for (int j = i + 1; j < clusters.size(); j++) {
//                V c1 = clusters.get(i);
//                V c2 = clusters.get(j);
//                double d = c1.distance(c2);
//                //System.out.println("" + c1.id + " - " + c2.id + "  =  " + d);
//                if (d == 0) {
//                    combine.add(new Pair(i, j));
//                }
//            }
//        }
//    }


    static class V {

        static byte[] foo = {'a', 'c', 't', 'g', '_'};

        int id;
        int n;
        int total;
        Map<Byte, int[]> counts;
        byte[] label;

        Set<String> allLabels;
        Set<String> previousLabels;

        public V(int id, String s) {
            this.id = id;
            this.label = s.toLowerCase().getBytes();

            n = label.length;

            counts = new HashMap<>();
            for (byte b : foo) {
                counts.put(b, new int[n]);
            }

            allLabels = new HashSet<>();
            previousLabels = new HashSet();

            add(s);
        }

        public void add(String s) {

            byte[] m = s.getBytes();

            if (m.length != n) {
                System.err.println("Wrong length");
                return;
            }

            total++;

            for (int i = 0; i < n; i++) {

                byte b = m[i];
                if (b < 95) b += 32;  // a fast "toLowercase"

                int[] cts = counts.get(b);
                if (cts != null) {
                    cts[i]++;
                } else {
                    System.err.println("Unknown nuc: " + ((char) m[i]));
                }
            }

            updateLabel();

            allLabels.add(s);
        }

        void updateLabel() {

            for (int i = 0; i < n; i++) {

                byte bMax = 0;
                int cMax = 0;

                for (byte b : foo) {
                    if (b < 95) b += 32;  // a fast "toLowercase"
                    int[] cts = counts.get(b);
                    if (cts == null) {
                        System.out.println("Null: " + ((char) b));
                    }
                    if (cts[i] > cMax) {
                        cMax = counts.get(b)[i];
                        bMax = b;
                    }
                }

                label[i] = bMax;
            }
        }

        double distance(String s) {

            if (s.length() != n) {
                System.out.println("Unequal lengths");
                return -Integer.MAX_VALUE;
            }

            // Lowercase
            byte[] b = s.toLowerCase().getBytes();

            double d = 0;
            for (int i = 0; i < b.length; i++) {

                // Uncomment for Hamming distance
                //  char c = (char) label[i];
                //   char c1 = s.charAt(i);
                //  if(c != c1) {
                //      d += 1.0;
                //  }

                // % of counts matching
                int[] cts = counts.get(b[i]);
                double pct = cts[i] / ((double) total);
                d += 1 - pct;

            }
            return d;
        }

        double distance(V v) {

            byte[] s = v.label;

            if (s.length != n) {
                System.out.println("Unequal lengths");
                return -Integer.MAX_VALUE;
            }

            double d = 0;
            for (int i = 0; i < s.length; i++) {
                byte c = label[i];
                byte c1 = s[i];
                if (c != c1) {
                    d += 1.0;
                }
            }
            return d;
        }

        public void reset() {

            counts = new HashMap<>();
            for (byte b : foo) {
                counts.put(b, new int[n]);
            }

            previousLabels = new HashSet(allLabels);
            allLabels = new HashSet<>();

            for (int i = 0; i < label.length; i++) {
                counts.get(label[i])[i] = 1;
            }

            total = 1;

        }

        public boolean movement() {

            if (allLabels.size() != previousLabels.size()) {
                return true;
            }
            for (String l : previousLabels) {
                if (!allLabels.contains(l)) return true;
            }
            return false;
        }
    }

}
