from abc import ABC
from typing import Iterable

from .polymer_specs import PolymerClassifier
import random
import os
from aldepyde.distribution.distribution import Distribution, load
from aldepyde.distribution.distributions.standards import distribution_head
import json
import math

class ResidueGenerator(ABC):
    def __init__(self, classifier:PolymerClassifier):
        self.classifier = classifier

    def _generate(self):
        pass

    # Just pure, random generation. Nothing fancy
    def random(self, length:int, n:int=1) -> list:
        alphabet = list(self.classifier.alphabet)
        sequences = []
        for i in range(n):
            sequences.append("".join(random.choices(alphabet, k=length)))
        return sequences

    def _list_remove(self, l, element) -> list:
        new_list = []
        for val in l:
            if val != element:
                new_list.append(val)
        return new_list

    def possible_mutations(self, sequence:str|Iterable[str], k=1, mutable:Iterable|None=None, alphabet:Iterable|None=None, strict:bool=False):
        r_alphabet = list(self.classifier.alphabet) if alphabet is None else alphabet
        mutable_sequence = sequence if mutable is None else "".join([x for x in sequence if x in mutable])
        # power operation
        product = 1
        for c in mutable_sequence:
            product *= len(r_alphabet) if c not in r_alphabet else len(r_alphabet) - 1
        return math.comb(len(mutable_sequence), k) * product if len(mutable_sequence) > k else product

    def random_mutation(self, sequence:str|Iterable[str], k=1, mutable:Iterable|None=None, alphabet:Iterable|None=None, strict:bool=False):
        r_alphabet = list(self.classifier.alphabet) if alphabet is None else alphabet
        # Assume an iterable
        if isinstance(sequence, str):
            all_values = [sequence]
        else:
            all_values = sequence

        rets = []
        for sequence in all_values:
            allowed_indeces = range(0, len(sequence)) if mutable is None else [x for x, val in enumerate(sequence) if val in mutable]
            s = list(sequence)
            if strict:
                k_n = k
            else:
                k_n = k if k <= len(allowed_indeces) else len(allowed_indeces)
            for index in  random.sample(allowed_indeces, k_n):
                s[index] = random.choice(self._list_remove(r_alphabet, s[index]))
            rets.append("".join(s))
        return rets

    def random_from_distribution(self, distribution: os.PathLike|str|dict[str, float]|Distribution, length:int|tuple[int,int], n:int=1, exclude:tuple=()) -> list:
        if isinstance(distribution, str) and os.path.isfile(distribution): # Given a file
            with open(distribution) as fp:
                data = json.load(fp)
                if distribution_head not in data.keys():
                    raise ValueError(f"Key '{distribution_head}' must be in the top level of your provided file")
            usable_distribution = Distribution(None, data[distribution_head])
        elif isinstance(distribution, str) and not os.path.isfile(distribution): # Given a preload
            usable_distribution = load(distribution)
        elif isinstance(distribution, Distribution): # Given a distribution
            usable_distribution = distribution
        elif isinstance(distribution, dict): # Given a dictionary
            usable_distribution = Distribution(None, distribution)
        else:
            raise ValueError('Invalid distribution')
        for r in exclude:
            usable_distribution.eliminate_entry(r)
        if not isinstance(length, int):
            size = lambda: random.randint(length[0], length[1])
        else:
            size = lambda: length
        usable_distribution.normalize_map()
        sequences = []
        for _ in range(n):
            sequence = random.choices(population=list(usable_distribution.frequency_map.keys()), weights=list(usable_distribution.frequency_map.values()), k=size())
            sequences.append("".join(sequence))
        return sequences

    def random_mutation_from_distribution(self, sequences:str|Iterable[str], distribution: os.PathLike|str|dict[str, float]|Distribution, k:int=1, exclude:tuple=()) -> list:
        if isinstance(sequences, str):
            all_values = [sequences]
        else:
            all_values = sequences

        if isinstance(distribution, str) and os.path.isfile(distribution): # Given a file
            with open(distribution) as fp:
                data = json.load(fp)
                if distribution_head not in data.keys():
                    raise ValueError(f"Key '{distribution_head}' must be in the top level of your provided file")
            usable_distribution = Distribution(None, data[distribution_head])
        elif isinstance(distribution, str) and not os.path.isfile(distribution): # Given a preload
            usable_distribution = load(distribution)
        elif isinstance(distribution, Distribution): # Given a distribution
            usable_distribution = distribution
        elif isinstance(distribution, dict): # Given a dictionary
            usable_distribution = Distribution(None, distribution)
        else:
            raise ValueError('Invalid distribution')
        for r in exclude:
            usable_distribution.eliminate_entry(r)
        usable_distribution.normalize_map()

        rets = []
        for sequence in all_values:
            s = list(sequence)
            for index in  random.sample(range(0, len(sequence)), k):
                while True:
                    selection = random.choices(population=list(usable_distribution.frequency_map.keys()),
                                   weights=list(usable_distribution.frequency_map.values()), k=1)[0]
                    if selection != s[index]:
                        s[index] = selection
                        break
            rets.append("".join(s))
        return rets

        # for _ in range(n):
        #     sequence = random.choices(population=list(usable_distribution.frequency_map.keys()), weights=list(usable_distribution.frequency_map.values()), k=size())
        #     sequences.append("".join(sequence))
        return s

