#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from singlecellmultiomics.statistic import *
import pandas as pd
import matplotlib.pyplot as plt
import os
import sys
import pysam
import collections
import argparse

from singlecellmultiomics.bamProcessing import bam_is_processed_by_program

from colorama import Fore
from colorama import Back
from colorama import Style
import singlecellmultiomics.pyutils as pyutils
from singlecellmultiomics.tagtools import tagtools
from pysamiterators import MatePairIteratorIncludingNonProper
import gzip
import pickle
import subprocess
from glob import glob

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


def select_bam_file(lookup):
    for l in lookup:
        if os.path.exists(l):
            print(f'Found file at {l}')
            return l

    return None


def select_fastq_file(lookup):
    for paths in lookup:
        if isinstance(paths, tuple):
            for path in paths:
                if os.path.exists(path):
                    return paths
        elif os.path.exists(paths):
            return (paths,)
    return None


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Obtain statistics from your libraries')
    argparser.add_argument(
        'libraries',
        type=str,
        nargs='*',
        help="either a library structured folder, or the tagged BAM file")
    argparser.add_argument('-t', type=str, default='all-stats',
                           help="type of staistics to produce. options are \
            'demult-stats', 'meth-stats', 'chic-stats' and 'all-stats'")
    argparser.add_argument('-o', type=str, help="output file prefix")
    argparser.add_argument(
        '--plotsOnly',
        action='store_true',
        help="only make plots")
    argparser.add_argument(
        '--fatal',
        action='store_true',
        help="Fatal error on any issue")
    argparser.add_argument(
        '--sl',
        action='store_true',
        help="Show lookup paths")


    argparser.add_argument(
        '--tablesOnly',
        action='store_true',
        help="only make tables")

    argparser.add_argument('-head', type=int)
    argparser.add_argument(
        '-tagged_bam',
        type=str,
        help='Alias of subpath to tagged bam file. For example /tagged/sorted.bam')
    argparser.add_argument('--v', action='store_true')
    argparser.add_argument('--nort', action='store_true')
    argparser.add_argument('--nolorenz', action='store_true')
    args = argparser.parse_args()

    for library in args.libraries:
        if library.endswith('.bam'):
            # the library is a bam file..
            bamFile = os.path.abspath(library)
            library = os.path.dirname(os.path.abspath(bamFile))
            library_name = os.path.basename(os.path.abspath(bamFile))
            print("Bam file was supplied:")
            print(bamFile)
        else:
            print("A library was supplied, automatically detecting files ..")
            bamFile = None
            library_name = os.path.basename(library)
        rc = ReadCount(args)  # Is also mappability

        statistics = [
            rc,
            FragmentSizeHistogram(args),
            RejectionReasonHistogram(args),
            MappingQualityHistogram(args),
            OversequencingHistogram(args),
            CellReadCount(args)
        ]


        full_file_statistics = []

        if not args.nolorenz:
            full_file_statistics.append( Lorenz(args) )

        if(args.t in ['meth-stats', 'all-stats']):
            statistics.extend([
                MethylationContextHistogram(args),
                ConversionMatrix(args)
            ])

        if(args.t in ['chic-stats', 'all-stats']):
            statistics.extend([ScCHICLigation(args)])

        if(args.t in ['demult-stats', 'all-stats']):
            statistics.extend([
                TrimmingStats(args),
                AlleleHistogram(args),
                DataTypeHistogram(args),
                TagHistogram(args),
                PlateStatistic(args),
                PlateStatistic2(args)
            ])

        demuxFastqFilesLookup = [
            (f'{library}/demultiplexedR1.fastq.gz',
             f'{library}/demultiplexedR2.fastq.gz'),
            (f'{library}/demultiplexedR1_val_1.fq.gz',
             f'{library}/demultiplexedR2_val_2.fq.gz'),
            (f'{library}/demultiplexedR1_val_1.fq',
             f'{library}/demultiplexedR2_val_2.fq')]

        rejectFilesLookup = [
            (f'{library}/rejectsR1.fastq.gz', f'{library}/rejectsR2.fastq.gz'),
            (f'{library}/rejectsR1.fastq', f'{library}/rejectsR2.fastq'),
            (f'{library}/rejectsR1.fq.gz', f'{library}/rejectsR2.fq.gz'),
            (f'{library}/rejectsR1.fq', f'{library}/rejectsR2.fq'),
        ]

        taggedFilesLookup = [
            f'{library}/tagged.bam',
            f'{library}/tagged/tagged.bam',
            f'{library}/tagged/marked_duplicates.bam',
            f'{library}/tagged/resorted.featureCounts.bam',
            f'{library}/tagged/STAR_mappedAligned.sortedByCoord.out.featureCounts.bam',
            f'{library}/tagged/STAR_mappedAligned.sortedByCoord.out.bam',
            f'{library}/tagged/sorted.bam']
        if args.tagged_bam:
            taggedFilesLookup = [
                library + '/' + args.tagged_bam] + taggedFilesLookup

        if args.sl:
            print(f'{Style.BRIGHT}Demux file lookup paths{Style.RESET_ALL}')
            print(demuxFastqFilesLookup)
            print(f'{Style.BRIGHT}Reject fastq file lookup paths{Style.RESET_ALL}')
            print(rejectFilesLookup)
            print(f'{Style.BRIGHT}Tagged bam lookup paths{Style.RESET_ALL}')
            print(taggedFilesLookup)

        if 'cluster' in library:
            continue
        print(f'{Style.BRIGHT}Library {library}{Style.RESET_ALL}')
        # Check if the bam file is present
        if bamFile is None:
            bamFile = select_bam_file(taggedFilesLookup)

        if bamFile is None:
            # Perform glob expansion
            bams = list(glob(f'{library}/*.bam'))+list(glob(f'{library}/*/*.bam'))
            for bam_path in bams:
                print(f"Trying {bam_path}",end="\t")
                try:
                    with pysam.AlignmentFile(bam_path) as a:
                        if bam_is_processed_by_program(a, program='bamtagmultiome'):
                            bamFile = bam_path
                            print("[TAGGED]")
                            break
                        else:
                            print("[NOT TAGGED]")
                except Exception as e:
                    print(f"[ERROR] {e}")

        if bamFile is None:
            print(f'{Fore.RED}BAM FILE MISSING {library}{Style.RESET_ALL}')
            exit()
        else:
            print(f'{Fore.GREEN}Bam file at {bamFile}{Style.RESET_ALL}')

        demuxFastqFiles = select_fastq_file(demuxFastqFilesLookup)
        rejectFastqFiles = select_fastq_file(rejectFilesLookup)

        print("Selected files:")
        print(f'demultiplexed reads: {demuxFastqFiles}')
        print(f'rejected reads: {rejectFastqFiles}')
        print(f'tagged bam: {bamFile}')

        demuxReads = None
        rejectedReads = None
        if demuxFastqFiles is not None:
            firstMate = demuxFastqFiles[0]
            print(f'\tDemuxed > {firstMate}')
            if firstMate.endswith('.gz'):
                demuxReads = pyutils.wccountgz(firstMate) / 4
            else:
                demuxReads = pyutils.wccount(firstMate) / 4

        if rejectFastqFiles is not None:
            firstMate = rejectFastqFiles[0]
            print(f'\tRejects > {firstMate}')
            if firstMate.endswith('.gz'):
                rejectedReads = pyutils.wccountgz(firstMate) / 4
            else:
                rejectedReads = pyutils.wccount(firstMate) / 4

        if demuxReads is not None:
            rc.setRawDemuxCount(demuxReads, paired=True)

            if rejectedReads is not None:
                rc.setRawReadCount(rejectedReads + demuxReads, paired=True)

        if bamFile is not None and os.path.exists(bamFile):
            print(f'\tTagged > {bamFile}')
            with pysam.AlignmentFile(bamFile) as f:

                for i, (R1,R2) in enumerate(MatePairIteratorIncludingNonProper(f)):
                    for statistic in statistics:
                        statistic.processRead(R1,R2)
                    if args.head is not None and i >= (args.head - 1):
                        break
        else:
            print(f'Did not find a bam file at {bamFile}')

        statDict = {}

        if os.path.exists(
                f'{library}/tagged/STAR_mappedAligned.sortedByCoord.out.featureCounts.bam'):
            cmd = f'samtools view {bamFile} -F 4 -f 64 | cut -f 1 | sort | uniq | wc -l'
            out = subprocess.Popen(cmd,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.STDOUT,
                                   shell=True
                                   ).communicate()[0]
            read1mapped = int(out.partition(b' ')[0])
            cmd = f'samtools view {bamFile} -F 4 -f 128 | cut -f 1 | sort | uniq | wc -l'
            out = subprocess.Popen(cmd,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.STDOUT,
                                   shell=True
                                   ).communicate()[0]
            read2mapped = int(out.partition(b' ')[0])

            rc.totalMappedReads['R1'] = read1mapped
            rc.totalMappedReads['R2'] = read2mapped
            # Deduplicated reads have RC:i:1 set, -f 64 selects for R2
            cmd = f'samtools view {bamFile} -F 4 -f 64 | grep RC:i:1 | cut -f 1 | sort | uniq | wc -l'
            out = subprocess.Popen(cmd,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.STDOUT,
                                   shell=True
                                   ).communicate()[0]
            read1mappeddedup = int(out.partition(b' ')[0])
            # Deduplicated reads have RC:i:1 set, -f 128 selects for R2
            cmd = f'samtools view {bamFile} -F 4 -f 128 | grep RC:i:1 |  cut -f 1 | sort | uniq | wc -l'
            out = subprocess.Popen(cmd,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.STDOUT,
                                   shell=True
                                   ).communicate()[0]
            read2mappeddedup = int(out.partition(b' ')[0])

            rc.totalDedupReads['R1'] = read1mappeddedup
            rc.totalDedupReads['R2'] = read2mappeddedup

        for statistic in statistics:
            try:
                print(f'\t{statistic.__class__.__name__}')
                print(f'\t\t{statistic}\n')
                statDict[statistic.__class__.__name__] = dict(statistic)
                print(dict(statistic))
            except Exception as e:
                if args.v:
                    print(e)
                if args.fatal:
                    raise

        if bamFile is not None:
            for statistic in full_file_statistics:
                try:
                    with pysam.AlignmentFile(bamFile ) as alignments:
                        statistic.process_file(alignments)
                except Exception as e:
                    if args.v:
                        print(e)
                    if args.fatal:
                        raise



        # Make plots:
        if args.o is None:
            plot_dir = f'{library}/plots'
            table_dir = f'{library}/tables'
            statFile = f'{library}/statistics.pickle.gz'
        else:
            plot_dir = f'{args.o}/plots'
            table_dir = f'{args.o}/tables'
            statFile = f'{args.o}/statistics.pickle.gz'

        if not args.tablesOnly:

            if not os.path.exists(plot_dir):
                os.makedirs(plot_dir)
            for statistic in statistics + full_file_statistics:
                if not hasattr(statistic, 'plot'):
                    print(
                        f'Not making a plot for {statistic.__class__.__name__} as no plot method is defined')
                    continue
                try:
                    statistic.plot(
                        f'{plot_dir}/{statistic.__class__.__name__}.png',
                        title=library_name)
                except Exception as e:
                    if args.v:
                        import traceback
                        traceback.print_exc()

        # Make tables:
        if not args.plotsOnly:
            with gzip.open(statFile, 'wb') as f:
                pickle.dump(statDict, f)
            if os.path.exists(statFile):
                with gzip.open(statFile, 'rb') as f:
                    try:
                        statDict.update(pickle.load(f))
                    except Exception as e:
                        if args.fatal:
                            raise


            if not os.path.exists(table_dir):
                os.makedirs(table_dir)
            for statistic in statistics + full_file_statistics:
                if not hasattr(statistic, 'to_csv'):
                    print(
                        f'Not making a table for {statistic.__class__.__name__} as to_csv method is not defined')
                    continue
                try:
                    statistic.to_csv(
                        f'{table_dir}/{statistic.__class__.__name__}_{library_name}.csv')
                except Exception as e:
                    if args.fatal:
                        raise
                    if args.v:
                        import traceback
                        traceback.print_exc()


        # Make RT reaction plot:
        if bamFile is not None and os.path.exists(bamFile):
            if not args.nort:
                os.system(
                    f"bamPlotRTstats.py {bamFile} -head 2_000_000 --notstrict -o {plot_dir}/RT_")