import matplotlib.patheffects as path_effects

import math

import string

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

from .statistic import StatisticHistogram

import singlecellmultiomics.pyutils as pyutils

from collections import defaultdict, Counter

from singlecellmultiomics.utils.plotting import plot_plate

import matplotlib

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

matplotlib.use('Agg')

def counter_defaultdict():

    return defaultdict(Counter)

class PlateStatistic2(object):

    def __init__(self, args):

        self.args = args

        self.stats = defaultdict(counter_defaultdict)

    def to_csv(self, path):

        export = {}

        for library, libd in self.stats.items():

            for (stat,metric), data in  libd.items():

                print(data)

                export[library,stat,metric] = data

        pd.DataFrame(

            export).to_csv(

            path.replace(

                '.csv',

                f'plate_metrics.csv'))

    def processRead(self, R1,R2=None, default_lib=None):

        for read in [R1,R2]:

            if read is None:

                continue

            cell_index = int(read.get_tag('bi'))

            if read.has_tag('MX'):

                mux= read.get_tag('MX')

                format = 384 if ('384' in mux or mux.startswith('CS2')) else 96

                if format==384:

                    n_rows = 16

                    n_cols = 24

                else:

                    n_rows = 8

                    n_cols = 12

            else:

                n_rows = 16

                n_cols = 24

            if default_lib is None:

                if read.has_tag('LY'):

                    library = read.get_tag('LY')

                else:

                    library = '_'

            else:

                library = default_lib

            cell_index = int(cell_index)-1

            row = int(cell_index/n_cols)

            col = cell_index - row*n_cols

            cell= read.get_tag('SM')

            self.stats[library][('total mapped','# reads')][(row,col, cell)] += 1

            if read.is_qcfail:

                self.stats[library]['qcfail', '# reads'][(row,col,cell)] += 1

                continue

            if read.is_duplicate:

                continue

            self.stats[library][('total mapped','# molecules')][(row,col,cell)] += 1

            if read.has_tag('lh') and  read.get_tag('lh') in ('TA','TT','AA'):

                self.stats[library][read.get_tag('lh'), 'ligated molecules'][(row,col,cell)] += 1

            break

    def __repr__(self):

        return 'Plate statistic'

    def __iter__(self):

        for library, libd in self.stats.items():

            for name, data in libd.items():

                for k,v in data.items():

                    yield (library,name,k),v

    def plot(self, target_path, title=None):

        for library, libd in self.stats.items():

            for (stat,metric), data in  libd.items():

                for log in [True,False]:

                    fig, ax, cax = plot_plate(data, log=log, cmap_name='viridis',vmin=1 if  log else 0)

                    cax.set_ylabel(metric)

                    ax.set_title(f'{stat}, {metric}')

                    plt.savefig(

                        target_path.replace(

                            '.png',

                            '') +

                        f'_{library}_{stat}_{metric.replace("# ","n_")}.{"log" if log else "linear"}.png', bbox_inches='tight', dpi=250)

                    plt.close()

            # Create oversequencing plot: (reads / molecule)

            overseq = {}

            for coordinate, n_molecules in libd[('total mapped','# molecules')].items():

                n_reads = libd[('total mapped','# reads')].get(coordinate,0)

                overseq[coordinate] = (n_reads/n_molecules) if n_reads>0 else 0

            fig, ax, cax = plot_plate(overseq, log=False, cmap_name='viridis', vmin=np.percentile(list(overseq.values()),5))

            cax.set_ylabel('reads / molecule')

            ax.set_title(f'Oversequencing')

            plt.savefig(

                target_path.replace(

                    '.png',

                    '') +

                f'_{library}_oversequencing.png', bbox_inches='tight', dpi=250)

            plt.close()

            # Create TA ratio plot: (ta molecules / total molecules)

            ta_ratio = {}

            for coordinate, n_ta_molecules in libd[('TA','ligated molecules')].items():

                n_molecules = libd[('total mapped','# molecules')].get(coordinate,0)

                ta_ratio[coordinate] = (n_ta_molecules/n_molecules) if n_molecules>0 else 0

            fig, ax, cax = plot_plate(ta_ratio, log=False, cmap_name='viridis', vmin=np.percentile(list(ta_ratio.values()),5))

            cax.set_ylabel('TA molecules / molecule')

            ax.set_title(f'TA fraction')

            plt.savefig(

                target_path.replace(

                    '.png',

                    '') +

                f'_{library}_ta_fraction.png', bbox_inches='tight', dpi=250)

            plt.close()