#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import collections
import matplotlib.pyplot as plt
from .statistic import Statistic
import singlecellmultiomics.pyutils as pyutils
import matplotlib
matplotlib.rcParams['figure.dpi'] = 160
matplotlib.use('Agg')


class ReadCount(Statistic):
    def __init__(self, args):
        Statistic.__init__(self, args)
        self.totalMappedReads = collections.Counter()
        self.rawReadCount = None
        self.unmappedReads = collections.Counter()
        self.totalDedupReads = collections.Counter()
        self.totalAssignedSiteReads = collections.Counter({'R1': 0, 'R2': 0})
        self.rejectionReasons = collections.Counter()
        self.demuxReadCount = 0
        self.rawReadCount = 0

    def get_df(self):
        df = pd.DataFrame.from_dict(dict(self))
        # They are paired and only present one time in the dict but are
        # expanded by pandas
        df['Raw reads'] /= 2
        df['Demultiplexed reads'] /= 2  # same

        df = df[['Raw reads',
                 'Demultiplexed reads',
                 'Mapped reads',
                 'AssignedSiteReads',
                 'Deduplicated reads']]
        return df

    def to_csv(self, path):
        self.get_df().to_csv(path)

    def plot(self, target_path, title=None):
        df = self.get_df()  # ,'UnmappedReads']]

        print(df)

        df.plot.bar(figsize=(10, 4)).legend(bbox_to_anchor=(1, 0.98))
        if title is not None:
            plt.title(title)

        plt.subplots_adjust(right=0.6)
        plt.tight_layout()
        plt.savefig(target_path)

        ax = plt.gca()
        ax.set_ylim(1, None)
        ax.set_ylabel('Frequency')
        ax.set_yscale('log')
        plt.tight_layout()
        plt.savefig(target_path.replace('.png', '.log.png'))

        plt.close()

    def processRead(self, R1,R2):

        for read in [R1,R2]:
            if read is None:
                continue
            # Count every read only once
            if read.is_supplementary or read.is_secondary:
                return

            if not read.is_unmapped:

                if read.is_read1:
                    self.totalMappedReads['R1'] += 1
                elif read.is_read2:
                    self.totalMappedReads['R2'] += 1
                else:
                    self.totalMappedReads['R?'] += 1
            else:
                if read.is_read1:
                    self.unmappedReads['R1'] += 1
                elif read.is_read2:
                    self.unmappedReads['R2'] += 1
                else:
                    self.unmappedReads['R?'] += 1

            # Compatibility with picard --TAG_DUPLICATE_SET_MEMBERS
            """
            java -jar `which picard.jar` MarkDuplicates I=sorted.bam O=marked_duplicates.bam M=marked_dup_metrics.txt TAG_DUPLICATE_SET_MEMBERS=1
            """
            if not read.has_tag('MX'):  # Bulk sample
                if not read.is_unmapped:
                    if read.is_duplicate:
                        pass
                    else:
                        if read.is_read1:
                            self.totalDedupReads['R1'] += 1
                        else:
                            self.totalDedupReads['R2'] += 1

            else:
                if not read.is_unmapped:
                    if (read.has_tag('DS') or read.has_tag('GN')) and not read.is_qcfail:
                        if read.is_read1:
                            self.totalAssignedSiteReads['R1'] += 1
                        elif read.is_read2:
                            self.totalAssignedSiteReads['R2'] += 1
                        else:
                            self.totalAssignedSiteReads['R?'] += 1

                    if not read.is_duplicate:

                        if read.is_read1:
                            self.totalDedupReads['R1'] += 1
                        elif read.is_read2:
                            self.totalDedupReads['R2'] += 1
                        else:
                            self.totalDedupReads['R?'] += 1

    def setRawReadCount(self, readCount, paired=True):
        self.rawReadCount = readCount * (2 if paired else 1)

    def setRawDemuxCount(self, readCount, paired=True):
        self.demuxReadCount = readCount * (2 if paired else 1)

    def mappability(self):
        if self.rawReadCount > 0:
            return sum(self.totalMappedReads.values()) / self.rawReadCount
        return None

    def __repr__(self):
        return f"Input:{int(self.rawReadCount)} raw reads, demultiplexed:{self.demuxReadCount}, final mapped: {self.totalMappedReads['R1']} R1 reads, {self.totalMappedReads['R2']} R2 reads, {self.totalMappedReads['R?']} unknown pairing reads, mappability:{self.mappability() if self.mappability() else 'Cannot be determined'}"

    def __iter__(self):
        yield 'Raw reads', self.rawReadCount
        yield 'Demultiplexed reads', self.demuxReadCount
        yield 'Mapped reads', self.totalMappedReads
        yield 'UnmappedReads', self.unmappedReads
        # if self.totalAssignedSiteReads['R1']>0 or
        # self.totalAssignedSiteReads['R2']>0 or
        # self.totalAssignedSiteReads['R?']>0:
        yield 'AssignedSiteReads', self.totalAssignedSiteReads
        yield 'Deduplicated reads', self.totalDedupReads