from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum, auto
from functools import cache, cached_property, total_ordering
from itertools import chain, combinations
from re import search
from typing import TYPE_CHECKING, Annotated, Any, Literal

from pydantic import PlainValidator

if TYPE_CHECKING:
    from collections.abc import Hashable, Iterable, Sequence

    from chemex.nmr.basis import Basis


SPIN_ALIASES = "isx"


class Nucleus(Enum):
    """Define all the different types of atoms."""

    H1 = auto()
    N15 = auto()
    C13 = auto()
    X = auto()


# Conversion dictionary from atom letter to corresponding Nucleus
STR_TO_NUCLEUS: dict[str, Nucleus] = {
    "H": Nucleus.H1,
    "Q": Nucleus.H1,
    "M": Nucleus.H1,
    "N": Nucleus.N15,
    "C": Nucleus.C13,
    "X": Nucleus.X,
}

# Conversion dictionary from 3-letter to 1-letter amino-acid convention
AAA_TO_A = {
    "ALA": "A",
    "ARG": "R",
    "ASN": "N",
    "ASP": "D",
    "CYS": "C",
    "GLN": "Q",
    "GLU": "E",
    "GLY": "G",
    "HIS": "H",
    "ILE": "I",
    "LEU": "L",
    "LYS": "K",
    "MET": "M",
    "PHE": "F",
    "PRO": "P",
    "SER": "S",
    "THR": "T",
    "TRP": "W",
    "TYR": "Y",
    "VAL": "V",
}


# Conversion dictionary to correct different ways to spell nuclei
CORRECT_ATOM_NAME = {"HN": "H", "C'": "C", "CO": "C"}

# fmt: off
STANDARD_ATOM_NAMES_PROTEIN = {
    "C", "CA", "CB", "CD", "CD1", "CD2", "CE", "CE1", "CE2", "CE3", "CG", "CG1", "CG2",
    "CH2", "CQD", "CQE", "CQG", "CZ", "CZ2", "CZ3", "H", "H2", "H3", "HA", "HA2",
    "HA3", "HB", "HB1", "HB2", "HB3", "HD", "HD1", "HD11", "HD12", "HD13", "HD2",
    "HD21", "HD22", "HD23", "HD3", "HE", "HE1", "HE2", "HE21", "HE22", "HE3", "HG",
    "HG1", "HG11", "HG12", "HG13", "HG2", "HG21", "HG22", "HG23", "HG3", "HH", "HH1",
    "HH11", "HH12", "HH2", "HH21", "HH22", "HZ", "HZ1", "HZ2", "HZ3", "MB", "MD", "MD1",
    "MD2", "ME", "MG", "MG1", "MG2", "MZ", "N", "ND1", "ND2", "NE", "NE1", "NE2", "NH",
    "NH1", "NH2", "NQH", "NZ", "QA", "QB", "QD", "QD2", "QE", "QE2", "QG", "QG1", "QH1",
    "QH2", "QMD", "QMG", "QQH", "QR", "QZ",
}

STANDARD_ATOM_NAMES_DNA_RNA = {
    "C1'", "C2", "C2'", "C3'", "C4", "C4'", "C5", "C5'", "C6", "C7", "C8",
    "H1", "H1'", "H2", "H2'", "H2'1", "H2'2", "H21", "H22", "H3", "H3'", "H4'", "H41",
    "H42", "H5", "H5'1", "H5'2", "H6", "H61", "H62", "H71", "H72", "H73", "H8",
    "M7", "N1", "N2", "N3", "N4", "N6", "N7", "N9"
}
# fmt: on
STANDARD_ATOM_NAMES = STANDARD_ATOM_NAMES_PROTEIN | STANDARD_ATOM_NAMES_DNA_RNA


@dataclass(order=True)
class Atom:
    name: str
    nucleus: Nucleus = field(init=False, compare=False)
    search_keys: set[Hashable] = field(init=False, default_factory=set, compare=False)

    def __post_init__(self):
        name = self.name.strip().upper()
        self.name = CORRECT_ATOM_NAME.get(name, name)
        self.nucleus = STR_TO_NUCLEUS.get(self.name[:1], Nucleus.X)
        self.search_keys.update({self, self.nucleus})

    def match(self, other: Atom) -> bool:
        return other.name.startswith(self.name)

    def __hash__(self) -> int:
        return hash(self.name)

    def __bool__(self) -> bool:
        return bool(self.name)

    def __str__(self):
        return self.name


@total_ordering
class Group:
    NO_NUMBER: int = -100000000

    def __init__(self, name: str) -> None:
        self.symbol, self.number, self.suffix = self.parse_group(name.strip().upper())
        self.symbol = AAA_TO_A.get(self.symbol, self.symbol)
        self.search_keys: set[Hashable] = {self} if self else set()

    def parse_group(self, name: str) -> tuple[str, int, str]:
        if found := search("[0-9]+", name.strip().upper()):
            return name[: found.start()], int(found.group()), name[found.end() :]
        return name, self.NO_NUMBER, ""

    @property
    def name(self) -> str:
        number = "" if self.number == self.NO_NUMBER else self.number
        return f"{self.symbol}{number}{self.suffix}"

    def match(self, other: Group) -> bool:
        symbol = other.symbol == self.symbol or not self.symbol
        number = self.number in (other.number, self.NO_NUMBER)
        suffix = other.suffix == self.suffix or not self.suffix
        return number and symbol and suffix

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, Group):
            return NotImplemented
        return self.number < other.number

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Group):
            return NotImplemented
        return self.name == other.name

    def __hash__(self) -> int:
        return hash(self.name)

    def __bool__(self) -> bool:
        return bool(self.name)

    def __str__(self) -> str:
        return self.name


@total_ordering
class Spin:
    def __init__(self, name: str, group_for_completion: Group | None = None) -> None:
        self.group, self.atom = self.split_group_atom(name.strip().upper())
        if not self.group and group_for_completion:
            self.group = group_for_completion
        self.search_keys = self.group.search_keys | self.atom.search_keys

    @staticmethod
    def split_group_atom(name: str) -> tuple[Group, Atom]:
        if name == "?":
            return Group(""), Atom("")
        found_digit = search("[0-9]", name)
        first_digit = found_digit.start() if found_digit else 0
        found_atom = search("[HCNQM]", name[first_digit:])
        if not found_atom:
            if name in STANDARD_ATOM_NAMES:
                return Group(""), Atom(name)
            return Group(name), Atom("")
        atom_index = first_digit + found_atom.start()
        return Group(name[:atom_index]), Atom(name[atom_index:])

    @property
    def name(self) -> str:
        return f"{self.group}{self.atom}"

    def match(self, other: Spin):
        return self.group.match(other.group) and self.atom.match(other.atom)

    def __lt__(self, other: Spin) -> bool:
        self_tuple = (self.atom.nucleus.name, self.group, self.atom)
        other_tuple = (other.atom.nucleus.name, other.group, other.atom)
        return self_tuple < other_tuple

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Spin):
            return NotImplemented
        return self.name == other.name

    def __hash__(self) -> int:
        return hash(self.name)

    def __bool__(self) -> bool:
        return bool(self.name)

    def __repr__(self) -> str:
        return self.name

    def __str__(self) -> str:
        return self.name


def powerset(iterable: Iterable[Any]) -> Iterable[Any]:
    "powerset([1,2,3]) --> (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)."
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(1, len(s) + 1))


@total_ordering
class SpinSystem:
    @staticmethod
    @cache
    def _parse_spin_system(name: str) -> dict[str, Spin]:
        if not name:
            return {}
        split = name.split("-")
        spins: dict[str, Spin] = {}
        last_group = None
        for short_name, name_spin in zip(SPIN_ALIASES, split, strict=False):
            spin = Spin(name_spin, last_group)
            spins[short_name] = spin
            last_group = spin.group
        return spins

    @staticmethod
    def _spins2name(spins: Iterable[Spin]) -> str:
        spin_names: list[str] = []
        last_group: Group = Group("")
        for spin in spins:
            spin_name = str(spin.atom) if spin.group == last_group else str(spin)
            spin_names.append(spin_name)
            last_group = spin.group
        return "-".join(spin_names)

    def __init__(self, name: str | int | None = None):
        if name is None:
            name = ""
        if isinstance(name, int):
            name = str(name)
        self.spins = self._parse_spin_system(name.strip().upper())
        self.name = self._spins2name(self.spins.values())
        self.search_keys: set[Hashable] = set()
        self.search_keys = self.search_keys.union(
            *(spin.search_keys for spin in self.spins.values())
        )

    def __deepcopy__(self, memo: dict[Any, Any]) -> SpinSystem:
        return SpinSystem(self.name)

    @cached_property
    def names(self) -> dict[str, str]:
        result: dict[str, str] = {}
        for alias_set in powerset(SPIN_ALIASES):
            if set(alias_set).issubset(self.spins):
                key = "".join(alias_set)
                name = self._spins2name(self.spins[alias] for alias in alias_set)
                result[key] = name
        return result

    @cached_property
    def groups(self) -> dict[str, Group]:
        return {alias: spin.group for alias, spin in self.spins.items()}

    @cached_property
    def symbols(self) -> dict[str, str]:
        return {alias: group.symbol for alias, group in self.groups.items()}

    @cached_property
    def numbers(self) -> dict[str, int]:
        return {alias: group.number for alias, group in self.groups.items()}

    @cached_property
    def atoms(self) -> dict[str, Atom]:
        return {alias: spin.atom for alias, spin in self.spins.items()}

    @cached_property
    def nuclei(self) -> dict[str, Nucleus]:
        return {alias: atom.nucleus for alias, atom in self.atoms.items()}

    def match(self, other: SpinSystem) -> bool:
        spin_pairs = zip(self.spins.values(), other.spins.values(), strict=False)
        return all(spin.match(other_spin) for spin, other_spin in spin_pairs)

    def part_of(self, selection: Sequence[SpinSystem] | str) -> bool:
        if isinstance(selection, str):
            return selection.lower() in ("all", "*")
        return any(item.match(self) for item in selection)

    def correct(self, basis: Basis) -> SpinSystem:
        spins: list[Spin] = []
        last_spin = Spin("")
        for letter, atom in basis.atoms.items():
            spin = Spin(self.spins.get(letter, last_spin).name)
            if not spin.atom.name.startswith(atom.upper()):
                spin.atom = Atom(f"{atom}{spin.atom.name[1:]}")
            last_spin = spin
            spins.append(spin)
        return SpinSystem(self._spins2name(spins))

    def build_sub_spin_system(
        self, spin_system_part: Literal["i", "s", "is", "g", ""]
    ) -> SpinSystem:
        if spin_system_part == "g":
            return SpinSystem(str(self.groups.get("i", "")))
        if spin_system_part == "i":
            return SpinSystem(str(self.spins.get("i", "")))
        if spin_system_part == "s":
            return SpinSystem(str(self.spins.get("s", "")))
        if spin_system_part == "is":
            return self
        return SpinSystem()

    def __and__(self, other: Any) -> SpinSystem:
        if not isinstance(other, SpinSystem):
            return NotImplemented
        if self == other:
            return SpinSystem(self.name)
        if spins := set(self.spins.values()) & set(other.spins.values()):
            return SpinSystem("-".join(spin.name for spin in spins))
        groups = set(self.groups.values()) & set(other.groups.values())
        if len(groups) == 1:
            return SpinSystem("-".join(group.name for group in groups))
        return SpinSystem("")

    def __lt__(self, other: object) -> bool:
        if not isinstance(other, SpinSystem):
            return NotImplemented
        return tuple(self.spins.values()) < tuple(other.spins.values())

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, SpinSystem):
            return NotImplemented
        return self.name == other.name

    def __hash__(self) -> int:
        return hash(self.name)

    def __bool__(self) -> bool:
        return bool(self.name)

    def __str__(self) -> str:
        return self.name


# We now create an `Annotated` wrapper that we'll use as the annotation for
# fields on `BaseModel`
PydanticSpinSystem = Annotated[SpinSystem, PlainValidator(lambda x: SpinSystem(x))]


# if __name__ == "__main__":
#     print(SpinSystem("G23N-G23HN"), SpinSystem("GLY023N-HN"), SpinSystem("g23N-hN"))
#     print(SpinSystem("G23N-G23HN") == SpinSystem("GLY023N-HN"))
#     print(SpinSystem("G23N-G23HN").match(SpinSystem("GLY023N-HN")))
#     print(SpinSystem("G23N-G23HN") & SpinSystem("G23C"))
#     print(SpinSystem("") == SpinSystem())
#     group = Group("L99")
#     spin = Spin("HD1", group)
#     print(f"spin = {spin}, spin.group = {spin.group}, spin.atom = {spin.atom}")
#     print(SpinSystem("GLY023N-HN").search_keys)
#     print(SpinSystem("GLY023N-HN").names)
#     print(SpinSystem("G23C1'").names)
