#!/usr/bin/env python3

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

from matplotlib.ticker import MaxNLocator

import matplotlib.pyplot as plt

from .statistic import StatisticHistogram

import singlecellmultiomics.pyutils as pyutils

import collections

import pandas as pd

import seaborn as sns

import matplotlib

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

matplotlib.use('Agg')

class ScCHICLigation():

    def __init__(self, args):

        # cell -> { A_start: count, total_cuts: count }

        self.per_cell_a_obs = collections.defaultdict(collections.Counter)

        # cell -> { TA_start: count, total_cuts: count }

        self.per_cell_ta_obs = collections.defaultdict(collections.Counter)

    def processRead(self, R1,R2):

        if R1 is None:

            return

        read = R1

        if read.has_tag('RZ') and not read.is_duplicate and read.is_read1:

            sample = read.get_tag('SM')

            first = read.get_tag('RZ')[0]

            if read.get_tag('RZ') == 'TA':

                self.per_cell_ta_obs[sample]['TA_start'] += 1

            if first == 'A':

                self.per_cell_a_obs[sample]['A_start'] += 1

            self.per_cell_ta_obs[sample]['total'] += 1

            self.per_cell_a_obs[sample]['total'] += 1

    def __repr__(self):

        return 'ScCHICLigation: no description'

    def __iter__(self):

        for cell, cell_data in self.per_cell_ta_obs.items():

            yield cell_data['total'],  cell_data['TA_start'] / cell_data['total']

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

        ########### TA ###########

        fig, ax = plt.subplots(figsize=(4, 4))

        x = []

        y = []

        for cell, cell_data in self.per_cell_ta_obs.items():

            x.append(cell_data['total'])

            y.append(cell_data['TA_start'] / cell_data['total'])

        ax.scatter(x, y, s=3,c='k')

        ax.set_xscale('log')

        if title is not None:

            ax.set_title(title)

        ax.set_ylabel("Fraction unique cuts starting with TA")

        ax.set_xlabel("# Molecules")

        ax.set_xlim(1, None)

        ax.set_ylim(-0.1, 1.05)

        sns.despine()

        plt.tight_layout()

        plt.savefig(target_path.replace('.png', '.TA.png'))

        plt.close()

        ########### A ###########

        fig, ax = plt.subplots(figsize=(4, 4))

        x = []

        y = []

        for cell, cell_data in self.per_cell_ta_obs.items():

            x.append(cell_data['total'])

            y.append(cell_data['A_start'] / cell_data['total'])

        ax.scatter(x, y, s=3,c='k')

        ax.set_xscale('log')

        if title is not None:

            ax.set_title(title)

        ax.set_ylabel("Fraction unique cuts starting with A")

        ax.set_xlabel("# Molecules")

        ax.set_xlim(1, None)

        ax.set_ylim(-0.1, 1.05)

        plt.tight_layout()

        sns.despine()

        plt.savefig(target_path.replace('.png', '.A.png'))

        plt.close()

    def to_csv(self, path):

        pd.DataFrame(

            self.per_cell_ta_obs).sort_index().to_csv(

            path.replace(

                '.csv',

                'TA_obs_per_cell.csv'))