from __future__ import annotations

from typing import Type

from imm import lprob_is_valid

from ._alphabet import BaseAlphabet
from ._cdata import CData
from ._codon import Codon
from ._codon_lprob import CodonLprob
from ._ffi import ffi, lib

__all__ = ["CodonMarg"]


class CodonMarg:
    """
    Codon marginals.

    Compute marginal and non-marginal codon probabilities.

    Parameters
    ----------
    nmm_codon_marg
        Codon marginals.
    alphabet
        Four-nucleotides alphabet.
    """

    def __init__(self, nmm_codon_marg: CData, alphabet: BaseAlphabet):
        self._nmm_codon_marg = nmm_codon_marg
        if self._nmm_codon_marg == ffi.NULL:
            raise RuntimeError("`nmm_codon_marg` is NULL.")
        self._alphabet = alphabet

    @classmethod
    def create(cls: Type[CodonMarg], codonp: CodonLprob) -> CodonMarg:
        """
        Create codon marginals.

        Parameters
        ----------
        codonp
            Non-marginal codon probabilities.
        """
        return cls(lib.nmm_codon_marg_create(codonp.nmm_codon_lprob), codonp.alphabet)

    @property
    def alphabet(self) -> BaseAlphabet:
        return self._alphabet

    @property
    def nmm_codon_marg(self) -> CData:
        return self._nmm_codon_marg

    def lprob(self, codon: Codon) -> float:
        lprob: float = lib.nmm_codon_marg_lprob(self._nmm_codon_marg, codon.nmm_codon)
        if not lprob_is_valid(lprob):
            raise RuntimeError("Could not get probability.")
        return lprob

    def __del__(self):
        if self._nmm_codon_marg != ffi.NULL:
            lib.nmm_codon_marg_destroy(self._nmm_codon_marg)
