#!/usr/bin/env python

from multiprocessing import Pool

from singlecellmultiomics.bamProcessing.bamFunctions import mate_iter

import argparse

import pysam

from glob import glob

import pandas as pd

from singlecellmultiomics.bamProcessing import get_contig_sizes

from collections import Counter, defaultdict

from singlecellmultiomics.features import FeatureContainer

import os

from matplotlib.patches import Rectangle

import matplotlib as mpl

from scipy.ndimage import gaussian_filter

import seaborn as sns

import numpy as np

import matplotlib.pyplot as plt

from itertools import product

from singlecellmultiomics.bamProcessing import get_contigs_with_reads

def _generate_count_dict(args):

    bam_path, bin_size, contig, start, stop = args #reference_path  = args

    #reference_handle = pysam.FastaFile(reference_path)

    #reference = CachedFasta(reference_handle)

    cut_counts = defaultdict(Counter )

    i = 0

    with pysam.AlignmentFile(bam_path) as alignments:

        for R1,R2 in mate_iter(alignments, contig=contig, start=start, stop=stop):

            if R1 is None or R1.is_duplicate or not R1.has_tag('DS') or R1.is_qcfail:

                continue

            cut_pos = R1.get_tag('DS')

            sample = R1.get_tag('SM')

            bin_idx=int(cut_pos/bin_size)*bin_size

            cut_counts[(contig,bin_idx)][sample] += 1

    return cut_counts, contig, bam_path

def get_binned_counts(bams, bin_size, regions=None):

    fs = 1000

    if regions is None:

        regions = [(c,None,None) for c in get_contig_sizes(bams[0]).keys()]

    else:

        for i,r in enumerate(regions):

            if type(r)==str:

                regions[i] = (r,None,None)

            else:

                contig, start, end =r

                if type(start)==int:

                    start = max(0,start-fs)

                regions[i] = (contig,start,end)

    jobs = [(bam_path, bin_size, *region) for region, bam_path in product(regions, bams)]

    cut_counts = defaultdict(Counter)

    with Pool() as workers:

        for i, (cc, contig, bam_path) in enumerate(workers.imap(_generate_count_dict,jobs)):

            for k,v in cc.items():

                cut_counts[k] += v

            print(i,'/', len(jobs), end='\r')

    return pd.DataFrame(cut_counts).T

def plot_region(counts, features, contig, start, end, sigma=2, target=None, caxlabel='Molecules per spike-in'):

    if target is None:

        target = f'{contig}_{start}_{end}.png'

    def create_gene_models(start,end,ax):

        exon_height = 0.010

        gene_height = 0.0002

        spacer = 0.035

        overlap_dist = 200_000

        gene_y = {}

        ymax = 0

        for fs,fe,name,strand, feature_meta in features.features[contig]:

            if not (((fs>=start or fe>=start) and (fs<=end or fe<=end))):

                continue

            feature_meta = dict(feature_meta)

            if feature_meta.get('type') == 'gene':

                if not 'gene_name' in  feature_meta or feature_meta.get('gene_name').startswith('AC'):

                    continue

                # Determine g-y coordinate:

                gy_not_avail = set()

                for gene,(s,e,loc) in gene_y.items():

                    if (s+overlap_dist>=fs and s-overlap_dist<=fe) or (e+overlap_dist>=fs and e-overlap_dist<=fe):

                        # Overlap:

                        gy_not_avail.add(loc)

                gy = 0

                while gy in gy_not_avail:

                    gy+=1

                gene_y[name] = (fs,fe,gy)

                y_offset = gy * spacer

                ymax = max(y_offset+gene_height,ymax)

                r = Rectangle((fs,-gene_height*0.5 + y_offset), fe-fs, gene_height, angle=0.0, color='k')

                ax.add_patch( r )

                ax.text((fe+fs)*0.5,-1.6*exon_height + y_offset,feature_meta.get('gene_name'),horizontalalignment='center',

              verticalalignment='center',fontsize=3)

                #print(feature_meta)

        if False:

            for xx in range(3):

                for fs,fe,name,strand, feature_meta in features.features[contig]:

                    if not (((fs>=start or fe>=start) and (fs<=end or fe<=end))):

                        continue

                    feature_meta = dict(feature_meta)

                    if not name in gene_y:

                        continue

                    if feature_meta.get('type') == 'exon':

                        y_offset = gene_y[name][2]*spacer

                        ymax = max(y_offset+exon_height,ymax)

                        r = Rectangle((fs,-exon_height*0.5 + y_offset), fe-fs, exon_height, angle=0.0,color='k', lw=0)

                        ax.add_patch( r )

        ax.set_xlim(start,end)

        ax.set_ylim(-0.1,ymax)

        #ax.axis('off')

        ax.set_yticks([])

        ax.set_xlabel(f'chr{contig} location bp', fontsize=6)

        #print([t.get_text() for t in ax.get_xticklabels()])

        #ax.set_xticklabels([t.get_text() for t in ax.get_xticklabels()],fontsize=4)

        ax.set_xticklabels(ax.get_xticks(), fontsize=4)

        ax.tick_params(length=0.5)

    for sigma in range(2,3):

        mpl.rcParams['figure.dpi'] = 300

        font = {'family' : 'helvetica',

                'weight' : 'normal',

                'size'   : 8}

        mpl.rc('font', **font)

        if end - start < 3_000_000:

            mode ='k'

            stepper = 100_000

            res = 100

        else:

            mode='M'

            stepper=1_000_000

            res = 1

        qf = counts.loc[:, [(c,p) for c,p in counts if c==contig and p>=start and p<=end] ].sort_index()

        qf = qf.sort_index(1).sort_index(0)

        qf = pd.DataFrame(gaussian_filter(qf, sigma=(0.00001,sigma)), index=qf.index, columns=qf.columns)

        qf = qf.sort_index(1).sort_index(0)

        cm = sns.clustermap(qf,

               #z_score=0,

                row_cluster=False,

                col_cluster=False,

                vmax=np.percentile(qf,99.5),#0.0005,

                #vmax=10,

                dendrogram_ratio=0.1,

                #row_colors=row_colors.loc[qf.index].drop('LOWESS_STAGE',1),

                figsize=(8,4), cmap='Greys', cbar_kws={"shrink": .1},

                cbar_pos=(0.0, 0.5, 0.01, 0.16),)

        ax = cm.ax_col_dendrogram

        qf.mean().plot.bar(ax=ax,color='k',width=1)

        ax.set_yticks([])

        cm.ax_heatmap.set_xticks([]) #np.arange(start,end, 1_000_000))

        cm.ax_heatmap.set_yticks([])

        cm.ax_heatmap.set_ylabel(f'{qf.shape[0]} single cells', fontsize=8)

        cm.ax_heatmap.tick_params(length=0.5)

        cm.ax_heatmap.set_xlabel(None)

        ax.grid()

        cm.cax.set_ylabel(caxlabel,fontsize=6)

        cm.cax.tick_params(labelsize=4)

        #plt.suptitle(mark, x=0.05)

        fig = plt.gcf()

        heatmap_start_x,heatmap_start_y, heatmap_end_x, heatmap_end_y = cm.ax_heatmap.get_position().bounds

        width = heatmap_end_x #-heatmap_start_x

        height = 0.2 if features is not None else 0.05

        ax = fig.add_axes(  (heatmap_start_x, heatmap_start_y-height-0.02, width, height)  )

        ax.ticklabel_format(axis='x',style='sci')

        sns.despine(fig=fig, ax=ax)

        if features is not None:

            create_gene_models(start,end,ax=ax)

        else:

            ax.set_xlim(start,end)

            #ax.axis('off')

            ax.set_yticks([])

            ax.set_xlabel(f'chr{contig} location bp', fontsize=6)

            #ax.set_xticklabels(ax.get_xticks(), fontsize=4)

            plt.xticks(fontsize=4)

            ax.tick_params(length=0.5)

        plt.savefig(target)

        plt.close()

if __name__=='__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description='Plot a genomic region')

    argparser.add_argument('bams', type=str, nargs='+', help='(X) Training bam files')

    argparser.add_argument('-regions', type=str, help='Regions to plot, with a bin size behind it, for example: 1:1000-100000:1000 , will be a single region plotted with a 1000bp bin size split regions by commas without a space')

    argparser.add_argument('-features', type=str, help='Gene models to plot (.gtf file or .gtf.gz)', required=False)

    argparser.add_argument('-norm', type=str, help='Normalize to, select from : total-molecules,spike', default='total-molecules')

    argparser.add_argument('-prefix', type=str, help='Prefix for output file',default='')

    argparser.add_argument('-format', type=str, help='png or svg',default='png')

    args = argparser.parse_args()

    regions = []

    contigs = set()

    for region in args.regions.split(','):

        contig = region.split(':')[0]

        if not '-' in region:

            start, end = None, None

        else:

            start, end = region.split(':')[1].split('-')

            start = int(start)

            end = int(end)

        bin_size = int(region.split(':')[-1])

        if start is not None:

            print(f'Region: {contig} from {start} to {end} with bin size : {bin_size}')

        else:

            print(f'Region: {contig} with bin size : {bin_size}')

        contigs.add(contig)

        regions.append( ((contig,start,end), bin_size))

    contigs=list(contigs)

    bams = args.bams

    if args.features is not None:

        print('Reading features')

        features = FeatureContainer()

        if len(contigs)==1:

            print(f'Reading only features from {contigs[0]}')

            features.loadGTF(args.features,store_all=True,contig=contigs[0])

        else:

            features.loadGTF(args.features,store_all=True)

    else:

        features = None

    print('Counting')

    # Obtain counts per cell

    norm = 'spike'

    if norm == 'spike':

        normalize_to_counts = get_binned_counts(bams, bin_size=10_000_000, regions=['J02459.1'])

    elif norm=='total-molecules':

        normalize_to_counts = get_binned_counts(bams, bin_size=10_000_000)

    for region, region_bin_size in regions:

        print(f'Plotting {region}')

        contig, start, end = region

        region_counts = get_binned_counts(bams, region_bin_size, regions=[ region ] )

        counts = (region_counts/normalize_to_counts.sum()).fillna(0).T.sort_index(1).sort_index(0)

        # Fill non intialized bins with zeros:

        add = []

        for i in np.arange(counts.columns[0][1], counts.columns[-1][1], region_bin_size):

            if not (contig,i) in counts.columns:

                add.append((contig,i))

        for a in add:

            counts[a] = 0

        counts = counts.sort_index(1)

        target = args.prefix+f'{contig}_{start}-{end}_{region_bin_size}.{args.format}'

        plot_region(counts, features, contig, start, end, sigma=2, target=target, caxlabel='Molecules per spike-in' if norm =='spike' else 'Molecules / total molecules')