#!/usr/bin/env python

# -*- coding: utf-8 -*-

import matplotlib

matplotlib.rcParams['figure.dpi'] = 160

matplotlib.use('Agg')

import matplotlib.pyplot as plt

import pandas as pd

import seaborn as sns

import multiprocessing

from singlecellmultiomics.bamProcessing.bamBinCounts import generate_commands, count_methylation_binned

import argparse

from colorama import Fore, Style

from singlecellmultiomics.utils import dataframe_to_wig

from singlecellmultiomics.methylation import MethylationCountMatrix

from singlecellmultiomics.bamProcessing.bamFunctions import get_reference_from_pysam_alignmentFile

from colorama import Fore,Style

from collections import defaultdict, Counter

from multiprocessing import Pool

from datetime import datetime

import pysam

from singlecellmultiomics.bamProcessing import get_contig_sizes, get_contig_size

from singlecellmultiomics.bamProcessing.bamBinCounts import generate_commands, read_counts

def sample_dict():

    return defaultdict(Counter)

def methylation_to_cut_histogram(args):

    (alignments_path, bin_size, max_fragment_size, \

     contig, start, end, \

     min_mq, alt_spans, key_tags, dedup, kwargs) = args

    distance_methylation = defaultdict(sample_dict) # sample - > distance -> context(ZzHhXx) : obs

    max_dist = 1000

    # Define which reads we want to count:

    known =  set()

    if 'known' in kwargs and kwargs['known'] is not None:

        # Only ban the very specific TAPS conversions:

        try:

            with pysam.VariantFile(kwargs['known']) as variants:

                for record in variants.fetch(contig, start, end):

                    if record.ref=='C' and 'T' in record.alts:

                        known.add( record.pos)

                    if record.ref=='G' and 'A' in record.alts:

                        known.add(record.pos)

        except ValueError:

            # This happends on contigs not present in the vcf

            pass

    p = 0

    start_time = datetime.now()

    with pysam.AlignmentFile(alignments_path, threads=4) as alignments:

        # Obtain size of selected contig:

        contig_size = get_contig_size(alignments, contig)

        if contig_size is None:

            raise ValueError('Unknown contig')

        # Determine where we start looking for fragments:

        f_start = max(0, start - max_fragment_size)

        f_end = min(end + max_fragment_size, contig_size)

        for p, read in enumerate(alignments.fetch(contig=contig, start=f_start,

                                                  stop=f_end)):

            if p%50==0 and 'maxtime' in kwargs and kwargs['maxtime'] is not None:

                if (datetime.now() - start_time).total_seconds() > kwargs['maxtime']:

                    print(f'Gave up on {contig}:{start}-{end}')

                    break

            if not read_counts(read, min_mq=min_mq, dedup=dedup):

                continue

            tags = dict(read.tags)

            for i, (qpos, methylation_pos) in enumerate(read.get_aligned_pairs(matches_only=True)):

                # Don't count sites outside the selected bounds

                if methylation_pos < start or methylation_pos >= end:

                    continue

                call = tags['XM'][i]

                if call=='.':

                    continue

                sample = read.get_tag('SM')

                distance = abs(read.get_tag('DS') - methylation_pos)

                if distance>max_dist:

                    continue

                distance_methylation[sample][(read.is_read1, read.is_reverse, distance)][call] +=1

    return distance_methylation

threads = None

def get_distance_methylation(bam_path,

                                 bp_per_job: int,

                                 min_mapping_qual: int = None,

                                 skip_contigs: set = None,

                                 known_variants: str = None,

                                 maxtime: int = None,

                                 head: int=None,

                                 threads: int = None,

                                **kwargs

                                 ):

    all_kwargs = {'known': known_variants,

            'maxtime': maxtime,

            'threads':threads

            }

    all_kwargs.update(kwargs)

    commands = generate_commands(

        alignments_path=bam_path,

        key_tags=None,

        max_fragment_size=0,

        dedup=True,

        head=head,

        bin_size=bp_per_job,

        bins_per_job= 1, min_mq=min_mapping_qual,

        kwargs=all_kwargs,

        skip_contigs=skip_contigs

    )

    distance_methylation = defaultdict(sample_dict) # sample - > distance -> context(ZzHhXx) : obs

    with Pool(threads) as workers:

        for result in workers.imap_unordered(methylation_to_cut_histogram, commands):

            for sample, data_for_sample in result.items():

                for distance, context_obs in data_for_sample.items():

                    distance_methylation[sample][distance] += context_obs

    return distance_methylation

if __name__ == '__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description="""Extract methylation levels relative to cut site (DS tag) from bam file""")

    argparser.add_argument('bamfile', metavar='bamfile', type=str)

    argparser.add_argument('-bp_per_job', default=5_000_000, type=int, help='Amount of basepairs to be processed per thread per chunk')

    argparser.add_argument('-threads', default=None, type=int, help='Amount of threads to use for counting, None to use the amount of available threads')

    fi = argparser.add_argument_group("Filters")

    fi.add_argument('-min_mapping_qual', default=40, type=int)

    fi.add_argument('-head', default=None, type=int,help='Process the first n bins')

    fi.add_argument('-skip_contigs', type=str, help='Comma separated contigs to skip', default='MT,chrM')

    fi.add_argument('-known_variants',

                           help='VCF file with known variants, will be not taken into account as methylated/unmethylated',

                           type=str)

    og = argparser.add_argument_group("Output")

    og.add_argument('-prefix', default='distance_calls', type=str, help='Prefix for output files')

    args = argparser.parse_args()

    print('Obtaining counts ', end="")

    r = get_distance_methylation(bam_path = args.bamfile,

                                 bp_per_job = args.bp_per_job,

                                 known_variants = args.known_variants,

                                 skip_contigs = args.skip_contigs.split(','),

                                 min_mapping_qual=args.min_mapping_qual,

                                 head = args.head,

                                 threads=args.threads,

    )

    print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    for ctx in 'zhx':

        beta = {}

        met = {}

        un = {}

        for sample, sample_data in r.items():

            beta[sample] = {}

            met[sample] = {}

            un[sample] = {}

            for distance, contexts in sample_data.items():

                if ctx in contexts or ctx.upper() in contexts:

                    beta[sample][distance] = contexts[ctx.upper()]/(contexts[ctx.upper()]+contexts[ctx])

                    met[sample][distance] = contexts[ctx.upper()]

                    un[sample][distance] = contexts[ctx]

        pd.DataFrame(beta).sort_index().T.sort_index().to_csv(f'{args.prefix}_beta_{ctx}.csv')

        pd.DataFrame(beta).sort_index().T.sort_index().to_csv(f'{args.prefix}_beta_{ctx}.pickle.gz')

        pd.DataFrame(met).sort_index().T.sort_index().to_csv(f'{args.prefix}_counts_{ctx.upper()}.csv')

        pd.DataFrame(met).sort_index().T.sort_index().to_csv(f'{args.prefix}_counts_{ctx.upper()}.pickle.gz')

        pd.DataFrame(un).sort_index().T.sort_index().to_csv(f'{args.prefix}_counts_{ctx}.csv')

        pd.DataFrame(un).sort_index().T.sort_index().to_csv(f'{args.prefix}_counts_{ctx}.pickle.gz')

        # Make plots

        beta = {}

        met = {}

        un = {}

        for sample, sample_data in r.items():

            beta[sample] = {}

            met[sample] = {}

            un[sample] = {}

            for distance, contexts in sample_data.items():

                if distance[-1] > 500 or distance[-1] < 4: # Clip in sane region

                    continue

                if ctx in contexts or ctx.upper() in contexts:

                    beta[sample][distance] = contexts[ctx.upper()] / (contexts[ctx.upper()] + contexts[ctx])

                    met[sample][distance] = contexts[ctx.upper()]

                    un[sample][distance] = contexts[ctx]

        beta = pd.DataFrame(beta).sort_index().T.sort_index()

        met = pd.DataFrame(met).sort_index().T.sort_index()

        un = pd.DataFrame(un).sort_index().T.sort_index()

        for mate in [True, False]:

            for strand in [True, False]:

                fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)

                un[mate, strand].sum().rename('Unmethylated').plot(ax=ax1)

                met[mate, strand].sum().rename('Methylated').plot(ax=ax1)

                ax1.set_xlabel('distance to cut')

                ax1.set_ylabel('# molecules')

                ax1.legend()

                (met[mate, strand].sum() / (un[mate, strand].sum() + met[mate, strand].sum())).rename('Beta').plot(

                    ax=ax2)

                # ax2.set_ylim(0,0)

                sns.despine()

                ax1.set_title(f'Mate {"R1" if mate else "R2"}, strand:{"reverse" if strand else "forward"}')

                ax2.set_ylabel('Beta')

                plt.savefig(f'{args.prefix}_{ctx}_{"R1" if mate else "R2"}_{"reverse" if strand else "forward"}.png')

                plt.close('all')