from __future__ import annotations

import attrs
from biotite.sequence.align import SubstitutionMatrix

from apb.sequence.alphabet import CanonicalProteinSequence

DEFAULT_DELIM = ":"

BLOSUM90 = SubstitutionMatrix(
    CanonicalProteinSequence.alphabet,
    CanonicalProteinSequence.alphabet,
    "BLOSUM90",
)

BLOSUM62 = SubstitutionMatrix(
    CanonicalProteinSequence.alphabet,
    CanonicalProteinSequence.alphabet,
    "BLOSUM62",
)


@attrs.define
class Mutation:
    before: str
    pos: int
    after: str

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

    def as_pikaa(self, chain: str = "A") -> str:
        return f"{self.pos} {chain} PIKAA {self.after}"

    def as_str(self) -> str:
        return f"{self.before}{self.pos}{self.after}"

    def reverse(self) -> Mutation:
        return Mutation(self.after, self.pos, self.before)

    def blosum90(self) -> int:
        return int(BLOSUM90.get_score(self.before, self.after))

    def blosum62(self) -> int:
        return int(BLOSUM62.get_score(self.before, self.after))

    @classmethod
    def from_str(cls, string: str) -> Mutation:
        before = string[0]
        pos = int(string[1:-1])
        after = string[-1]
        return cls(before, pos, after)

    @classmethod
    def from_multistr(cls, string: str, delim: str = DEFAULT_DELIM) -> list[Mutation]:
        return [cls.from_str(substr) for substr in string.split(delim)]


def apply_mutations(seq: str, mutations: list[Mutation]) -> str:
    seq_list = list(seq)

    for mut in mutations:
        idx = mut.pos - 1

        if not (0 <= idx < len(seq)):
            raise IndexError(f"{mut} is outside sequence range [1-{len(seq)}]")

        if seq[idx] != mut.before:
            raise ValueError(f"{seq[idx]} in sequence doesn't match mutation {mut}")

        seq_list[idx] = mut.after

    return "".join(seq_list)


def find_mutations(seq1: str, seq2: str) -> list[Mutation]:
    mutations: list[Mutation] = []

    for idx, (before, after) in enumerate(zip(seq1, seq2, strict=True)):
        if before == after:
            continue

        mutations.append(Mutation(before, idx + 1, after))

    return mutations


def get_multi_mutation_string(mutation: list[Mutation], delim: str = DEFAULT_DELIM) -> str:
    return f"{delim}".join(f"{mut}" for mut in mutation)
