#!/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.export import dataframe_to_wig
from singlecellmultiomics.methylation import MethylationCountMatrix
from singlecellmultiomics.bamProcessing.bamFunctions import get_reference_from_pysam_alignmentFile
from colorama import Fore,Style


def prefilter(counts, cell_names, min_measurements, min_variance):
    if min_measurements>0:
        counts = counts.loc[:, (counts >= 0).sum() > min_measurements]
    if min_variance>0:
        return counts.loc[:, counts.var() >= min_variance].reindex(cell_names)
    else:
        return counts.reindex(cell_names)


def panda_and_prefilter(args):
    d, args = args # counts_dict ,(cell_names, min_measurements, min_variance)
    return prefilter(pd.DataFrame(d), *args)


def get_methylation_count_matrix(bam_path,
                                 bin_size: int,
                                 bp_per_job: int,
                                 min_samples: int = None,
                                 min_variance: int = None,
                                 min_mapping_qual: int = None,
                                 skip_contigs: set = None,
                                 known_variants: str = None,
                                 maxtime: int = None,
                                 head: int=None,
                                 threads: int = None,
                                 count_reads: bool = True,
                                 **kwargs
                                 ):


    all_kwargs = {'known': known_variants,
            'maxtime': maxtime,
            'single_location':bin_size==1,
            'min_samples':min_samples,
            'min_variance':min_variance,
            'threads':threads,
            'count_reads':count_reads
            }

    all_kwargs.update(kwargs)
    commands = generate_commands(
        alignments_path=bam_path,
        bin_size=bin_size if bin_size!=1 else bp_per_job,
        key_tags=None,
        max_fragment_size=0,
        dedup=True,
        head=head,
        bins_per_job= int(bp_per_job / bin_size) if bin_size!=1 else 1, min_mq=min_mapping_qual,
        kwargs=all_kwargs,
        skip_contigs=skip_contigs
    )


    count_mat = MethylationCountMatrix()
    read_count_mat = dict()

    if threads==1:
        for command in commands:
            result, result_read_counts = count_methylation_binned(command)
            #result.prune(min_samples=min_samples, min_variance=min_variance)
            count_mat.update( result )
            if count_reads:
                read_count_mat.update(result_read_counts)

    else:
        with multiprocessing.Pool(threads) as workers:

            for result,result_read_counts in workers.imap_unordered(count_methylation_binned, commands):
                #result.prune(min_samples=min_samples, min_variance=min_variance)
                count_mat.update(result)
                if count_reads:
                    read_count_mat.update(result_read_counts)

    return count_mat, read_count_mat



if __name__ == '__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="""Extract methylation calls from bam file
    """)
    argparser.add_argument('bamfile', metavar='bamfile', type=str)
    argparser.add_argument('-bin_size', default=500, type=int, help='bin size, set to 1 for single CpG')
    argparser.add_argument('-bp_per_job', default=1_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')
    argparser.add_argument('-threads_agg', default=1, type=int, help='Amount of threads to use for aggregation. Aggregation is very memory intensive, so this amount of threads should probably be lower than -threads')
    argparser.add_argument('-reference', default=None, type=str, help='Path to reference fasta file')

    argparser.add_argument('--mirror_cpg_dyad',action='store_true', help='Count CpG methylation of a single dyad as one position')
    argparser.add_argument('--stranded',action='store_true', help='Perform strand specific methylation calls')


    fi = argparser.add_argument_group("Filters")
    fi.add_argument('-min_variance', default=None, type=float)
    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('-min_samples', default=1, type=int)
    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)

    fi.add_argument('-contexts_to_capture', default=None, type=str, help='Contexts to capture, comma separated')\

    og = argparser.add_argument_group("Output")
    #og.add_argument('-bed', type=str, help='Bed file to write methylation calls to')
    og.add_argument('-default_sample_name', type=str, help='Default sample name to use when SM tag is not available in the read')
    og.add_argument('-wig_beta', type=str, help='WIG file to write mean methylation per bin to')
    og.add_argument('-wig_n_samples', type=str, help='WIG file to write amount of samples covering to')
    og.add_argument('-capture_read_depth', default=None, type=str, help='Capture read depth into this csv file')

    og.add_argument('-betas', type=str, help='CSV or pickle file to write single cell methylation betas to')
    #og.add_argument('-mets', type=str, help='CSV or pickle file to write single cell methylation unmethylated frame to')
    #og.add_argument('-unmets', type=str, help='CSV or pickle file to write single cell methylation frame to')

    og.add_argument('-bismark_tabfile', type=str, help='Tabulated file to write to, contains: chr | start | end | unmethylated_counts | methylated_counts | beta_value')
    og.add_argument('-tabfile', type=str,
                    help='Tabulated file to write to, contains: chr | start | end | unmethylated_counts | methylated_counts | beta_value | variance | n_samples')
    og.add_argument('-distmat', type=str, help='CSV or pickle file to write single cell distance matrix to')
    og.add_argument('-distmat_plot', type=str, help='.PNG file to write distance matrix image to')

    args = argparser.parse_args()

    if args.contexts_to_capture is None:
        contexts_to_capture = None
        context_radius = None
    else:
        assert args.mirror_cpg_dyad is False, '--mirror_cpg_dyad can only be used when no contexts_to_capture are supplied. (CpG mode) '
        contexts_to_capture = args.contexts_to_capture.split(',')
        lengths = set( (len(context) for context in contexts_to_capture ) )
        assert len(lengths)==1, 'Please only supply contexts of the same length'
        context_radius = int((list(lengths)[0] - 1) / 2)
        for context in contexts_to_capture:
            assert context[context_radius]=='C', 'Supply only contexts with a C in the middle'

        assert args.reference is not None, 'Please supply supply a fasta file using -reference'


    if args.stranded:
        assert args.mirror_cpg_dyad is False, '--mirror_cpg_dyad can only be used in non-strand specific mode '
        assert args.wig_beta is None, '-wig_beta cannot be used in stranded mode as WIG files are not strand specific'
        assert args.wig_n_samples is None, '-wig_n_samples cannot be used in stranded mode as WIG files are not strand specific'

    if args.distmat_plot is not None and not args.distmat_plot.endswith('.png'):
        args.distmat_plot += '.png'

    print('Obtaining counts ', end="")
    counts, read_dictionary = get_methylation_count_matrix(bam_path=args.bamfile,
                                          bin_size=args.bin_size,
                                          bp_per_job=args.bp_per_job,
                                          min_samples=args.min_samples,
                                          min_variance=args.min_variance,
                                          known_variants=args.known_variants,
                                          skip_contigs=args.skip_contigs.split(','),
                                          min_mapping_qual=args.min_mapping_qual,
                                          head=args.head,
                                          threads=args.threads,
                                          single_location=(args.bin_size==1),
                                          contexts_to_capture=contexts_to_capture,
                                          context_radius=context_radius,
                                          reference_path=args.reference,
                                          dyad_mode=args.mirror_cpg_dyad,
                                          stranded=args.stranded,
                                          count_reads=args.capture_read_depth is not None

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

    if args.capture_read_depth:
        print('Writing read depth file ', end="")
        pd.DataFrame(read_dictionary).sort_index().sort_index(1).to_csv(args.capture_read_depth)
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")


    print(counts)

    if args.betas is not None:
        print('Writing counts ', end="")
        if args.betas.endswith('.pickle.gz'):
            counts.to_pickle(args.betas)
        else:
            counts.get_frame('beta').to_csv(args.betas)
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    if args.wig_beta is not None:
        # Calculate mean beta value and write to wig:
        print('Writing WIG beta', end="")
        dataframe_to_wig(counts.get_bulk_frame()[['beta']], args.wig_beta, span=args.bin_size)
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    if args.wig_n_samples is not None:
        # Calculate mean beta value and write to wig:
        print('Writing WIG n_samples', end="")
        dataframe_to_wig(counts.get_bulk_frame()[['n_samples']], args.wig_n_samples, span=args.bin_size)
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    if args.distmat is not None or args.distmat_plot is not None:
        print('Calculating distance matrix', end="")
        dmat = counts.get_sample_distance_matrix()
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")
        if args.distmat_plot is not None:
            print('Writing distmat_image', end="")
            try:
                sns.clustermap(dmat)
                plt.tight_layout()
                plt.savefig(args.distmat_plot)
                print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")
            except Exception as e:
                print(e)
                print(f" [ {Fore.RED}FAIL{Style.RESET_ALL} ] ")

        if args.distmat is not None:
            print('Writing distance matrix', end="")
            if args.distmat.endswith('.pickle.gz'):
                dmat.to_pickle(args.distmat)
            else:
                dmat.to_csv(args.distmat, sep='\t')
            print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    if args.bismark_tabfile is not None:
        print('Writing bismark_tabfile matrix', end="")
        bf = counts.get_bulk_frame()

        bf.index.set_names(['chr','start', 'end'] + (['strand'] if args.stranded else []), inplace=True)
        bf[['unmethylated', 'methylated', 'beta']].to_csv(args.bismark_tabfile, sep='\t')
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")

    if args.tabfile is not None:
        print('Writing tabfile', end="")
        if args.tabfile.endswith('.pickle.gz'):
            counts.get_bulk_frame().to_pickle(args.tabfile)
        else:
            counts.get_bulk_frame().to_csv(args.tabfile, sep='\t')
        print(f" [ {Fore.GREEN}OK{Style.RESET_ALL} ] ")