#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
from .statistic import StatisticHistogram
import singlecellmultiomics.pyutils as pyutils
import collections
import pandas as pd
import matplotlib
matplotlib.rcParams['figure.dpi'] = 160
matplotlib.use('Agg')


class MethylationContextHistogram(StatisticHistogram):
    def __init__(self, args):
        StatisticHistogram.__init__(self, args)
        self.context_obs = collections.Counter()  # (bismark_call_tag)=> observations
        self.hasdata = False
    def processRead(self, R1,R2):

        for read in [R2,R2]:
            if read is None:
                continue

            if not read.is_paired or not read.has_tag(
                    'XM') or not read.is_read1 or read.is_duplicate or read.has_tag('RR'):
                return
            self.hasdata = True
            tags = dict(read.tags)
            for tag in 'zhx':
                self.context_obs[tag] += tags.get(f's{tag}', 0)
                self.context_obs[tag.upper()] += tags.get(f's{tag.upper()}', 0)

    def __repr__(self):
        return f'Methylation status.'

    def get_df(self):

        x = self.context_obs

        return pd.DataFrame(
            {'ratio':
             {
                 'CpG': (x['Z'] / (x['z'] + x['Z'])) if x['z'] + x['Z'] > 0 else 0 ,
                 'CHG': (x['X'] / (x['x'] + x['X'])) if x['x'] + x['X'] >0 else 0,
                 'CHH': (x['H'] / (x['h'] + x['H'])) if x['h'] + x['H'] > 0 else 0},
             'methylated': {
                 'CpG': x['Z'],
                 'CHG': x['X'],
                 'CHH': x['H']
             },
             'unmethylated': {
                 'CpG': x['z'],
                 'CHG': x['x'],
                 'CHH': x['h']
             }
             })

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

        df = self.get_df()
        # Methylation ratio plot:
        name = 'methylation_pct'

        (df['ratio'] * 100).plot.bar()
        ax = plt.gca()
        for p in ax.patches:
            ax.annotate(
                f'{p.get_height():.1f}%',
                (p.get_x() + p.get_width() / 2, p.get_height() * 1.005),
                va='bottom', ha='center')
        ax.set_ylim(0, 110)
        ax.set_xlabel('Methylation context')
        ax.set_ylabel('% methylated')
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        plt.show()
        if title is not None:
            plt.title(title)
        plt.tight_layout()
        plt.savefig(target_path.replace('.png', f'.{name}.png'))

        ax.set_yscale('log')
        ax.set_ylim(0.01, 110)

        plt.tight_layout()
        plt.savefig(target_path.replace('.png', f'.{name}.log.png'))
        plt.close()
        ########
        name = 'methylation_absolute'

        (df[['methylated', 'unmethylated']]).plot.bar()
        ax = plt.gca()
        for p in ax.patches:
            ax.annotate(
                f'{p.get_height()}',
                (p.get_x() + p.get_width() / 2, p.get_height() * 1.005),
                va='bottom', ha='center')

        ax.set_xlabel('Methylation context')
        ax.set_ylabel('Bases total')
        plt.tight_layout()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        if title is not None:
            plt.title(title)
        plt.tight_layout()
        plt.savefig(target_path.replace('.png', f'.{name}.png'))

        ax.set_yscale('log')
        plt.tight_layout()
        plt.savefig(target_path.replace('.png', f'.{name}.log.png'))
        plt.close()

    def to_csv(self, path):

        self.get_df().to_csv(path)