"""
Algorithmic steps, functions take tables, scalars, etc. as input and also output tables, scalars etc.
Does not deal with IO. Functions need to be stitched together to do something useful.
"""
import time
import logging
import pysam
from collections import Counter, defaultdict, deque
import functools
import regex
from typing import Optional, Dict
from phantombuster.error_corrector import ErrorCorrector, Umi, UmiCount
from phantombuster.io_ import open_sequencing_file
import numpy as np
from itertools import tee
import scipy.stats
from dataclasses import dataclass
from scipy.stats._discrete_distns import binom
import polars as pl

from phantombuster import bamindexer

def collect(df):
    if hasattr(df, "collect"):
        return df.collect()
    return df

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

_marker = object()

try:
    StringCounter = Counter[str] # type: ignore
except:
    StringCounter = Counter # type: ignore

featuretype = tuple

# copied from more-itertools https://github.com/more-itertools/more-itertools 
class peekable:
    """Wrap an iterator to allow lookahead and prepending elements.

    Call :meth:`peek` on the result to get the value that will be returned
    by :func:`next`. This won't advance the iterator:

        >>> p = peekable(['a', 'b'])
        >>> p.peek()
        'a'
        >>> next(p)
        'a'

    Pass :meth:`peek` a default value to return that instead of raising
    ``StopIteration`` when the iterator is exhausted.

        >>> p = peekable([])
        >>> p.peek('hi')
        'hi'

    peekables also offer a :meth:`prepend` method, which "inserts" items
    at the head of the iterable:

        >>> p = peekable([1, 2, 3])
        >>> p.prepend(10, 11, 12)
        >>> next(p)
        10
        >>> p.peek()
        11
        >>> list(p)
        [11, 12, 1, 2, 3]

    peekables can be indexed. Index 0 is the item that will be returned by
    :func:`next`, index 1 is the item after that, and so on:
    The values up to the given index will be cached.

        >>> p = peekable(['a', 'b', 'c', 'd'])
        >>> p[0]
        'a'
        >>> p[1]
        'b'
        >>> next(p)
        'a'

    Negative indexes are supported, but be aware that they will cache the
    remaining items in the source iterator, which may require significant
    storage.

    To check whether a peekable is exhausted, check its truth value:

        >>> p = peekable(['a', 'b'])
        >>> if p:  # peekable has items
        ...     list(p)
        ['a', 'b']
        >>> if not p:  # peekable is exhausted
        ...     list(p)
        []

    """

    def __init__(self, iterable):
        self._it = iter(iterable)
        self._cache = deque()

    def __iter__(self):
        return self

    def __bool__(self):
        try:
            self.peek()
        except StopIteration:
            return False
        return True

    def peek(self, default=_marker):
        """Return the item that will be next returned from ``next()``.

        Return ``default`` if there are no items left. If ``default`` is not
        provided, raise ``StopIteration``.

        """
        if not self._cache:
            try:
                self._cache.append(next(self._it))
            except StopIteration:
                if default is _marker:
                    raise
                return default
        return self._cache[0]

    def prepend(self, *items):
        """Stack up items to be the next ones returned from ``next()`` or
        ``self.peek()``. The items will be returned in
        first in, first out order::

            >>> p = peekable([1, 2, 3])
            >>> p.prepend(10, 11, 12)
            >>> next(p)
            10
            >>> list(p)
            [11, 12, 1, 2, 3]

        It is possible, by prepending items, to "resurrect" a peekable that
        previously raised ``StopIteration``.

            >>> p = peekable([])
            >>> next(p)
            Traceback (most recent call last):
              ...
            StopIteration
            >>> p.prepend(1)
            >>> next(p)
            1
            >>> next(p)
            Traceback (most recent call last):
              ...
            StopIteration

        """
        self._cache.extendleft(reversed(items))

    def __next__(self):
        if self._cache:
            return self._cache.popleft()

        return next(self._it)

    def _get_slice(self, index):
        # Normalize the slice's arguments
        step = 1 if (index.step is None) else index.step
        if step > 0:
            start = 0 if (index.start is None) else index.start
            stop = maxsize if (index.stop is None) else index.stop
        elif step < 0:
            start = -1 if (index.start is None) else index.start
            stop = (-maxsize - 1) if (index.stop is None) else index.stop
        else:
            raise ValueError('slice step cannot be zero')

        # If either the start or stop index is negative, we'll need to cache
        # the rest of the iterable in order to slice from the right side.
        if (start < 0) or (stop < 0):
            self._cache.extend(self._it)
        # Otherwise we'll need to find the rightmost index and cache to that
        # point.
        else:
            n = min(max(start, stop) + 1, maxsize)
            cache_len = len(self._cache)
            if n >= cache_len:
                self._cache.extend(islice(self._it, n - cache_len))

        return list(self._cache)[index]

    def __getitem__(self, index):
        if isinstance(index, slice):
            return self._get_slice(index)

        cache_len = len(self._cache)
        if index < 0:
            self._cache.extend(self._it)
        elif index >= cache_len:
            self._cache.extend(islice(self._it, index + 1 - cache_len))

        return self._cache[index]


def _get_type(c, item):
    """
    Get the type of a column depending on its name and one example item.
    Returns None for undetermined type.
    """
    name, type = c['name'], c['type']
    if type == 'reference':
        return pl.String
    if name == 'reads':
        return pl.UInt64
    return pl.String


def deduplicator_to_table(deduplicator, barcode_hierarchy):
    """
    Transform a collections.Counter to a table.
    The counter has a tuple as keys and the readcount as values,
    the order of the columns must be the same order as the tuple-key + "reads"
    """
    barcode_hierarchy = barcode_hierarchy+ [{'name': "reads", 'type': 'count'}]
    column_names = [c['name'] for c in barcode_hierarchy]
    rows = ((*items, readcount) for items, readcount in deduplicator.items())
    rows = peekable(rows)
    try:
        first_row = rows.peek()
    except StopIteration:
        return None

    arrays = zip(*rows)
    types = [_get_type(c, value) for c, value in zip(barcode_hierarchy, first_row)]

    assert len(first_row) == len(barcode_hierarchy)

    try:
        arrays = [pl.Series(bc["name"], x, dtype=type) for x, type, bc in zip(arrays, types, barcode_hierarchy)]
    except Exception as e:
        logging.error('First row %s', first_row)
        logging.error('Column names %s', column_names)
        logging.error(f'First entry of each array: {[col[0] for col in arrays]}, types: {types}')
        raise e
    table = pl.DataFrame(arrays)
    return table

def table_to_counter(table, columns):
    """
    Transform a table to a collections.counter.
    `columns` determines the order of presence of columsn in the counter.
    """
    d = table.to_dict()
    counter = Counter(dict(zip(zip(*(d[col] for col in columns)), d["reads"])))
    return counter


def get_sectioned_iterator(inputgroup, start, end):
    if len(inputgroup.files) == 1:
        inputfile = inputgroup.files[0]
        with open_sequencing_file(inputfile.path) as f:
            if start != 0:
                f.seek(start)
            for read in f:
                yield read
                if f.tell() >= end:
                    break
    else:
        input_files = [(iter(open_sequencing_file(inputfile.path)), inputfile.prefix) for inputfile in inputgroup.files]

        running = True

        while running:
            read = {}
            for f, prefix in input_files:
                try:
                    partial_read = next(f)
                except StopIteration:
                    running = False
                    break
                for tag in partial_read:
                    read[prefix+tag] = partial_read[tag]
            if running:
                yield read


def deduplicate_section(inputgroup, section,
        regex_dictionary,
        barcode_hierarchy,
        count_qc: bool):
    logger = logging.getLogger()

    if section is None:
        start, end = 0, np.inf
    else:
        start, end = section

    regexes = regex_dictionary.get_regexes_for_group(inputgroup.group)

    feature_deduplicator : Counter = Counter()

    # aggregate some information while iterating through all reads
    processed_reads = 0                   # count processed reads
    remaining_reads = 0                   # count surviving reads
    counters_success = {bc["name"]: StringCounter() for bc in barcode_hierarchy if bc["type"] == "reference" }
    counters_fail = StringCounter()

    logging_buffer = 5000000

    logger.debug(f"Processing group {inputgroup.group}")

    opened_input_file = get_sectioned_iterator(inputgroup, start, end)

    for read in opened_input_file:
        insufficient = False
        processed_reads += 1
        if processed_reads % logging_buffer == 0:
            logger.info(f"Processed {processed_reads} reads from group {inputgroup.group}")

        features = extract_read(regexes, read)

        for bc in barcode_hierarchy:
            tag = bc["name"]
            if tag not in features or len(features[tag]) == 0:
                counters_fail[tag] += 1
                insufficient = True

        if not insufficient:
            for bc in barcode_hierarchy:
                col = bc["name"]
                if bc["min_length"] is not None:
                    if len(features[col]) < bc["min_length"]:
                        insufficient = True
                        counters_fail[col+"_length"] += 1
                if bc ["max_length"] is not None:
                    features[col] = features[col][:bc["max_length"]]

        if not insufficient:
            for bc in barcode_hierarchy:
                if bc["type"] == "reference":
                    col = bc["name"]
                    correct = deterministic_correct(features[col], bc["reference"], None, bc["threshold"])
                    features[col] = correct

                    if correct is None:
                        counters_fail.update([col+"_correction_failed"])
                        insufficient = True

        if not insufficient:
            for bc in barcode_hierarchy:
                if bc["type"] == "reference":
                    col = bc["name"]
                    counters_success[col].update([features[col]])

            data = tuple(features.get(bc["name"], "") for bc in barcode_hierarchy)
            feature_deduplicator.update([data])
            remaining_reads += 1

    logger.info(f"Finished parsing {inputgroup.group}.")

    logger.info("Sorting each sample")

    table = deduplicator_to_table(feature_deduplicator, barcode_hierarchy)
    if table is not None:
        table = table.sort([bc["name"] for bc in barcode_hierarchy])

    logger.info("demultiplexing thread finished")

    return table, (processed_reads, remaining_reads, counters_success, counters_fail) 


def extract_read(regex_dictionary, read):
    """Extract features from read

    Reads are made out of three regions of interest: The query, the b2tag and the bctag.
    Features can be scattered across these regions, e.g. part1 barcode is in the query region
    and part2 is in the b2tag.

    The function applies the user defined regex for each region and then aggregates
    the features. 
    """
    groups = {}

    for tag in read:
        re = regex_dictionary.get(tag, None)
        if re:
            match_ = re.search(read[tag])
        else:
            continue

        if not match_:
            continue

        groups.update(match_.groupdict())

    tag_stems = [tag.rstrip('0123456789') for tag in groups.keys()]

    # normalize ["sample", "sample1"] to ["sample0", "sample1"] so that sort handles this correctly
    for tag0 in tag_stems:
        if tag0 in groups:
            groups[tag0 + "0"] = groups[tag0]
            del groups[tag0]

    features = {}
    for tag in tag_stems:
        combined_tag = functools.reduce(
            lambda x, y: x + groups[y],
            sorted(name for name in groups.keys() if name.startswith(tag)),
            "",
        )
        features[tag] = combined_tag

    return features


def _combine_duplicators_polars2(main_table, partial_table, barcode_hierarchy):
    if main_table is None:
        return partial_table
    if partial_table is None:
        return main_table

    table = pl.concat([main_table, partial_table])

    table = table.group_by([bc["name"] for bc in barcode_hierarchy]).agg(pl.col("reads").sum())
    table = table.sort([bc["name"] for bc in barcode_hierarchy])
    return table


def _combine_duplicators_polars(main_table, partial_table, barcode_hierarchy):
    if main_table is None:
        return partial_table
    if partial_table is None:
        return main_table

    bchi = 0
    while bchi < len(barcode_hierarchy) and barcode_hierarchy[bchi]["type"] == "random":
        bchi += 1

    if bchi < len(barcode_hierarchy) and barcode_hierarchy[bchi]["type"] == "reference":
        anchor = pl.col(barcode_hierarchy[bchi]["name"])
        df1_anchor = collect(main_table.select(anchor.unique(maintain_order=True).alias("anchor")))
        df2_anchor = collect(partial_table.select(anchor.unique(maintain_order=True).alias("anchor")))

        anchor_values = pl.concat([df1_anchor, df2_anchor]).select(pl.col("anchor").unique(maintain_order=True))["anchor"]

        tables = []
        for a in anchor_values:
            df = (main_table.filter(anchor == a).join(partial_table.filter(anchor == a), on=[bc["name"] for bc in barcode_hierarchy if bc["name"] != "reads"], how="full", coalesce=True))
            df = df.select(*[bc["name"] for bc in barcode_hierarchy if bc["name"] != "reads"],
                                       (pl.col("reads").fill_null(0) + pl.col("reads_right").fill_null(0)).alias("reads"))
            df = df.sort([bc["name"] for bc in barcode_hierarchy])
            tables.append(df)
        return pl.concat(tables)
    else:
        table = (main_table.join(partial_table, on=[bc["name"] for bc in barcode_hierarchy if bc["name"] != "reads"], how="full", coalesce=True))
        table = table.select(*[bc["name"] for bc in barcode_hierarchy if bc["name"] != "reads"],
                                   (pl.col("reads").fill_null(0) + pl.col("reads_right").fill_null(0)).alias("reads"))

        table = table.sort([bc["name"] for bc in barcode_hierarchy])
        return table


def combine(deduplications, barcode_hierarchy):
    main_table, processed_reads, remaining_reads, counters_success, counters_fail = None, 0, 0, defaultdict(Counter), Counter()

    for i, deduplication in enumerate(deduplications):
        partial_table, (_processed_reads, _remaining_reads,  _counters_success, _counters_fail) = deduplication

        main_table = _combine_duplicators_polars(main_table, partial_table, barcode_hierarchy)

        processed_reads += _processed_reads
        remaining_reads += _remaining_reads
        counters_success = {name: counters_success[name] + Counter(_counters_success[name]) for name in set(counters_success.keys()) | set(_counters_success.keys())}
        counters_fail += _counters_fail

        del deduplication
        del partial_table

    return main_table, (processed_reads, remaining_reads, counters_success, counters_fail)


def error_correct(table, barcode_hierarchy, threshold):
    partitions = [table]
    tag_names = [tag["name"] for tag in barcode_hierarchy[1:]]
    for tag_idx, tag in enumerate(barcode_hierarchy):
        logging.debug("Error correcting barcode %s of type %s", type["name"], type["type"])
        new_partitions = []
        for partition in partitions:
            if tag["type"] == "random":
                partition = error_correct_column(partition, tag['name'], threshold)

            sort_necessary = tag['type'] == 'random'
            ps = sort_and_partition(partition, sort_columns=[tag['name']], partition_columns=[tag['name']], unwrap_if_single=False, sort=sort_necessary)
            for partition in ps.values():
                new_partitions.append(partition)
        partitions = new_partitions
    if len(partitions) > 0:
        table = pl.concat(partitions)
    return table

def error_correct_column(partition, tag, threshold):
    logging.debug("START Error correcting column %s", tag)
    tag_length = len(partition[tag][0])

    wrong_lengths = partition.filter(pl.col(tag).str.len_chars() != tag_length)

    if len(wrong_lengths) > 0:
        logging.error(f"There are {len(wrong_lengths)} many barcodes that do not have length {tag_length}, aborting")
        logging.error(f"One example: {wrong_lengths.head(1)}")
        raise ValueError(f"There are {len(wrong_lengths)} many barcodes that do not have length {tag_length}, aborting")

    tag_and_reads = partition.drop([column for column in partition.columns if not column in [tag, 'reads']])
    tag_and_reads = deduplicate_single(tag_and_reads)

    correction = calculate_corrections(tag_and_reads[tag], tag_and_reads['reads'], tag_length,  threshold)
    partition = apply_correction(partition, tag, correction)
    logging.debug("DONE Error correcting column %s", tag)
    return partition


def remove_ambigious(df, barcode_hierarchy):
    for tag in barcode_hierarchy:
        if tag['type'] == 'random':
            df = df.filter(~pl.col(tag['name']).str.contains_any(["R", "Y", "S", "W", "K", "M", "B", "D", "H", "V", "N"]))
    return df


def deduplicate_single(table):
    return table.group_by([bc for bc in table.columns if bc != "reads"]).agg(pl.col("reads").sum())


def calculate_corrections(tags, reads, tag_length, threshold):
    ec = ErrorCorrector(tag_length, 0, 0, threshold)
    tags = [(Umi(0, 0, tag), UmiCount(r, 1)) for tag, r in zip(tags, reads)]
    tags_dict = dict(tags)
    result = ec.process(tags_dict)

    corrections = result[1]
    corrected   = result[0]

    correction = {a.tag: b.tag for a, b in corrections.items()}

    return correction


def apply_correction(table, correction_column, correction):
    table = table.with_columns(pl.col(correction_column).replace_strict(correction, default=None))
    table = table.filter(~pl.col(correction_column).is_null())
    table = deduplicate_single(table)
    return table


def calculate_hopping_threshold(table, hopping_barcodes, alpha_threshold=0.05):

    columns = table.columns
    anchor_barcodes = [col for col in table.columns if col not in hopping_barcodes]
    if 'reads' in anchor_barcodes:
        anchor_barcodes.remove('reads')

    table = table.select(pl.col("reads").sum().over(anchor_barcodes).alias("n"), (pl.col("reads") >= 2).sum().over(anchor_barcodes).alias("sig"))
    table = table.with_columns(pl.when(pl.col("sig") > 0).then(1/pl.col("sig")).otherwise(pl.lit(1.0)).alias("p"))
    threshold = table.select(pl.struct(["n", "p"]).map_batches(lambda x: binom.isf(alpha_threshold, x.struct.field("n"), x.struct.field("p"))).alias("threshold"))
    return threshold


def call_hopping(table, threshold_list):

    thresholds = pl.concat([df.select(pl.col("threshold").alias(f"treshold{i}")) for i, df in enumerate(threshold_list)], how="horizontal")
    thresholds = thresholds.select(pl.max_horizontal(*[c for c in thresholds.columns]).alias("threshold"))

    df = pl.concat([table, thresholds], how="horizontal")
    df = df.filter(pl.col("reads") >= pl.col("threshold"))
    df = df.select(table.columns)

    return df


def sort_and_partition(table, sort_columns=("samplename", "reads"), partition_columns=("samplename",), unwrap_if_single=True, return_bounds=False, sort=True):
    partition_keys, partitions = _sort_and_partition(table, sort_columns, partition_columns, sort=sort)

    if not return_bounds:
        partitions = [p for p, start, end in partitions]
    d = dict(zip([tuple(key) for key in partition_keys], partitions))

    # in case of only one partition key, make keys the values instead of tuple with length 1
    if unwrap_if_single:
        if len(list(d.keys())[0]) == 1:
            d = {key[0]: value for key, value in d.items()}

    return d

def _sort_and_partition(table, sort_columns=("samplename", "reads"), partition_columns=("samplename",), sort=True):
    if isinstance(partition_columns, str):
        partition_columns = [partition_columns]

    p_col_idx = [sort_columns.index(col) for col in partition_columns]
    table_columns = table.columns

    # all partition columns must appear as sorting columns. More sorting columns are allowed.
    assert all(i != -1 for i in p_col_idx)
    # partition columns must be in same order as sort columns
    assert sorted(p_col_idx) == p_col_idx
    assert all(col in table_columns for col in sort_columns)

    table_length = len(table)

    if sort:
        table = table.sort(*sort_columns)

    partitions     = [(table, 0, len(table))]
    partition_keys = [[]]

    for col in partition_columns:
        new_partitions = []
        new_partition_keys = []

        for part, key in zip(partitions, partition_keys):
            p, orig_start, orig_length = part

            new_keys  = p[col].unique(maintain_order=True)

            start = 0
            indexes = []
            for new_key in new_keys:
                idx = p[col].slice(start).index_of(new_key)
                idx = idx + start
                indexes.append(idx)
                start = idx
            indexes.append(len(p))

            d = {}
            for new_key, (start, end) in zip(new_keys, pairwise(indexes)):
                new_partition_key = [*key, new_key]
                assert end-start >= 0
                new_partition = (p.slice(start, end-start), orig_start+start, end-start)

                new_partitions.append(new_partition)
                new_partition_keys.append(new_partition_key)

        partitions = new_partitions
        partition_keys = new_partition_keys
    assert sum(len(p[0]) for p in partitions) == table_length
    return partition_keys, partitions


def threshold_table(table, threshold):
    table = table.join(threshold, "sample")
    table = table.filter(pl.col("reads") >= pl.col("threshold"))
    return table


def _metric(a, b):
    return sum([aa != bb for aa, bb in zip(a,b)])


def deterministic_correct(barcode, references, default_value, threshold):
    # TODO do this better
    r = references.get(barcode, None)
    if r is not None:
        return r
    else:
        for reference_bc, reference_name in references.items():
            if _metric(barcode, reference_bc) <= threshold:
                return reference_name
        return default_value


def calculate_threshold(value_dict):
    th = np.inf
    for v1 in value_dict.keys():
        for v2 in value_dict.keys():
            if v1 == v2:
                continue
            th = min(th, int((_metric(v1, v2)-1)/2))
    return th
