#!/usr/bin/env python2.7

# This script clusters reads with unsupervised KMeans clustering and then aligns all reads from each cluster using
# MarkovAligner to assess their correspondence to the proposed haplotypes generated by Alex's black box.
# The top 3 SNPs/features at each variable position of aligned reads are used as slots in an expanded binary feature
# vector. At the moment, no attempt is made to measure the distance of non-canonical features to the top features.
# The set of clusters in H9ESC correspondent with Alex's haplotypes appears to be the most stable configuration, despite
# notable variation between runs.

# An exemplar is generated by collecting the top feature at each position for each cluster, see stdout for results.

from collections import defaultdict,Counter
import vcf
import pandas as pd
import pysam
import os
from sklearn.cluster import KMeans
from MarkovAligner import *
from matplotlib import pyplot
import sys


pyplot.style.use('simpleMatplotlibStyleSheet')


from scipy.cluster.hierarchy import dendrogram, linkage

# samFileName = "alignedReadsFrankANDYuliaANDMeghan.sam"
samFileName = "margin_align_longer_consensus.sam"
# samFileName = "alignedReadsJasonAdult.sam"
# samFileName = "alignedReadsJasonFetal.sam"
# samFileName = "alignedReadsMeghan.sam"

outFolder = "RESULTS_"+samFileName.split('.')[0]

if not os.path.exists(outFolder):
    os.makedirs(outFolder)

readsByName = dict()

with open(samFileName) as samfile:
    for line in samfile:
        if not line.startswith('@'):
            line = line.split('\t')
            readID = line[0]
            readSequence = line[9]

            # print(readID)

            readsByName[readID] = readSequence


#----------------------------------------------------------------------------------------------------------------------#
# FROM IAN FIDDES

reads = list(pysam.Samfile(samFileName))

recs = list(vcf.Reader(open('variants_longer.vcf')))
recs = [x for x in recs if x.POS <= 1100]
intervals = []

for r in recs:
    start, stop = r._compute_coordinates_for_indel()
    intervals.append([start - 1, stop, r])  # convert to 0-based

# how many distinct feature vectors are there?
feature_vectors = defaultdict(list)  # why defaultdict? <- initializes all values as an empty list
for r in recs:
    for sample in r.samples[:-1]:  # exclude consensus
        gt = sample.gt_bases.split('/')[0]  # remove duplicate SNP (homozygous)
        feature_vectors[sample.sample].append(gt)  # store sample name:[SNP1,SNP2...SNPn] in dictionary

feature_df = pd.DataFrame.from_dict(feature_vectors).T
feature_df = feature_df.drop_duplicates()

# print(feature_df)

featureVectors = []
for a in reads:  #nanopore MarginAlign-ed reads
    r = [a.qname]
    m = {y: x for x, y in a.get_aligned_pairs()}  #a list of aligned read (query) and reference positions
    for start, stop, rec in intervals:  #only collect the SNPs/features from these alignments
        start = m.get(start, None)
        stop = m.get(stop, None)
        if start is not None and stop is not None:
            seq = a.seq[start:stop]
        else:
            seq = None
        r.append(seq)  #store the features as a list
    featureVectors.append(r)



#----------------------------------------------------------------------------------------------------------------------#

characterSet = set(['A','C','G','T'])
nonCharacter = 'N'

#assimilate all characters in feature vectors so that clustering is simplified
for r,row in enumerate(featureVectors):
    assimilatedFeatureVector = list()
    for item in row[1:]:
        if item != None:
            assimilatedFeatureVector.append(''.join([x if x in characterSet else nonCharacter for x in item]))
        else:
            assimilatedFeatureVector.append(nonCharacter)

    featureVectors[r] = [featureVectors[r][0]]+assimilatedFeatureVector


# results = np.array(results)
# observedFeatures = list()
#
# for i in range(results.shape[1]-1):
#     observedFeatures.append(sorted(list(set(results[:,1+i])),key=lambda x: len(x)))
#     print(observedFeatures[i])

# for c,column in enumerate(observedFeatures):
#     for feature in column:

#----------------------------------------------------------------------------------------------------------------------#
# FROM IAN FIDDES (again)

# construct column names based on VCF records
columns = ['read_name']
for i, r in enumerate(recs):
    s = defaultdict(list)
    for sample in r.samples[:-1]:  # exclude consensus
        gt = sample.gt_bases.split('/')[0]  # inherently homozygous
        if 'NOTCH2-' in sample.sample:  # simplify names
            name = 'N2-'+sample.sample.split('-', 1)[1]
        else:
            name = sample.sample.split('-', 1)[1]
        s[gt].append(name)
    name_str = '_'.join([':'.join([x, ','.join(y)]) for x, y in s.iteritems()])
    columns.append('{}/{}'.format(i, name_str))

results_df = pd.DataFrame(featureVectors, columns=columns)
results_df = results_df.set_index('read_name')

#----------------------------------------------------------------------------------------------------------------------#

names = feature_df.index.tolist()[:]
featureArray = feature_df.as_matrix()[:,:]

model = MarkovAligner()
model.buildModel(featureArray,names)

# print(len(results_df))

#iniitialize expanded binary feature vectors, i.e.: [pos1==A, pos1==T, pos1==C, pos2==G, pos2==C...]
#using the top 3 most common features at each position
scoreVectors = [list() for i in range(len(results_df))]
scoreVectorKey = list()


print("\nTOP FEATURES AT EACH VARIANT SITE:")
for col in results_df:
    topCounts = Counter(results_df[col]).most_common(10)
    topFeatures = [x[0] for x in topCounts if x[0] != nonCharacter]
    scoreVectorKey += topFeatures
    # print topFeatures
    print(topCounts)

    for i,observedFeature in enumerate(results_df[col]):

        for referenceFeature in topFeatures:
            # alignment = align.globalms(observedFeature, referenceFeature, 5, -4, -2, -1)
            # score = alignment[0][2]

            if observedFeature == referenceFeature:
                score = 1
            else:
                score = 0

            scoreVectors[i].append(score)

print("\nBINARY FEATURE VECTOR KEY:")
print(scoreVectorKey)
# print(scoreVectors[0:10])

nClusters = 6

print("\nk parameter: %d"%nClusters)

# use binary vectors in KMeans clustering
clustering = KMeans(algorithm="full",n_clusters=nClusters)
kmeansResults = clustering.fit_predict(scoreVectors) #clusterResults =


clusterBins = [list() for c in range(nClusters)]    # initialize one bin for each cluster/haplotype

# storage for alignment results after each cluster is aligned to proposed haplotypes (from VCF)
clusterHaplotypes = [{haplotypeName:list() for haplotypeName in names} for c in range(nClusters)]

# bin the reads into their clusters
for i,read in enumerate(featureVectors):
    clusterBins[kmeansResults[i]].append(read[1:])


print("\nCLUSTER EXEMPLARS (based on frequency per position):")
# align the clusters back to proposed haplotypes (from VCF) and determine an exemplar for each cluster
for b,bin in enumerate(clusterBins):
    bin_df = pd.DataFrame(bin,columns=columns[1:])

    # use MarkovAligner to asses the similarity of each read in the cluster to the proposed haplotypes
    for featureVector in bin:
        # readName = entry[0]
        # print(len(readFeatureVector))

        paths,maxPosteriorProb = model.maximumProbabilityPath(featureVector)

        for path in paths:
            clusterHaplotypes[b][path].append(2**maxPosteriorProb)


    exemplar = list()

    #generate an "exemplar" which in this case is actually just the most frequent non-N character at each pos.
    for col in bin_df:
        topCounts = Counter(bin_df[col]).most_common(4)
        # topFeatures = [x[0] for x in topCounts if x[0] != nonCharacter]
        # print topCounts

        if len(topCounts)>1:
            if topCounts[0][0] != nonCharacter:
                exemplar.append(topCounts[0][0])
            else:
                exemplar.append(topCounts[1][0])
        else:
            exemplar.append(topCounts[0][0])

    print('\t'.join(exemplar))


# generate a plot with x:y subplots where x=nClusters and y=n haplotypes, show the score distributions for each cluster
fig, axes = pyplot.subplots(len(names), nClusters, sharex=True, sharey=True)

for c,cluster in enumerate(clusterHaplotypes):
    for h,bin in enumerate(sorted(cluster)):

        x = cluster[bin]
        axes[h][c].hist(x,25,normed=0,histtype='bar',color=[.118, .565, 1.00],linewidth=0.1,range=[0.125,1.0],log=False)
        if c == 0:
            axes[h][c].set_ylabel(bin[-3:])

        if h == len(axes)-1:
            axes[h][c].set_xlabel("Match Likelihood")
            axes[h][c].set_xticks([0.1,0.5,1.0])

pyplot.show()




