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
import collections

import matplotlib
matplotlib.rcParams['figure.dpi'] = 160
matplotlib.use('Agg')


def human_readable(value, targetDigits=2, fp=0):

    # Float:
    if value < 1 and value > 0:
        return('%.2f' % value)

    if value == 0.0:
        return('0')

    baseId = int(math.floor(math.log10(float(value)) / 3.0))
    suffix = ""
    if baseId == 0:
        sVal = str(round(value, targetDigits))
        if len(sVal) > targetDigits and sVal.find('.'):
            sVal = sVal.split('.')[0]

    elif baseId > 0:

        sStrD = max(0, targetDigits -
                    len(str('{:.0f}'.format((value / (math.pow(10, baseId * 3)))))))

        sVal = ('{:.%sf}' % min(fp, sStrD)).format(
            (value / (math.pow(10, baseId * 3))))
        suffix = 'kMGTYZ'[baseId - 1]
    else:

        sStrD = max(0, targetDigits -
                    len(str('{:.0f}'.format((value * (math.pow(10, -baseId * 3)))))))
        sVal = ('{:.%sf}' % min(fp, sStrD)).format(
            (value * (math.pow(10, -baseId * 3))))
        suffix = 'mnpf'[-baseId - 1]

        if len(sVal) + 1 > targetDigits:
            # :(
            sVal = str(round(value, fp))[1:]
            suffix = ''

    return('%s%s' % (sVal, suffix))


# Visualize the following:
# PER LIBRARY / DEMUX method
# total fragments
# total fragments with correct site
# unique molecules

# 384 well format:

well2index = collections.defaultdict(dict)
index2well = collections.defaultdict(dict)
rows = string.ascii_uppercase[:16]
columns = list(range(1, 25))

for ci in range(1, 385):
    i = ci - 1
    rowIndex = math.floor(i / len(columns))
    row = rows[rowIndex]
    column = columns[i % len(columns)]
    well2index[384][(row, column)] = ci
    index2well[384][ci] = (row, column)


rows96 = string.ascii_uppercase[:8]

columns96 = list(range(1, 13))

for ci in range(1, 97):
    i = ci - 1
    rowIndex = math.floor(i / len(columns96))
    row = rows96[rowIndex]
    column = columns96[i % len(columns96)]
    well2index[96][(row, column)] = ci
    index2well[96][ci] = (row, column)


class PlateStatistic(object):

    def __init__(self, args):
        self.args = args

        self.rawFragmentCount = collections.defaultdict(
            collections.Counter)  # (library, mux) -> cell -> counts
        self.usableCount = collections.defaultdict(
            collections.Counter)  # (library, mux) -> cell -> counts
        self.moleculeCount = collections.defaultdict(
            collections.Counter)  # (library, mux) -> cell -> counts
        self.skipReasons = collections.Counter()

    def to_csv(self, path):
        pd.DataFrame(
            self.moleculeCount).to_csv(
            path.replace(
                '.csv',
                'molecules.csv'))
        pd.DataFrame(
            self.usableCount).to_csv(
            path.replace(
                '.csv',
                'usable_reads.csv'))
        pd.DataFrame(
            self.rawFragmentCount).to_csv(
            path.replace(
                '.csv',
                'raw_fragments.csv'))

    def processRead(self, R1,R2):

        for read in [R1,R2]:

            if read is None:
                continue

            if not read.has_tag('MX'):
                return

            self.rawFragmentCount[(read.get_tag('LY'),
                                   read.get_tag('MX'))][read.get_tag('SM')] += 1

            if read.get_tag('MX').startswith('CS2'):
                if read.has_tag('XT') or read.has_tag('EX'):
                    if read.is_read1:  # We only count reads2
                        return
                    self.usableCount[(read.get_tag('LY'),
                                      read.get_tag('MX'))][read.get_tag('SM')] += 1

                    if read.has_tag('RC') and read.get_tag('RC') == 1:
                        self.moleculeCount[(read.get_tag('LY'), read.get_tag(
                            'MX'))][read.get_tag('SM')] += 1
            else:

                if read.has_tag('DS'):
                    if not read.is_read1:
                        self.skipReasons['Not R1'] += 1
                        return

                    self.usableCount[(read.get_tag('LY'),
                                      read.get_tag('MX'))][read.get_tag('SM')] += 1
                    if not read.is_duplicate:
                        self.moleculeCount[(read.get_tag('LY'), read.get_tag(
                            'MX'))][read.get_tag('SM')] += 1
                else:
                    self.skipReasons['No DS'] += 1
            break


    def __repr__(self):
        return 'Plate statistic'

    def cell_counts_to_dataframe(self, cell_counts, mux, name='raw_reads'):
        df = pd.DataFrame({name: cell_counts})

        offset = 0 # Offset is zero for all protocols since 0.1.12

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

        df['col'] = [index2well[format]
                     [(offset + int(x.rsplit('_')[-1]))][1] for x in df.index]
        df['row'] = [-rows.index(index2well[format]
                                 [(offset + int(x.rsplit('_')[-1]))][0]) for x in df.index]
        df['size'] = (df[name] / np.percentile(df[name], 99) * 200)

        return df

    def __iter__(self):
        for data, name in [
            (self.rawFragmentCount, 'raw_reads'),
            (self.usableCount, 'usable_reads'),
                (self.moleculeCount, 'unique_molecules')]:
            for (library, mux), cellCounts in data.items():
                df = self.cell_counts_to_dataframe(cellCounts, mux, name=name)
                for i, row in df.iterrows():
                    yield i, row

    def plot(self, target_path, title=None):
        for data, name in [
            (self.rawFragmentCount, 'raw_reads'),
            (self.usableCount, 'usable_reads'),
                (self.moleculeCount, 'unique_molecules')]:

            for (library, mux), cellCounts in data.items():

                df = self.cell_counts_to_dataframe(cellCounts, mux, name=name)
                df.plot.scatter(x='col', y='row', s=df['size'],
                                c=[(0.2, 0.2, 0.5, 0.9)]
                                )

                # Annotate the outliers with values:
                ax = plt.gca()
                for ii, row in df.iterrows():
                    if row[name] > 0 and (
                        row[name] < np.percentile(
                            df[name],
                            5) or row[name] > np.percentile(
                            df[name],
                            95)):
                        text = ax.annotate(human_readable(int(row[name])), (row['col'], row['row']),
                                           ha='center', va='baseline', color='w', size=7)
                        text.set_path_effects([path_effects.Stroke(
                            linewidth=3, foreground='black'), path_effects.Normal()])

                plt.yticks(
                    sorted(
                        df['row'].unique())[
                        ::-1],
                    sorted(rows),
                    rotation=0)
                plt.xticks(
                    sorted(
                        df['col'].unique()),
                    sorted(columns),
                    rotation=0)
                plt.title(fr'{name} with ${mux}$ adapter' + f'\n{library}')

                # Create legend:
                #ld = []
                # for x in np.linspace(1, max(df[name]), 4):
            #        size = (x/np.percentile(df[name],99))*200
                #    ld.append( mlines.Line2D([], [], color='blue', marker='.', linestyle='None',
                # markersize=np.sqrt(size), label=f'{int(x)}:{size}'))
                # plt.legend(handles=ld)

                plt.savefig(
                    target_path.replace(
                        '.png',
                        '') +
                    f'{name}_{mux}_{library}.png')
                plt.close()