# src/svy/regression/_results.py
"""
Result containers for GLM.
"""

from __future__ import annotations

import logging

from typing import TYPE_CHECKING, Any

import msgspec
import numpy as np
import polars as pl

from msgspec import field

from svy.core.containers import FDist, TDist
from svy.utils.formats import _fmt_fixed, _fmt_p, _fmt_smart


if TYPE_CHECKING:
    pass

log = logging.getLogger(__name__)


# =============================================================================
# UI Integration
# =============================================================================


def _ui_module():
    """Try to import central UI module (optional)."""
    try:
        from svy.ui import printing as _p

        return _p
    except Exception:
        return None


def _resolve_width_for(obj, default: int = 88) -> int:
    """Robust width resolution delegating to central UI."""
    P = _ui_module()
    if P:
        try:
            return int(P.resolve_width(owner=obj, default=default))
        except Exception:
            pass
    return default


# =============================================================================
# Result Structs
# =============================================================================


class GLMCoef(msgspec.Struct, frozen=True):
    """A single coefficient in a regression table."""

    term: str
    est: float
    se: float
    lci: float
    uci: float
    wald: TDist | None = None
    wald_adj: TDist | None = None

    def to_dict(self) -> dict[str, Any]:
        return msgspec.to_builtins(self)


class GLMStats(msgspec.Struct, frozen=True):
    """Model-level goodness-of-fit statistics."""

    n: int
    wald: FDist
    wald_adj: FDist
    scale: float
    deviance: float
    aic: float | None = None
    bic: float | None = None
    r_squared: float | None = None
    r_squared_adj: float | None = None
    iterations: int | None = None

    def to_dict(self) -> dict[str, Any]:
        return msgspec.to_builtins(self)


class GLMFit(msgspec.Struct, frozen=True):
    """Immutable container for fitted GLM results."""

    y: str
    family: str
    link: str
    stats: GLMStats
    coefs: list[GLMCoef] = field(default_factory=list)
    cov_matrix: np.ndarray | None = None
    term_info: dict | None = None  # Encoding info for prediction
    feature_names: list[str] = field(default_factory=list)  # Original feature order

    def to_dict(self) -> dict[str, Any]:
        d = msgspec.to_builtins(self)
        # Don't serialize large matrices
        d.pop("cov_matrix", None)
        d.pop("term_info", None)
        return d

    def to_polars(self) -> pl.DataFrame:
        """Convert coefficients to DataFrame."""
        data = []
        for c in self.coefs:
            row: dict[str, Any] = {
                "term": c.term,
                "estimate": c.est,
                "std_err": c.se,
                "conf_low": c.lci,
                "conf_high": c.uci,
            }
            if c.wald:
                row.update(
                    {
                        "statistic": c.wald.value,
                        "p_value": c.wald.p_value,
                        "df": c.wald.df,
                    }
                )
            else:
                row.update({"statistic": None, "p_value": None, "df": None})

            if c.wald_adj:
                row["adj_statistic"] = c.wald_adj.value
                row["adj_p_value"] = c.wald_adj.p_value
                row["adj_df"] = c.wald_adj.df

            data.append(row)
        return pl.DataFrame(data)

    def __rich_console__(self, console, options):
        from rich import box
        from rich.console import Group
        from rich.table import Table as RTable
        from rich.text import Text

        if not self.coefs:
            yield Text("Empty GLM Model", style="red")
            return

        # 1. Stats Grid
        st = self.stats
        df_resid = st.wald.df_den if st.wald else "-"

        stats_rows = [
            ("Observations", str(st.n), "AIC", _fmt_smart(st.aic)),
            ("Df Residuals", str(df_resid), "BIC", _fmt_smart(st.bic)),
            ("Deviance", _fmt_smart(st.deviance), "Scale", _fmt_smart(st.scale)),
        ]

        if st.r_squared is not None:
            stats_rows.append(
                (
                    "R-squared",
                    _fmt_fixed(st.r_squared),
                    "R-sq (adj)",
                    _fmt_fixed(st.r_squared_adj),
                )
            )

        if st.iterations:
            stats_rows.append(("", "", "Iterations", str(st.iterations)))

        if st.wald_adj:
            stats_rows.append(
                (
                    "F-stat (adj)",
                    _fmt_fixed(st.wald_adj.value),
                    "Prob (F-adj)",
                    _fmt_p(st.wald_adj.p_value),
                )
            )

        stats_grid = RTable.grid(padding=(0, 2))
        stats_grid.add_column(style="bold")
        stats_grid.add_column(justify="right")
        stats_grid.add_column(style="bold")
        stats_grid.add_column(justify="right")

        for r in stats_rows:
            stats_grid.add_row(*r)

        # 2. Coefficients Table
        coef_tbl = RTable(
            show_header=True,
            header_style="bold",
            box=box.SIMPLE_HEAVY,
            show_edge=True,
            show_lines=False,
            pad_edge=False,
            expand=False,
        )

        cols = [
            ("Term", "left"),
            ("Coef.", "right"),
            ("Std.Err.", "right"),
            ("t", "right"),
            ("P>|t|", "right"),
            ("[0.025", "right"),
            ("0.975]", "right"),
        ]
        for name, justify in cols:
            coef_tbl.add_column(name, justify=justify)  # type: ignore[arg-type]

        for row in self.coefs:
            t_val = row.wald.value if row.wald else 0.0
            p_val = row.wald.p_value if row.wald else 1.0

            p_style = "bold red" if p_val < 0.05 else ""
            coef_tbl.add_row(
                row.term,
                _fmt_fixed(row.est),
                _fmt_fixed(row.se),
                _fmt_fixed(t_val),
                Text(_fmt_p(p_val), style=p_style),
                _fmt_fixed(row.lci),
                _fmt_fixed(row.uci),
            )

        # 3. Assemble Content
        content = Group(
            Text(f"Modeling: {self.y}", style="dim"),
            Text(""),
            stats_grid,
            Text(""),
            coef_tbl,
        )

        # 4. Wrap in Panel via Central UI
        P = _ui_module()
        title_text = f"GLM: {self.family} ({self.link})"

        if P and hasattr(P, "make_panel"):
            yield P.make_panel([content], title=title_text, kind="estimate")
        else:
            from rich.panel import Panel  # always available inside __rich_console__

            yield Panel(
                content, title=title_text, title_align="left", border_style="cyan", padding=(0, 1)
            )

    def __plain_str__(self) -> str:
        """Plain-text fallback when rich is not installed. Never calls str(self)."""
        lines = [
            f"GLM: {self.family} ({self.link})",
            f"  y          : {self.y}",
            f"  n          : {self.stats.n}",
            f"  deviance   : {_fmt_smart(self.stats.deviance)}",
            f"  AIC        : {_fmt_smart(self.stats.aic)}",
            "",
            "Coefficients:",
            f"  {'Term':<30}  {'Coef':>10}  {'Std.Err':>10}  {'t':>8}  {'P>|t|':>8}  {'[0.025':>10}  {'0.975]':>10}",
            f"  {'-' * 30}  {'-' * 10}  {'-' * 10}  {'-' * 8}  {'-' * 8}  {'-' * 10}  {'-' * 10}",
        ]
        for c in self.coefs:
            t_val = c.wald.value if c.wald else 0.0
            p_val = c.wald.p_value if c.wald else 1.0
            lines.append(
                f"  {c.term:<30}  {_fmt_fixed(c.est):>10}  {_fmt_fixed(c.se):>10}"
                f"  {_fmt_fixed(t_val):>8}  {_fmt_p(p_val):>8}"
                f"  {_fmt_fixed(c.lci):>10}  {_fmt_fixed(c.uci):>10}"
            )
        return "\n".join(lines)

    def __str__(self) -> str:
        P = _ui_module()
        if P:
            return P.render_rich_to_str(self, width=_resolve_width_for(self))
        return self.__plain_str__()

    def __repr__(self) -> str:
        return f"GLMFit(y={self.y!r}, family={self.family!r}, n={self.stats.n}, coefs={len(self.coefs)})"
