#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import collections
import glob
import sys
from colorama import Fore
from colorama import Back
from colorama import Style
import importlib
import inspect
import traceback
import singlecellmultiomics.modularDemultiplexer.demultiplexModules as dm
import singlecellmultiomics.fastqProcessing.fastqIterator as fastqIterator
from singlecellmultiomics.modularDemultiplexer.baseDemultiplexMethods import NonMultiplexable, IlluminaBaseDemultiplexer
import logging


class DemultiplexingStrategyLoader:
    def __init__(
            self,
            barcodeParser,
            moduleSearchDir='demultiplexModules',
            indexParser=None,
            ignoreMethods=None,
            only_detect_methods=None, #
            indexFileAlias=None):
        package = f'singlecellmultiomics.modularDemultiplexer.{moduleSearchDir}'
        moduleSearchPath = os.path.join(
            os.path.dirname(
                os.path.realpath(__file__)),
            moduleSearchDir).replace(
            '/./',
            '/')
        self.barcodeParser = barcodeParser
        self.indexParser = indexParser
        self.only_detect_methods = only_detect_methods
        moduleSearchPath = moduleSearchPath

        #print(f'{Style.DIM}Current script location: {__file__}')
        #print(f'Searchdir: {moduleSearchDir}')
        #print(f'Looking for modules in {moduleSearchPath}{Style.RESET_ALL}')
        self.demultiplexingStrategies = []
        self.demux_classes = [
            IlluminaBaseDemultiplexer,

            dm.CELSeq1_c8_u4,
            dm.CELSeq2_c8_u6,
            dm.CELSeq2_c8_u6_NH,
            dm.CELSeq2_c8_u8,
            dm.CELSeq2_c8_u8_NNLAIII,

            dm.CELSeq2_c8_u6_swapped_reads,

            dm.NLAIII_384w_c8_u3,
            dm.NLAIII_96w_c8_u3,
            dm.Nla_384w_u8_c8_ad3_is15,
            dm.NLAIII_384w_c8_u3_SINGLE_END,
            dm.NLAIII_96w_c8_u3_SINGLE_END,

            dm.SCCHIC_384w_c8_u3,
            dm.SCCHIC_384w_c8_u3_cs2,
            dm.SCCHIC_384w_c8_u3_pdt,
            dm.SCCHIC_384w_c8_u3_direct_ligation,
            dm.SCCHIC_384w_c8_u3_direct_ligation_SINGLE_END,

            dm.MSPJI_c8_u3,
            dm.ScartraceR2,
            dm.ScartraceR1,
            dm.ScartraceR2RP4,

            dm.chrom10x_c16_u12,

            dm.DamID2,
            dm.DamID2_c8_u3_cs2,
            dm.DamID2_SCA,
            dm.DamID2andT_SCA,
            dm.DamID2andT_SCA6,
            dm.DamID2_NO_OVERHANG


        ]
        for c in self.demux_classes:


            initialised_demux = c(
                barcodeFileParser=barcodeParser,
                indexFileParser=indexParser,
                indexFileAlias=indexFileAlias)

            if self.only_detect_methods is not None:

                if initialised_demux.shortName in self.only_detect_methods:
                    print(f'Only loading {initialised_demux.shortName}')
                else:
                    continue

            self.demultiplexingStrategies.append(initialised_demux)


    def getSelectedStrategiesFromStringList(self, strList, verbose=True):
        selectedStrategies = []

        resolved = {part: False for part in strList}
        for strategy in self.demultiplexingStrategies:
            if strategy.shortName in strList:
                selectedStrategies.append(strategy)
                if verbose:
                    print('Selected strategy %s' % strategy)
                resolved[strategy.shortName] = True

        if any([v is False for v in resolved.values()]):
            for strat in strList:
                if resolved[strat] is False:
                    print(f'{Fore.RED}Could not resolve {strat}{Style.RESET_ALL}')
                    print('Available:')
                    for s in self.demultiplexingStrategies:
                        print(s.shortName)
                    raise ValueError(f'Strategy {strat} not found')

            raise ValueError()
        return selectedStrategies

    def list(self):
        print(f"{Style.BRIGHT}Available demultiplexing strategies:{Style.RESET_ALL}")
        #print('Name, alias, will be auto detected, description')
        for strategy in self.demultiplexingStrategies:

            try:
                print(
                    f'{Style.BRIGHT}{strategy.shortName}{Style.RESET_ALL}\t{strategy.longName}\t' +
                    (
                        f'{Fore.GREEN}Will be autodetected' if strategy.autoDetectable else f'{Fore.RED}Will not be autodetected') +
                    Style.RESET_ALL +
                    Style.DIM +
                    f' {strategy.barcodeFileParser.getTargetCount(strategy.barcodeFileAlias) if hasattr(strategy,"barcodeFileParser") else "NA"} targets\n ' +
                    Style.DIM +
                    strategy.description +
                    '\n' +
                    strategy.getParserSummary() +
                    Style.RESET_ALL +
                    '\n')
            except Exception as e:
                print(
                    f"{Fore.RED}{Style.BRIGHT}Error in: {strategy.shortName}\nException: {e}{Style.RESET_ALL}\nTraceback for the error:\n")
                import traceback
                traceback.print_exc()
                from os import stat
                from pwd import getpwuid
                try:
                    modulePath = sys.modules[strategy.__module__].__file__

                    print(
                        f'Contact {Style.BRIGHT}%s{Style.RESET_ALL} for help\n' %
                        getpwuid(
                            stat(modulePath).st_uid).pw_name)
                    print(
                        'The error only affects this module.\nProceeding to load more modules...\n')
                except Exception as e:
                    pass

    def getAutodetectStrategies(self):
        return [
            strategy for strategy in self.demultiplexingStrategies if strategy.autoDetectable]

    def getDemultiplexingSelectedStrategies(self):
        if self.selectedStrategies is None:
            raise ValueError('No strategies selected')
        return self.selectedStrategies

    def demultiplex(
            self,
            fastqfiles,
            maxReadPairs=None,
            strategies=None,
            library=None,
            targetFile=None,
            rejectHandle=None,
            log_handle=None,
            probe=None
            ):

        useStrategies = strategies if strategies is not None else self.getAutodetectStrategies()
        strategyYields = collections.Counter()
        processedReadPairs = 0
        baseDemux = IlluminaBaseDemultiplexer(
            indexFileParser=self.indexParser,
            barcodeParser=self.barcodeParser,
            probe=probe)

        for p, reads in enumerate(
                fastqIterator.FastqIterator(*fastqfiles)):
            processedReadPairs = p+1

            for strategy in useStrategies:
                try:
                    recodedRecords = strategy.demultiplex(
                        reads, library=library, probe=probe)

                    if targetFile is not None:
                        targetFile.write(recodedRecords)

                except NonMultiplexable as reason:
                    # print('NonMultiplexable')
                    if rejectHandle is not None:

                        try:
                            to_write = baseDemux.demultiplex(
                                reads, library=library, reason=reason)
                            rejectHandle.write(to_write)

                        except NonMultiplexable as e:
                            # we cannot read the header of the read..
                            reads = [
                                '\n'.join(
                                    (read.header +
                                     f';RR:{reason};Rr:{e}',
                                     read.sequence,
                                     read.plus,
                                     read.qual)) for read in reads]
                            rejectHandle.write(reads)

                    continue
                except Exception as e:
                    if probe:
                        continue
                    traceback_str = str(traceback.format_exc())
                    print(traceback_str)
                    print(
                        f'{Fore.RED}Fatal error. While demultiplexing strategy {strategy.longName} yielded an error, the error message was: {e}')
                    print('The read(s) causing the error looked like this:')
                    for read in reads:
                        print(str(read))
                    print(Style.RESET_ALL)
                    if log_handle is not None:
                        log_handle.write(
                            f"Error occured using {strategy.longName}\n{traceback_str}\n")
                # print(recodedRecord)
                strategyYields[strategy.shortName] += 1
            if (maxReadPairs is not None and (
                    processedReadPairs) >= maxReadPairs):
                break
        # write yields to log file if applicable:
        if log_handle is not None:
            log_handle.write(f'processed {processedReadPairs} read pairs\n')
            log_handle.write(f'Reads obtained per protocol\n')
            log_handle.write(f'Strategy\tReads\n')
            for strategy, used_reads in strategyYields.items():
                log_handle.write(f'{strategy}\t{used_reads}\n')
        return processedReadPairs, strategyYields

    def detectLibYields(
            self,
            libraries,
            strategies=None,
            testReads=100000,
            maxAutoDetectMethods=1,
            minAutoDetectPct=5,
            verbose=False):
        libYields = dict()

        for lib, lanes in libraries.items():
            for lane, readPairs in lanes.items():

                for readPair in readPairs:
                    if len(readPairs) == 1:
                        processedReadPairs, strategyYields = self.demultiplex(
                            [ readPairs[
                                list(readPairs.keys())[0]
                                ][0] ], maxReadPairs=testReads, strategies=strategies, probe=True)
                    elif len(readPairs) == 2:
                        processedReadPairs, strategyYields = self.demultiplex(
                            (readPairs['R1'][0], readPairs['R2'][0]), maxReadPairs=testReads, strategies=strategies, probe=True)
                    else:
                        raise ValueError('Error: %s' % readPairs.keys())

                if verbose:
                    print(f'Report for {lib}:')
                    self.strategyYieldsToFormattedReport(
                        processedReadPairs,
                        strategyYields,
                        maxAutoDetectMethods=maxAutoDetectMethods,
                        minAutoDetectPct=minAutoDetectPct)
                libYields[lib] = {
                    'processedReadPairs': processedReadPairs,
                    'strategyYields': strategyYields}
                break
        return processedReadPairs, libYields

    def strategyYieldsToFormattedReport(
            self,
            processedReadPairs,
            strategyYields,
            selectedStrategies=None,
            maxAutoDetectMethods=1,
            minAutoDetectPct=5):
        print(
            f'Analysed {Style.BRIGHT}{processedReadPairs}{Style.RESET_ALL} read pairs')

        if selectedStrategies is None:
            selectedStrategies = {}
        #selectedStrategies = self.selectedStrategiesBasedOnYield(processedReadPairs, strategyYields)
        for i, (strategy, strategyYield) in enumerate(
                strategyYields.most_common()):
            yieldRatio = strategyYield / (0.001 + processedReadPairs)
            print(
                (
                    Style.BRIGHT +
                    Fore.GREEN if (
                        (strategy in selectedStrategies) or i < maxAutoDetectMethods) else (
                        Fore.YELLOW if yieldRatio *
                        100 >= minAutoDetectPct else Style.DIM)) +
                f'\t {strategy}:%.2f%%{Style.RESET_ALL}' %
                (100.0 *
                 yieldRatio))

    def selectedStrategiesBasedOnYield(
            self,
            processedReadPairs,
            strategyYields,
            maxAutoDetectMethods=1,
            minAutoDetectPct=0.05):
        selectedStrategies = []
        for strategy, strategyYield in strategyYields.most_common(
                maxAutoDetectMethods):
            yieldRatio = strategyYield / (0.001 + processedReadPairs) * 100.0
            if yieldRatio >= minAutoDetectPct:
                selectedStrategies.append(strategy)
        return selectedStrategies