#!/usr/bin/env python3

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

import seaborn as sns

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 ConversionMatrix(StatisticHistogram):

    def __init__(self, args, process_reads=200_000):

        StatisticHistogram.__init__(self, args)

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

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

        self.stranded_base_conversions = collections.defaultdict(

            collections.Counter)

        self.processed_reads = 0

        self.process_reads = process_reads

    def processRead(self, R1,R2):

        for read in [R1,R2]:

            if read is None:

                continue

            if self.processed_reads >= self.process_reads or read is None or read.is_unmapped or read.mapping_quality < 30:

                return

            self.processed_reads += 1

            try:

                for index, reference_pos, reference_base in read.get_aligned_pairs(

                        with_seq=True, matches_only=True):

                    query_base = read.seq[index]

                    reference_base = reference_base.upper()

                    if reference_base != 'N' and query_base != 'N':

                        k = (

                            query_base, 'R1' if read.is_read1 else (

                                'R2' if read.is_read2 else 'R?'), ('forward', 'reverse')[

                                read.is_reverse])

                        self.base_obs[reference_base][k] += 1

                        if reference_base != query_base:

                            self.conversion_obs[reference_base][k] += 1

                            if read.has_tag('RS'):

                                k = (

                                    query_base,

                                    'R1' if read.is_read1 else (

                                        'R2' if read.is_read2 else 'R?'),

                                    read.get_tag('RS'))

                                self.stranded_base_conversions[reference_base][k] += 1

            except ValueError: # Fails when the MD tag is not present

                continue

    def __repr__(self):

        return f'Observed base conversions'

    def get_df(self):

        return pd.DataFrame(

            self.base_obs).fillna(0).sort_index(

            axis=0).sort_index(

            axis=1)

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

        df = pd.DataFrame(

            self.conversion_obs).fillna(0).sort_index(

            axis=0).sort_index(

            axis=1)

        cm = sns.clustermap(df, row_cluster=True, col_cluster=True)

        ax = cm.ax_heatmap

        ax.set_yticklabels(ax.get_yticklabels(), rotation=0)

        cm.ax_heatmap.set_title('Raw conversions')

        cm.ax_heatmap.set_xlabel('reference base')

        cm.ax_heatmap.set_ylabel('sequenced base')

        plt.subplots_adjust(left=0, right=0.8)

        if title is not None:

            cm.ax_heatmap.set_title(title)

        #

        plt.savefig(target_path.replace('.png', f'.conversions.png'))

        plt.close()

        df = self.get_df()

        cm = sns.clustermap(df, row_cluster=False, col_cluster=False)

        ax = cm.ax_heatmap

        ax.set_yticklabels(ax.get_yticklabels(), rotation=0)

        cm.ax_heatmap.set_title('Raw conversions')

        cm.ax_heatmap.set_xlabel('reference base')

        cm.ax_heatmap.set_ylabel('sequenced base')

        #plt.tight_layout(pad=0.4, w_pad=0.8, h_pad=1.0)

        plt.subplots_adjust(left=0, right=0.8)

        if title is not None:

            cm.ax_heatmap.set_title(title)

        plt.savefig(target_path.replace('.png', f'.base_obs.png'))

        plt.close()

        if len(self.stranded_base_conversions) == 0:

            return

        df = pd.DataFrame(

            self.stranded_base_conversions).fillna(0).sort_index(

            axis=0).sort_index(

            axis=1)

        cm = sns.clustermap(df, row_cluster=False, col_cluster=False)

        ax = cm.ax_heatmap

        ax.set_yticklabels(ax.get_yticklabels(), rotation=0)

        cm.ax_heatmap.set_title('Raw conversions')

        cm.ax_heatmap.set_xlabel('reference base')

        cm.ax_heatmap.set_ylabel('sequenced base')

        #plt.tight_layout(pad=0.4, w_pad=0.8, h_pad=1.0)

        plt.subplots_adjust(left=0, right=0.8)

        if title is not None:

            cm.ax_heatmap.set_title(title)

        plt.savefig(

            target_path.replace(

                '.png',

                f'.RS_TAG_strand_conversions.png'))

        plt.close()

    def to_csv(self, path):

        self.get_df().to_csv(path)