from __future__ import annotations

from enum import Enum
from typing import Tuple, Type

from imm import Sequence, State

from ._alphabet import BaseAlphabet
from ._base_lprob import BaseLprob
from ._cdata import CData
from ._codon import Codon
from ._codon_lprob import CodonLprob
from ._codon_marg import CodonMarg
from ._ffi import ffi, lib

__all__ = [
    "CodonState",
    "FrameState",
    "StateType",
]


class StateType(Enum):
    CODON = 0x10
    FRAME = 0x11


class FrameState(State):
    def __init__(
        self,
        nmm_frame_state: CData,
        basep: BaseLprob,
        codonm: CodonMarg,
    ):
        """
        Frame state.

        Parameters
        ----------
        nmm_frame_state
            State pointer.
        basep
            Base of log-probabilities.
        codonm
            Codon marginals.
        """
        self._nmm_frame_state = nmm_frame_state
        if self._nmm_frame_state == ffi.NULL:
            raise RuntimeError("`nmm_frame_state` is NULL.")
        self._basep = basep
        self._codonm = codonm
        alphabet = basep.alphabet
        super().__init__(lib.nmm_frame_state_super(self._nmm_frame_state), alphabet)

    @classmethod
    def create(
        cls: Type[FrameState],
        name: bytes,
        basep: BaseLprob,
        codonm: CodonMarg,
        epsilon: float,
    ) -> FrameState:
        """
        Create frame state.

        Parameters
        ----------
        name
            State name.
        basep
            Base log-probabilities.
        codonm
            Codon marginals.
        epsilon
            Epsilon.
        """
        ptr = lib.nmm_frame_state_create(
            name, basep.nmm_base_lprob, codonm.nmm_codon_marg, epsilon
        )
        return cls(ptr, basep, codonm)

    @property
    def base_lprob(self) -> BaseLprob:
        return self._basep

    @property
    def codon_marg(self) -> CodonMarg:
        return self._codonm

    def decode(self, seq: Sequence) -> Tuple[float, Codon]:
        state = self._nmm_frame_state
        any_symbol = self.alphabet.any_symbol
        assert isinstance(self.alphabet, BaseAlphabet)
        codon = Codon.create(any_symbol * 3, self.alphabet)
        lprob = lib.nmm_frame_state_decode(state, seq.imm_seq, codon.nmm_codon)
        return lprob, codon

    @property
    def epsilon(self) -> float:
        return lib.nmm_frame_state_epsilon(self._nmm_frame_state)

    def __del__(self):
        if self._nmm_frame_state != ffi.NULL:
            lib.nmm_frame_state_destroy(self._nmm_frame_state)

    def __repr__(self):
        return f"<{self.__class__.__name__}:{str(self)}>"


class CodonState(State):
    def __init__(self, nmm_codon_state: CData, codonp: CodonLprob):
        """
        Codon state.

        Parameters
        ----------
        nmm_codon_state
            State pointer.
        codonp
            Codon probabilities.
        """
        self._nmm_codon_state = nmm_codon_state
        if self._nmm_codon_state == ffi.NULL:
            raise RuntimeError("`nmm_codon_state` is NULL.")
        self._codonp = codonp
        alphabet = codonp.alphabet
        super().__init__(lib.nmm_codon_state_super(nmm_codon_state), alphabet)

    @classmethod
    def create(cls: Type[CodonState], name: bytes, codonp: CodonLprob) -> CodonState:
        """
        Create codon state.

        Parameters
        ----------
        name
            State name.
        codonp
            Codon probabilities.
        """
        ptr = lib.nmm_codon_state_create(name, codonp.nmm_codon_lprob)
        return cls(ptr, codonp)

    def __del__(self):
        if self._nmm_codon_state != ffi.NULL:
            lib.nmm_codon_state_destroy(self._nmm_codon_state)

    def __repr__(self):
        return f"<{self.__class__.__name__}:{str(self)}>"
