from typing import Dict, Optional

import imm

from ._alphabet import Alphabet, AlphabetType, AminoAlphabet, BaseAlphabet
from ._base_lprob import BaseLprob
from ._cdata import CData
from ._codon_lprob import CodonLprob
from ._codon_marg import CodonMarg
from ._ffi import ffi, lib
from ._profile import Profile
from ._state import CodonState, FrameState, StateType

__all__ = ["imm_abc", "imm_state", "nmm_profile"]


def imm_abc(ptr: CData):
    try:
        alphabet_type = AlphabetType(imm.lib.imm_abc_type_id(ptr))
    except ValueError:
        return imm.Alphabet(ptr)

    if alphabet_type == AlphabetType.BASE:
        nmm_base_abc = lib.nmm_base_abc_derived(ptr)
        return BaseAlphabet(nmm_base_abc)

    if alphabet_type == AlphabetType.AMINO:
        nmm_amino_abc = lib.nmm_amino_abc_derived(ptr)
        return AminoAlphabet(nmm_amino_abc)

    raise RuntimeError("It should not get here.")


def imm_state(
    ptr: CData,
    alphabet: Alphabet,
    base_lprobs: Dict[CData, BaseLprob],
    codon_lprobs: Dict[CData, CodonLprob],
    codon_margs: Dict[CData, CodonMarg],
) -> imm.State:
    try:
        state_type = StateType(imm.lib.imm_state_type_id(ptr))
    except ValueError:
        return imm.wrap.imm_state(ptr, alphabet)

    if state_type == StateType.CODON:
        nmm_codon_state = lib.nmm_codon_state_derived(ptr)
        if nmm_codon_state == ffi.NULL:
            raise RuntimeError("`nmm_codon_state` is NULL.")

        nmm_codon_lprob = lib.nmm_codon_state_codon_lprob(nmm_codon_state)
        codonp = codon_lprobs[nmm_codon_lprob]
        return CodonState(nmm_codon_lprob, codonp)

    if state_type == StateType.FRAME:
        nmm_frame_state = lib.nmm_frame_state_derived(ptr)
        if nmm_frame_state == ffi.NULL:
            raise RuntimeError("`nmm_frame_state` is NULL.")

        nmm_base_lprob = lib.nmm_frame_state_base_lprob(nmm_frame_state)
        nmm_codon_marg = lib.nmm_frame_state_codon_marg(nmm_frame_state)

        basep = base_lprobs[nmm_base_lprob]
        codonm = codon_margs[nmm_codon_marg]

        return FrameState(nmm_frame_state, basep, codonm)

    raise ValueError(f"Unknown state type: {imm.lib.imm_state_type_id(ptr)}.")


def nmm_profile(ptr: CData, abc: Optional[BaseAlphabet] = None) -> Profile:
    if abc is None:
        tmp_abc = imm_abc(lib.nmm_profile_abc(ptr))
        assert isinstance(tmp_abc, BaseAlphabet)
        abc = tmp_abc

    base_lprobs = read_base_lprobs(ptr, abc)
    codon_lprobs = read_codon_lprobs(ptr, abc)
    codon_margs = read_cond_margs(ptr, abc)

    prof = Profile(ptr, abc)
    for i in range(lib.nmm_profile_nmodels(ptr)):
        imm_model = lib.nmm_profile_get_model(ptr, i)
        states = {}
        for j in range(imm.lib.imm_model_nstates(imm_model)):
            state_ptr = imm.lib.imm_model_state(imm_model, j)
            states[state_ptr] = imm_state(
                state_ptr, abc, base_lprobs, codon_lprobs, codon_margs
            )

        hmm = imm.HMM(imm.lib.imm_model_hmm(imm_model), abc, states)
        dp = imm.DP(imm.lib.imm_model_dp(imm_model), hmm)
        model = imm.Model(imm_model, hmm, dp)
        prof.append_model(model)

    return prof


def read_base_lprobs(nmm_profile: CData, abc: BaseAlphabet) -> Dict[CData, BaseLprob]:
    base_lprobs: Dict[CData, BaseLprob] = {}
    for i in range(lib.nmm_profile_nbase_lprobs(nmm_profile)):
        nmm_base_lprob = lib.nmm_profile_base_lprob(nmm_profile, i)
        base_lprobs[nmm_base_lprob] = BaseLprob(nmm_base_lprob, abc)
    return base_lprobs


def read_codon_lprobs(nmm_profile: CData, abc: BaseAlphabet) -> Dict[CData, CodonLprob]:
    codon_lprobs: Dict[CData, CodonLprob] = {}
    for i in range(lib.nmm_profile_ncodon_lprobs(nmm_profile)):
        nmm_codon_lprob = lib.nmm_profile_codon_lprob(nmm_profile, i)
        codon_lprobs[nmm_codon_lprob] = CodonLprob(nmm_codon_lprob, abc)
    return codon_lprobs


def read_cond_margs(nmm_profile: CData, abc: BaseAlphabet) -> Dict[CData, CodonMarg]:
    codon_margs: Dict[CData, CodonMarg] = {}
    for i in range(lib.nmm_profile_ncodon_margs(nmm_profile)):
        nmm_codon_marg = lib.nmm_profile_codon_marg(nmm_profile, i)
        codon_margs[nmm_codon_marg] = CodonMarg(nmm_codon_marg, abc)
    return codon_margs
