"""
bishop.graphical_models – PRML Chapter 8
=========================================
Wraps pgmpy for Bayesian networks and hmmlearn for HMMs.

    from ml_ai_library.bishop.graphical_models import (
        BayesianNetwork,   # pgmpy wrapper
        GaussianHMM,       # hmmlearn wrapper
        CategoricalHMM,    # hmmlearn wrapper
        viterbi_decode,    # convenience helper
    )
"""

from __future__ import annotations
from typing import List, Tuple, Sequence

import numpy as np

# pgmpy – Bayesian networks (PRML §8.1) ------------------------------------
from pgmpy.models import BayesianNetwork                          # PRML §8.1
from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination, BeliefPropagation

# hmmlearn – Hidden Markov Models (PRML §13.2) ------------------------------
from hmmlearn.hmm import GaussianHMM, CategoricalHMM


# ---------------------------------------------------------------------------
# Bayesian Network helpers
# ---------------------------------------------------------------------------

def make_bayesian_network(
    edges: List[Tuple[str, str]],
    cpds: List[TabularCPD],
) -> BayesianNetwork:
    """
    Build and validate a pgmpy BayesianNetwork.

    Parameters
    ----------
    edges : list of (parent, child) tuples
    cpds  : list of TabularCPD objects (one per node)

    Example
    -------
    >>> edges = [("Rain", "Wet"), ("Sprinkler", "Wet")]
    >>> cpd_rain = TabularCPD("Rain", 2, [[0.2], [0.8]])
    >>> net = make_bayesian_network(edges, [cpd_rain, ...])
    >>> infer = VariableElimination(net)
    >>> result = infer.query(["Wet"], evidence={"Rain": 1})
    """
    model = BayesianNetwork(edges)
    model.add_cpds(*cpds)
    assert model.check_model(), "CPDs do not form a valid BN"
    return model


def exact_inference(
    model: BayesianNetwork, query: List[str], evidence: dict
):
    """Variable elimination exact inference (PRML §8.4.2)."""
    return VariableElimination(model).query(query, evidence=evidence,
                                            show_progress=False)


def loopy_bp_inference(
    model: BayesianNetwork, query: List[str], evidence: dict
):
    """Loopy belief propagation approximate inference (PRML §8.4.7)."""
    return BeliefPropagation(model).query(query, evidence=evidence,
                                          show_progress=False)


# ---------------------------------------------------------------------------
# HMM helpers
# ---------------------------------------------------------------------------

def viterbi_decode(model: GaussianHMM, obs: np.ndarray) -> np.ndarray:
    """
    Most-probable state sequence via Viterbi algorithm (PRML §13.2.5).

    Parameters
    ----------
    model : fitted GaussianHMM or CategoricalHMM
    obs   : (T, D) observation sequence

    Returns
    -------
    states : (T,) integer array of most-probable hidden states
    """
    _, states = model.decode(obs, algorithm="viterbi")
    return states
