import math
from abc import ABC, abstractmethod
from collections import Counter
from dataclasses import dataclass


@dataclass
class ConstraintSpecs():
    categories:dict[str,str]
    constraints:dict[str,float]=None

    def __post_init__(self):
        if self.constraints is None:
            self.set_constraints_from_categories()

    def set_constraints_from_categories(self):
        self.constraints = {k:1 for k in self.categories}

    def normalize(self, establish=True):
        non_zero = [x for x,v in self.categories.items() if len(v) > 0]
        total = sum([self.constraints[x] for x in non_zero])
        normalized = {k:(v/total) if len(self.categories.get(k)) > 0 else 0 for k, v in self.constraints.items() }
        if establish:
            self.constraints = normalized
        return normalized

    def clear_all(self):
        self.categories = dict()
        self.constraints = dict()

    def clear_category(self, category:str):
        self.categories[category] = ""

    def clear_constraint(self, category:str):
        self.constraints[category] = 0

    def clear_constraints(self):
        self.constraints = {k:0 for k in self.constraints}

    def clear_alphabet(self):
        self.categories = {k:"" for k in self.categories}

    def wipe(self, establish=False) -> dict:
        c =  {k:1 for k in self.constraints}
        if establish:
            self.constraints = c
        return self.constraints

    def set_category_constraint(self, category:str, constraint:float):
        self.constraints[category] = constraint

    def hard_set_category_constraint(self, category:str, percentage:float):
        if percentage > 1:
            raise ValueError("Percentage must be less than 1.0")
        self.constraints[category] = 0
        self.normalize()
        self.constraints = {k:v*(1-percentage) for k,v in self.constraints.items()}
        self.constraints[category] = percentage


    #TODO Consider what you want to do with wiping/not wiping
    def combine_categories(self, category_1: str, category_2: str, new_category=None, constraint_ratio: float = 1.0):
        if new_category is None:
            new_category = category_1
        s1 = self.categories.pop(category_1)
        s2 = self.categories.pop(category_2)
        self.constraints.pop(category_1)
        self.constraints.pop(category_2)
        self.constraints[new_category] = constraint_ratio
        self.categories[new_category] = frozenset().union(s1, s2)
        if constraint_ratio == -1:
            self.wipe()

    def get_ratio(self, category: str) -> float:
        return self.constraints.get(category)

    def get_constraints(self) -> dict:
        return self.constraints

    def get_category(self, category:str):
        return self.categories.get(category)

    def get_categories(self) -> dict:
        return self.categories

    def get_expected(self, length:int):
        constraint_ratios = self.normalize()
        pairs = sorted([[k, length * v] for (k, v) in constraint_ratios.items()], key=lambda x: x[1] % 1, reverse=True)
        floored = [[x[0], math.floor(x[1])] for x in pairs]
        total_residues = sum([x[1] for x in floored])
        index = 0
        while total_residues < length:
            floored[index][1] += 1
            total_residues += 1
            index += 1
        return {k: math.floor(v) for (k, v) in floored}

    def _is_in(self, value) -> str:
        for key in self.categories:
            if value in self.categories.get(key):
                return key
        return '*'

    def is_expected(self, sequence:str, min_tolerance:int=1) -> bool:
        expected = self.get_expected(length=len(sequence))
        actual_per_category = {k:0 for k in expected}
        actual_count = Counter(sequence)
        for key in actual_count:
            actual_per_category[self._is_in(key)] += actual_count[key]
        normalized = self.normalize(establish=False)
        for key in actual_per_category:
            p_or_m = max(min_tolerance, math.floor(math.sqrt(len(sequence)*(normalized[key]))))
            if not expected[key] - p_or_m <= actual_per_category[key] <= expected[key] + p_or_m:
                return False
        return True



    @property
    def alphabet(self) -> set:
        return set().union(*[val for val in self.categories.values()])


    #TODO Implement the following methods:
    # get_ratio()
    # get_ratios()
    # get_category()
    # get_categories()
    # get_expected()
    # is_expected()


# Define constraints when generating a polymer
class PolymerConstraints(ABC):
    def __init__(self, specs:ConstraintSpecs):
        self.sps = specs

    # @abstractmethod
    # def get_default_classes(self):
    #     pass

    def get_specs(self) -> ConstraintSpecs:
        return self.sps

    @property
    def specs(self) -> ConstraintSpecs:
        return self.sps

    # def new_constraint_ratios(self) -> dict:
    #     # return self.normalize_table({k:1 for k in self.classes})
    #     return self.specs.normalize(self.specs.wipe(establish=False))

    # def set_constraint_classes(self, table:dict):
    #     self.classes = table
    #
    # def normalize_table(self, table:dict) -> dict:
    #     non_zero = [x for x,v in self.classes.items() if len(v) > 0]
    #     total = sum([table[x] for x in non_zero])
    #     return {k:(v/total) if len(self.classes.get(k)) > 0 else 0 for k, v in table.items() }
    #
    # def normalize(self):
    #     self.constraint_ratios = self.normalize_table(self.constraint_ratios)
    #
    # def set_constraint_ratios(self, table:dict):
    #     self.constraint_ratios = self.normalize_table(table)
    #
    # def set_class_ratio(self, class_name:str, value:float, normalize:bool=False):
    #     self.constraint_ratios[class_name] = value
    #     if normalize:
    #         self.constraint_ratios = self.normalize_table(self.constraint_ratios)
    #
    # def hard_set(self, class_name:str, percentage:float):
    #     if percentage > 1:
    #         raise ValueError("Percentage must be less than 1.0")
    #     self.constraint_ratios[class_name] = 0
    #     self.normalize()
    #     self.constraint_ratios = {k:v*(1-percentage) for k,v in self.constraint_ratios.items()}
    #     self.constraint_ratios[class_name] = percentage
    #
    # def combine(self, class_1:str, class_2:str, new_class:str, ratio:float=-1):
    #     s1 = self.classes.pop(class_1)
    #     s2 = self.classes.pop(class_2)
    #     self.constraint_ratios.pop(class_1)
    #     self.constraint_ratios.pop(class_2)
    #     self.constraint_ratios[new_class] = ratio
    #     self.classes[new_class] = frozenset().union(s1, s2)
    #     if ratio == -1:
    #         self.wipe()
    #
    # def alphabet(self) -> set:
    #     return set().union(*[val for val in self.classes.values()])
    #
    # #TODO Get this working
    #
    # # def remove(self, name:str):
    # #     self.classes.pop(name)
    # #     self.constraint_ratios = self.normalize_table(self.constraint_ratios)
    #
    # def negate(self, name:str):
    #     self.constraint_ratios[name] = 0
    #     self.constraint_ratios = self.normalize_table(self.constraint_ratios)
    #
    # def wipe(self):
    #     self.constraint_ratios = {k:1 for k in self.constraint_ratios}
    #
    # def clear_all(self):
    #     self.constraint_ratios = {k:0 for k in self.constraint_ratios}
    #
    # def expected_per_class(self, length):
    #     constraint_ratios = self.normalize_table(self.constraint_ratios)
    #     pairs = sorted([[k,length*v] for (k, v) in constraint_ratios.items()], key=lambda x: x[1]%1, reverse=True)
    #     floored = [[x[0], math.floor(x[1])] for x in pairs]
    #     total_residues = sum([x[1] for x in floored])
    #     index = 0
    #     while total_residues < length:
    #         floored[index][1] += 1
    #         total_residues += 1
    #         index += 1
    #     return {k:math.floor(v) for (k, v) in floored}
