#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import glob
import logging
from colorama import Fore
from colorama import Back
from colorama import Style
import os
import collections
import itertools
import gzip

#logging.basicConfig(level=logging.DEBUG)
# http://codereview.stackexchange.com/questions/88912/create-a-list-of-all-strings-within-hamming-distance-of-a-reference-string-with
def hamming_circle(s, n, alphabet):
    for positions in itertools.combinations(range(len(s)), n):
        for replacements in itertools.product(
                range(len(alphabet) - 1), repeat=n):
            cousin = list(s)
            for p, r in zip(positions, replacements):
                if cousin[p] == alphabet[r]:
                    cousin[p] = alphabet[-1]
                else:
                    cousin[p] = alphabet[r]
            yield ''.join(cousin)


class BarcodeParser():

    def path_to_barcode_alias(self, path):
        barcodeFileAlias = os.path.splitext(
            os.path.basename(path))[0].replace('.gz','').replace('.bc','')
        return barcodeFileAlias

    def parse_pending_barcode_file_of_alias(self, alias):
        if not alias in self.pending_files:
            raise ValueError(f"trying to load {alias}, but barcode file not found")

        logging.info(f"Loading promised barcode file {alias}")
        self.parse_barcode_file( self.pending_files[alias] )
        logging.info(f"Performing hamming extension for {alias}")
        self.expand(self.hammingDistanceExpansion, alias=alias)
        del self.pending_files[alias]

    def parse_barcode_file(self, barcodeFile):

        barcodeFileAlias = self.path_to_barcode_alias(barcodeFile)
        # Decide the file type (index first or name first)
        indexNotFirst = False
        with gzip.open(barcodeFile,'rt') if barcodeFile.endswith('.gz') else open(barcodeFile) as f :
            for i, line in enumerate(f):
                parts = line.strip().split()
                if len(parts) == 1 and ' ' in line:
                    parts = line.strip().split(' ')
                if len(parts) == 1:
                    pass
                elif len(parts) == 2:
                    indexFirst = not all((c in 'ATCGNX') for c in parts[0])
                    if not indexFirst:
                        indexNotFirst = True
                    # print(parts[1],indexFirst)

        nospec=False
        i=0
        with gzip.open(barcodeFile,'rt') if barcodeFile.endswith('.gz') else open(barcodeFile) as f :

            for i, line in enumerate(f):
                parts = line.strip().split()
                if len(parts) == 1 and ' ' in line:
                    parts = line.strip().split(' ')
                if len(parts) == 1:
                    self.addBarcode(
                        barcodeFileAlias, barcode=parts[0], index=i+1)
                    #self.barcodes[barcodeFileAlias][parts[0]] = i
                    if not nospec: # only show this once:
                        logging.info(
                            f"\t{parts[0]}:{i} (No index specified in file)")
                        nospec=True
                elif len(parts) == 2:
                    if indexNotFirst:
                        barcode, index = parts
                    else:
                        index, barcode = parts
                    #self.barcodes[barcodeFileAlias][barcode] = index

                    # When the index is only digits, convert to integer
                    try:
                        if int(index)==int(str(int(index))):
                            index = int(index)
                        else:
                            pass
                    except Exception as e:
                        pass
                    self.addBarcode(
                        barcodeFileAlias, barcode=barcode, index=index)
                    if not nospec: # only show this once:
                        logging.info(
                            f"\t{barcode}:{index} (index was specified in file, {'index' if indexFirst else 'barcode'} on first column)")
                        nospec=True
                else:
                    e = f'The barcode file {barcodeFile} contains more than two columns. Failed to parse!'
                    logging.error(e)
                    raise ValueError(e)
        logging.info(f'done loading {i} barcodes')


    def __init__(
            self,
            barcodeDirectory='barcodes',
            hammingDistanceExpansion=0,
            spaceFill=False,
            lazyLoad=None # these aliases wiill not be loaded until requested, '*' matches all files
            ):

        self.hammingDistanceExpansion = hammingDistanceExpansion

        barcodeDirectory = os.path.join(
            os.path.dirname(
                os.path.realpath(__file__)),
            barcodeDirectory)
        barcode_files = list(glob.glob(f'{barcodeDirectory}/*'))

        self.spaceFill = spaceFill
        self.hammingDistanceExpansion = hammingDistanceExpansion
        self.barcodes = collections.defaultdict(
            dict)  # alias -> barcode -> index
        # alias -> barcode -> (index, hammingDistance)
        self.extendedBarcodes = collections.defaultdict(dict)
        self.pending_files = dict()
        for barcodeFile in barcode_files:
            barcodeFileAlias = self.path_to_barcode_alias(barcodeFile)
            if lazyLoad is not None and (barcodeFileAlias in lazyLoad or lazyLoad =='*'):
                logging.info(f"Lazy loading {barcodeFile}, alias {barcodeFileAlias}")
                self.pending_files[barcodeFileAlias] = barcodeFile
                continue
            logging.info(f"Parsing {barcodeFile}, alias {barcodeFileAlias}")
            self.parse_barcode_file(barcodeFile)

            if hammingDistanceExpansion > 0:
                self.expand(hammingDistanceExpansion, alias=barcodeFileAlias)

    def getTargetCount(self, barcodeFileAlias):
        return(len(self.barcodes[barcodeFileAlias]), len(self.extendedBarcodes[barcodeFileAlias]))

    def expand(
            self,
            hammingDistanceExpansion,
            alias,
            reportCollisions=True,
            spaceFill=None):

        barcodes = self.barcodes[alias]
        # hammingBarcode -> ( ( distance,origin) )
        hammingSpace = collections.defaultdict(list)
        for barcode in barcodes:

            for hammingDistance in range(0, hammingDistanceExpansion + 1):
                for hammingInstance in hamming_circle(
                        barcode, hammingDistance, 'ACTGN'):
                    hammingSpace[hammingInstance].append(
                        (hammingDistance, barcode))
        # Resolve all
        for hammingBarcode in hammingSpace:

            # Check if there is a closest origin:
            sortedDistances = sorted(hammingSpace[hammingBarcode])
            if len(sortedDistances) > 1 and (
                    sortedDistances[0][0] == sortedDistances[1][0]):
                # We cannot resolve this, two or more origins are at the same distance:
                #print('Cannot resolve %s' % hammingBarcode)
                continue

            hammingDistance, origin = sortedDistances[0]

            self.addBarcode(
                alias,
                barcode=hammingBarcode,
                index=self.barcodes[alias][origin],
                hammingDistance=hammingDistance,
                originBarcode=origin)

    def addBarcode(
            self,
            barcodeFileAlias,
            barcode,
            index,
            hammingDistance=0,
            originBarcode=None):
        if hammingDistance == 0:
            self.barcodes[barcodeFileAlias][barcode] = index
        else:
            if originBarcode is None:
                raise ValueError()
            self.extendedBarcodes[barcodeFileAlias][barcode] = (
                index, originBarcode, hammingDistance)

    # get index and hamming distance to barcode  returns none if not Available
    def getIndexCorrectedBarcodeAndHammingDistance(self, barcode, alias, try_lazy_load_pending=True):

        # Check if the alias still needs to be loaded:


        if barcode in self.barcodes[alias]:
            return (self.barcodes[alias][barcode], barcode, 0)
        if barcode in self.extendedBarcodes[alias]:
            return self.extendedBarcodes[alias][barcode]

        if alias in self.pending_files:
            if not try_lazy_load_pending:
                raise RecursionError()
            self.parse_pending_barcode_file_of_alias(alias)
            return self.getIndexCorrectedBarcodeAndHammingDistance(barcode, alias, try_lazy_load_pending=False)

        return (None, None, None)

    def list(self, showBarcodes=5):  # showBarcodes=None show all
        for barcodeAlias, mapping in self.barcodes.items():
            print(
                f'{len(mapping)} barcodes{Style.DIM} obtained from {Style.RESET_ALL}{barcodeAlias}')
            if len(mapping):
                for bcId in list(mapping.keys())[:showBarcodes]:
                    try:
                        print(('%s%s%s%s%s%s' % (Fore.GREEN, bcId,
                                                 Fore.WHITE, '→', mapping[bcId], Style.RESET_ALL)))
                    except Exception as e:
                        print(('%s%s%s%s%s%s' % (Fore.GREEN, bcId,
                                                 Fore.WHITE, '->', mapping[bcId], Style.RESET_ALL)))
                if showBarcodes is not None and len(mapping) > showBarcodes:
                    print(
                        f'{Style.DIM} %s more ...\n{Style.RESET_ALL}' %
                        (len(mapping) - showBarcodes))

    def getBarcodeMapping(self):
        return self.barcodes

    def __getitem__(self, alias):
        if alias in self.pending_files:
            self.parse_pending_barcode_file_of_alias(alias)
        return self.barcodes.get(alias)