#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
from singlecellmultiomics.molecule import MoleculeIterator, CHICNLAMolecule, TAPSNlaIIIMolecule,TAPSCHICMolecule,TAPS
from singlecellmultiomics.fragment import CHICFragment, NlaIIIFragment
import pysam
from pysamiterators import CachedFasta
from singlecellmultiomics.variants.substitutions import conversion_dict_stranded
from collections import defaultdict
from singlecellmultiomics.utils import reverse_complement, complement
from glob import glob
from multiprocessing import Pool
from singlecellmultiomics.bamProcessing.bamFunctions import get_reference_path_from_bam
from collections import Counter
import pandas as pd
import matplotlib as mpl
from singlecellmultiomics.utils.sequtils import phred_to_prob, prob_to_phred
import seaborn as sns
import argparse

def update_mutation_dict(molecule,reference, conversions_per_library, context_obs):

    consensus = molecule.get_consensus(dove_safe=True,
                                       min_phred_score=22,
                                       skip_first_n_cycles_R1=10,
                                       skip_last_n_cycles_R1=20,
                                       skip_first_n_cycles_R2=10,
                                       skip_last_n_cycles_R2=20,
                                       dove_R2_distance=15,
                                       dove_R1_distance=15


                                      )

    nm = 0


    contexts_to_add = []


    for (chrom,pos), base in consensus.items():
        context = reference.fetch(chrom, pos-1, pos+2).upper()

        if len(context)!=3:
            continue

        # Check if the base matches or the refence contains N's
        if 'N' in context or len(context)!=3:
            continue

        # Ignore germline variants:
        #if might_be_variant(chrom, pos,  known):
        #    continue

        if not molecule.strand: # reverse template
            context = reverse_complement(context)
            base = complement(base)

        if context[1]!='C' and context[1]!=base:
            nm+=1

        contexts_to_add.append((context,base))


    if nm>5:
        nm=5

    k = tuple((*molecule.sample.rsplit('_',2), nm))
    for (context, base) in contexts_to_add:

        context_obs[ k ][context] += 1
        try:
            conversions_per_library[k][(context, base)] += 1
        except:
            pass


def get_conversion_counts(args):


    taps = TAPS()

    conversions_per_library = defaultdict( conversion_dict_stranded )
    context_obs = defaultdict( Counter )

    bam,refpath,method,every_fragment_as_molecule, spikein_name  = args


    if method=='nla':
        fragment_class=NlaIIIFragment
        molecule_class=TAPSNlaIIIMolecule
    else:
        fragment_class=CHICFragment
        molecule_class=TAPSCHICMolecule

    with pysam.FastaFile(refpath) as ref:
        reference = CachedFasta(ref)


        print(f'Processing {bam}')

        with pysam.AlignmentFile(bam, threads=8) as al:

            for molecule in MoleculeIterator(
                al,
                fragment_class=fragment_class,
                molecule_class=molecule_class,
                molecule_class_args={
                     'reference':reference,
                     'taps':taps,
                    'taps_strand':'R'
                },
                every_fragment_as_molecule=every_fragment_as_molecule,
                fragment_class_args={},
                contig = spikein_name
            ):
                update_mutation_dict(molecule, reference ,conversions_per_library, context_obs)

    return conversions_per_library, context_obs


def generate_taps_conversion_stats(bams, reference_path, prefix, method, every_fragment_as_molecule, spikein_name, n_threads=None):
    if reference_path is None:
        reference_path = get_reference_path_from_bam(bams[0])

    print(f'Reference at {reference_path}')
    if reference_path is None:
        raise ValueError('Please supply a reference fasta file')

    conversions_per_library = defaultdict( conversion_dict_stranded )
    context_obs = defaultdict( Counter )

    with Pool(n_threads) as workers:

        for cl, co in workers.imap(get_conversion_counts, [(bam, reference_path, method, every_fragment_as_molecule, spikein_name) for bam in bams] ):

            for lib, obs in cl.items():
                for k,v in obs.items():
                    conversions_per_library[lib][k] +=v

            for lib, obs in co.items():
                for k,v in obs.items():
                    context_obs[lib][k] += v


    qf = pd.DataFrame(context_obs)
    qf.to_csv(f'{prefix}_conversions_counts_raw_lambda.csv')
    ###
    indices = []
    for lib, qqf in qf.groupby(level=0,axis=1):

        ser = qqf.sum(level=(0,1),axis=1).sum().sort_values(ascending=False)
        ser = ser[ser>5000][::-1]
        indices += list(ser.index)


    ###

    normed_conversions_per_library = defaultdict( conversion_dict_stranded )

    for INDEX in context_obs:
        for (context, base),obs in conversions_per_library[INDEX].items():
            try:
                normed_conversions_per_library[INDEX][(context,base)] = obs/ context_obs[INDEX][context]
            except Exception:
                pass

    df = pd.DataFrame(normed_conversions_per_library)
    df = df[ [INDEX for INDEX in df if (INDEX[0],  INDEX[1]) in indices] ]

    df = df.loc[ [(context, base)for context, base in df.index if context[1]=='C' and base=='T' and context.endswith('CG')]  ]
    df = df.T

    df.to_csv(f'{prefix}_conversions_lambda.csv')

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

    samples = []

    for (lib, cell, nm), row in df.iterrows():

        if nm!=0:
            continue

        for context, base in [('ACG', 'T'),
            ('CCG', 'T'),
            ('GCG', 'T'),
            ('TCG', 'T')]:

            r = {
            'lib':lib,
            'cell':cell,
            'nm':nm,
                #'plate': int(lib.split('-')[-1].replace('pl','')),
            'group':  f'{nm},{context},{cell}',
            'context': f'{context}>{base}',
            'conversion rate':row[context,base]
            }

            samples.append(r)

    plot_table = pd.DataFrame(samples)
    print(plot_table)

    ph = 22
    fig, ax = plt.subplots(figsize=(12,5))
    sns.boxplot(data=plot_table.sort_values('lib'),x='context', y='conversion rate',hue='lib',whis=6, ax=ax)

    #ax = sns.swarmplot(data=plot_table,x='nm', y='conversion rate',hue='lib',)

    plt.legend()
    plt.ylabel('Lambda Conversion rate')
    sns.despine()
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5 ))

    plt.suptitle('Estimated TAPS conversion rate', y=1.05, fontsize=12)
    plt.title(f'Lambda spike-in, >{(1.0-(phred_to_prob(22)))*100 : .2f}% accuracy base calls', fontsize=10)
    plt.tight_layout()
    plt.savefig(f'{prefix}_conversion_rate_phred_{ph}.png', bbox_inches='tight')
    plt.close()


if __name__=='__main__':
    argparser = argparse.ArgumentParser(
      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
      description='Estimate the conversion efficiency of a TAPS converted file. ')
    argparser.add_argument('bams', type=str, nargs='+',help='Input bam files')
    argparser.add_argument('-o', type=str, help="output alias (Will be the prefix of the output files)", required=True)
    argparser.add_argument('-method', type=str, default='nla', help='Molecule class (nla or chic). Use chic when you are not sure or when another other protocol is used.')
    argparser.add_argument('--dedup', action='store_true',help='perform UMI deduplication and consensus calling. Do not use when the UMI\'s are (near) saturated')
    argparser.add_argument('-t', type=int, help='Amount of threads')
    argparser.add_argument('-spikein_name', type=str, help='Name of spikein contig',default='J02459.1')


    argparser.add_argument(
        '-ref',
        type=str,
        default=None,
        help="Path to reference fast (autodected if not supplied)")
    args = argparser.parse_args()


    generate_taps_conversion_stats(args.bams,
                                   args.ref,
                                   prefix=args.o,
                                   method=args.method,
                                   every_fragment_as_molecule=not args.dedup,
                                   spikein_name=args.spikein_name,
                                   n_threads=args.t)