"""
Diagnostics for mixed models: intra-class correlation, information
criteria, and the Nakagawa-Schielzeth R² (delegated to the
``MixedResult`` / ``MEGLMResult`` classes themselves).
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import Optional

import numpy as np
from scipy import stats


@dataclass
class ICCResult:
    """Container for an ICC estimate with a delta-method Wald CI."""

    estimate: float
    se: float
    ci_lower: float
    ci_upper: float
    alpha: float

    def summary(self) -> str:
        return (
            f"ICC = {self.estimate:.4f}  (SE = {self.se:.4f}, "
            f"{100*(1-self.alpha):.0f}% CI [{self.ci_lower:.4f}, "
            f"{self.ci_upper:.4f}])"
        )

    # Allow ``float(icc(...))`` to recover the point estimate.
    def __float__(self) -> float:
        return float(self.estimate)


def icc(
    result,
    component: str = "_cons",
    alpha: float = 0.05,
    n_boot: int = 0,
    seed: Optional[int] = None,
) -> ICCResult:
    """
    Intra-class correlation for a fitted mixed model.

    Parameters
    ----------
    result
        A ``MixedResult`` returned by :func:`statspai.mixed`.
    component
        Name of the random-effect variance to put in the numerator.
        Defaults to the random intercept (``"_cons"``).
    alpha
        Significance level for the confidence interval.  Default 0.05.
    n_boot
        Number of parametric bootstrap replicates used to compute the
        CI.  ``0`` (default) uses the delta-method approximation on the
        log-variance scale, which is faster and usually within a few
        decimals of the parametric-bootstrap answer for moderate N.
    seed
        RNG seed forwarded to :func:`numpy.random.default_rng`.

    Returns
    -------
    ICCResult
    """
    if not hasattr(result, "variance_components"):
        raise TypeError("icc() expects a MixedResult-like object")

    key = f"var({component})"
    if key not in result.variance_components:
        raise KeyError(f"variance component {key!r} not found; "
                       f"available: {list(result.variance_components)}")

    var_u = float(result.variance_components[key])
    var_e = float(result.variance_components.get("var(Residual)", np.nan))
    total = var_u + var_e
    if total <= 0 or not np.isfinite(total):
        return ICCResult(np.nan, np.nan, np.nan, np.nan, alpha)

    rho = var_u / total

    if n_boot and n_boot > 0:
        raise NotImplementedError(
            "parametric-bootstrap ICC intervals are not implemented yet; "
            "call icc(result, n_boot=0) for the delta-method CI."
        )

    # Delta method on logit scale keeps the CI inside [0,1].
    # Var(log var_u) ≈ 2 / dof_u ; approximate dof from group count.
    # This is a crude heuristic that treats var_u as χ²-distributed with
    # (n_groups - 1) df; it gets tight CIs wrong for small samples and
    # ignores the correlation between var_u and var_e under (RE)ML —
    # flagged below so users don't over-interpret the bounds.
    n_groups = max(getattr(result, "n_groups", 1), 1)
    if n_groups < 30:
        warnings.warn(
            f"icc(): delta-method CI is unreliable for n_groups="
            f"{n_groups} < 30 (treats var(u) as χ²-distributed with "
            "n_groups df).  Interpret the bounds cautiously.",
            RuntimeWarning,
            stacklevel=2,
        )
    var_log_var_u = 2.0 / n_groups
    var_log_var_e = 2.0 / max(result.n_obs - result.n_fixed, 1)
    # logit(rho) = log(var_u) - log(var_e)
    var_logit = var_log_var_u + var_log_var_e
    se_logit = np.sqrt(var_logit)
    z = stats.norm.ppf(1 - alpha / 2)
    logit_rho = np.log(rho / (1 - rho)) if 0 < rho < 1 else np.nan
    if np.isfinite(logit_rho):
        lo = 1.0 / (1.0 + np.exp(-(logit_rho - z * se_logit)))
        hi = 1.0 / (1.0 + np.exp(-(logit_rho + z * se_logit)))
    else:
        lo, hi = np.nan, np.nan
    se_rho = se_logit * rho * (1 - rho)
    return ICCResult(rho, se_rho, lo, hi, alpha)


__all__ = ["icc", "ICCResult"]
