# src/svy/regression/_glm.py
"""
GLM fitting for complex survey data.
"""

from __future__ import annotations

import logging
import math

from typing import TYPE_CHECKING, Sequence, cast

import numpy as np
import polars as pl

from scipy import stats


# Rust backend
try:
    from svy_rs import _internal as rs
except ImportError:
    import svy_rs as rs

from svy.core.constants import _INTERNAL_CONCAT_SUFFIX
from svy.core.containers import FDist, TDist
from svy.core.enumerations import DistFamily, LinkFunction
from svy.core.terms import Cat, Cross, Feature
from svy.errors.model_errors import ModelError
from svy.regression.base import link_inverse, link_mu_eta, resolve_link
from svy.regression.prediction import GLMPred
from svy.regression.results import GLMCoef, GLMFit, GLMStats


if TYPE_CHECKING:
    from svy.core.sample import Sample
    from svy.regression.margins import GLMMargins

log = logging.getLogger(__name__)


class GLM:
    """
    Generalized Linear Models for Complex Survey Data.

    Access via `sample.glm`.

    Examples
    --------
    >>> sample.glm.fit(y="outcome", x=["age", Cat("region")])
    >>> sample.glm.predict(new_data)
    """

    PRINT_WIDTH: int | None = None  # Class-level print width

    __slots__ = ("_sample", "fitted", "_print_width")

    def __init__(self, sample: Sample) -> None:
        self._sample = sample
        self.fitted: GLMFit | None = None
        self._print_width: int | None = None

    def _ensure_fitted(self) -> GLMFit:
        if self.fitted is None:
            raise ModelError.not_fitted(where="GLM")
        return self.fitted

    # --- Print Width Configuration ---

    def set_print_width(self, width: int | None) -> "GLM":
        """Set the print width for this specific instance."""
        if width is None:
            self._print_width = None
            return self
        try:
            w = int(width)
        except Exception as ex:
            raise TypeError(f"print width must be int or None; got {width!r}") from ex
        if w <= 20:
            raise ValueError("print width must be > 20 characters.")
        self._print_width = w
        return self

    @classmethod
    def set_default_print_width(cls, width: int | None) -> None:
        """Set the default print width for all GLM instances."""
        if width is None:
            cls.PRINT_WIDTH = None
            return
        try:
            w = int(width)
        except Exception as ex:
            raise TypeError(f"class print width must be int or None; got {width!r}") from ex
        if w <= 20:
            raise ValueError("class print width must be > 20 characters.")
        cls.PRINT_WIDTH = w

    # --- Display Methods ---

    def __plain_str__(self) -> str:
        """Plain-text fallback when rich is not installed. Never calls str(self)."""
        if self.fitted is None:
            return "GLM (not fitted)"
        fit = self.fitted
        lines = [
            f"GLM: {fit.family} ({fit.link})",
            f"  y = {fit.y!r}",
            f"  n = {fit.stats.n}",
            "",
            "Coefficients:",
        ]
        for c in fit.coefs:
            lines.append(
                f"  {c.term:<30}  est={c.est:>10.4f}  se={c.se:>10.4f}  [{c.lci:.4f}, {c.uci:.4f}]"
            )
        return "\n".join(lines)

    def __str__(self) -> str:
        from svy.ui.printing import render_rich_to_str, resolve_width

        return render_rich_to_str(self, width=resolve_width(self))

    def __repr__(self) -> str:
        if self.fitted:
            return (
                f"<GLM: Fitted {self.fitted.family} ({self.fitted.link}) on y '{self.fitted.y}'>"
            )
        return "<GLM: Unfitted>"

    def __rich_console__(self, console, options):
        if self.fitted:
            yield from self.fitted.__rich_console__(console, options)
        else:
            from rich.text import Text

            yield Text("GLM (not fitted)", style="dim")

    # --- Properties ---

    @property
    def coefs(self) -> list[GLMCoef]:
        return self._ensure_fitted().coefs

    @property
    def stats(self) -> GLMStats:
        return self._ensure_fitted().stats

    def fit(
        self,
        y: str,
        *,
        x: Sequence[Feature] | None = None,
        intercept: bool = True,
        family: DistFamily | str = "gaussian",
        link: LinkFunction | str | None = None,
        drop_nulls: bool = True,
        tol: float = 1e-8,
        max_iter: int = 100,
        alpha: float = 0.05,
    ) -> GLM:
        """
        Fit a GLM to the survey data.

        Parameters
        ----------
        y : str
            Response variable column name.
        x : sequence of Feature, optional
            Predictor variables. Can be strings or Cat/Cross terms.
        intercept : bool
            Include intercept term.
        family : str or DistFamily
            Distribution family (gaussian, binomial, poisson, gamma).
        link : str or LinkFunction, optional
            Link function. If None, uses canonical link.
        drop_nulls : bool
            Drop rows with missing values.
        tol : float
            Convergence tolerance.
        max_iter : int
            Maximum IRLS iterations.
        alpha : float
            Significance level for confidence intervals.

        Returns
        -------
        GLM
            Self, with fitted results in `.fitted`.
        """
        # Resolve family/link
        fam_str = str(family.value if isinstance(family, DistFamily) else family).lower()
        link_str = resolve_link(fam_str, link.value if isinstance(link, LinkFunction) else link)

        # Prepare data
        df_lazy = self._sample._data
        if isinstance(df_lazy, pl.DataFrame):
            df_lazy = df_lazy.lazy()

        design = self._sample._design

        # Handle design columns
        df_clean, (_, s_col, p_col, _) = self._sample._create_concatenated_cols_from_lists(
            data=df_lazy,
            design=design,
            by=None,
            null_token="__Null__",
            suffix=_INTERNAL_CONCAT_SUFFIX,
        )
        if isinstance(df_clean, pl.DataFrame):
            df_clean = df_clean.lazy()

        # Weight column
        w_col = design.wgt or "_svy_ones_"
        if design.wgt is None:
            df_clean = df_clean.with_columns(pl.lit(1.0).alias(w_col))
        else:
            df_clean = df_clean.filter(
                (pl.col(w_col) > 0) & pl.col(w_col).is_not_null() & pl.col(w_col).is_finite()
            )

        # Build features
        feature_specs = list(x) if x else []
        required_cols = {y}
        feature_exprs: list[pl.Expr] = []
        feature_names: list[str] = []
        term_info: dict = {}

        if intercept:
            df_clean = df_clean.with_columns(pl.lit(1.0).alias("_intercept_"))
            feature_exprs.append(pl.col("_intercept_"))
            feature_names.append("_intercept_")

        def _resolve_feature(feat: Feature) -> tuple[list[pl.Expr], list[str]]:
            if isinstance(feat, str):
                required_cols.add(feat)
                term_info[feat] = {"type": "continuous"}
                return [pl.col(feat).cast(pl.Float64)], [feat]

            elif isinstance(feat, Cat):
                required_cols.add(feat.name)
                unique_df: pl.DataFrame = cast(
                    pl.DataFrame, df_clean.select(feat.name).unique().collect()
                )
                levels = sorted(unique_df.get_column(feat.name).drop_nulls().to_list())

                if len(levels) < 2:
                    log.warning(f"Categorical '{feat.name}' has < 2 levels. Dropped.")
                    return [], []

                ref_val = feat.ref if feat.ref in levels else levels[0]
                levels = [ref_val] + [v for v in levels if v != ref_val]

                term_info[feat.name] = {
                    "type": "categorical",
                    "levels": levels,
                    "ref": ref_val,
                }

                exprs, names = [], []
                for level in levels[1:]:
                    name = f"{feat.name}_{level}"
                    expr = (pl.col(feat.name) == level).cast(pl.Float64).alias(name)
                    exprs.append(expr)
                    names.append(name)
                return exprs, names

            elif isinstance(feat, Cross):
                left_exprs, left_names = _resolve_feature(feat.left)
                right_exprs, right_names = _resolve_feature(feat.right)

                if not left_exprs or not right_exprs:
                    return [], []

                exprs, names = [], []
                for le, ln in zip(left_exprs, left_names):
                    for re, rn in zip(right_exprs, right_names):
                        name = f"{ln}:{rn}"
                        exprs.append((le * re).alias(name))
                        names.append(name)
                return exprs, names

            else:
                raise TypeError(f"Unknown feature type: {type(feat)}")

        for feat in feature_specs:
            exprs, names = _resolve_feature(feat)
            feature_exprs.extend(exprs)
            feature_names.extend(names)

        # Handle nulls
        if drop_nulls:
            df_clean = df_clean.drop_nulls(subset=list(required_cols))

        # Build final selection
        final_selects = (
            [pl.col(y).cast(pl.Float64)] + feature_exprs + [pl.col(w_col).cast(pl.Float64)]
        )

        # Unwrap list columns
        if isinstance(s_col, (list, tuple)):
            s_col = s_col[0] if s_col else None
        if isinstance(p_col, (list, tuple)):
            p_col = p_col[0] if p_col else None

        if s_col:
            final_selects.append(pl.col(s_col))
        if p_col:
            final_selects.append(pl.col(p_col))

        # Collect
        try:
            eng_df: pl.DataFrame = cast(pl.DataFrame, df_clean.select(final_selects).collect())
        except Exception as e:
            raise ValueError(f"Failed to prepare data: {e}")

        # Validate
        y_data = eng_df.get_column(y)
        self._validate_response(y_data, fam_str)

        # Call Rust
        try:
            res = rs.fit_glm_rs(
                y, feature_names, w_col, s_col, p_col, fam_str, link_str, tol, max_iter, eng_df
            )
        except Exception as e:
            raise RuntimeError(f"Rust GLM engine failed: {e}") from e

        beta, cov_flat, scale, df_resid, dev, null_dev, iters = res

        # Post-process
        k = len(feature_names)
        n = eng_df.height
        cov_mat = np.array(cov_flat).reshape(k, k)
        se_arr = np.sqrt(np.diag(cov_mat))
        params_arr = np.array(beta)

        # Design DF
        df_design = self._compute_design_df(eng_df, s_col, p_col, n)

        # Statistics
        stats_struct = self._build_stats(n, k, scale, dev, null_dev, iters, df_design, intercept)

        # Coefficients
        t_crit = stats.t.ppf(1 - alpha / 2, df_design)
        coef_list = []
        for i, name in enumerate(feature_names):
            est, se = params_arr[i], se_arr[i]
            t_val = est / se if se > 0 else 0.0
            p_val = 2 * stats.t.sf(abs(t_val), df_design)

            coef_list.append(
                GLMCoef(
                    term=name,
                    est=est,
                    se=se,
                    lci=est - t_crit * se,
                    uci=est + t_crit * se,
                    wald=TDist(value=t_val, df=df_design, p_value=p_val),
                )
            )

        self.fitted = GLMFit(
            y=y,
            family=fam_str.capitalize(),
            link=link_str,
            stats=stats_struct,
            coefs=coef_list,
            cov_matrix=cov_mat,
            term_info=term_info,
            feature_names=feature_names,
        )
        return self

    def predict(
        self,
        new_data: pl.DataFrame,
        alpha: float = 0.05,
        y_col: str | None = None,
    ) -> GLMPred:
        fit = self._ensure_fitted()

        # Build design matrix
        X = self._build_prediction_matrix(new_data, fit)

        # Extract model info
        beta = np.array([c.est for c in fit.coefs])
        cov = fit.cov_matrix
        df = fit.coefs[0].wald.df if fit.coefs[0].wald else 1e6
        if cov is None:
            raise ValueError("Covariance matrix not available")

        # Linear predictor
        eta = X @ beta
        var_eta = np.sum((X @ cov) * X, axis=1)
        se_eta = np.sqrt(np.maximum(var_eta, 0))

        t_crit = stats.t.ppf(1 - alpha / 2, df)
        lci_eta = eta - t_crit * se_eta
        uci_eta = eta + t_crit * se_eta

        # Use functions with fit.link
        yhat = link_inverse(fit.link, eta)
        lci = link_inverse(fit.link, lci_eta)
        uci = link_inverse(fit.link, uci_eta)

        dmu_deta = link_mu_eta(fit.link, eta)
        se = se_eta * np.abs(dmu_deta)

        # Residuals (if y available)
        residuals = None
        if y_col is not None and y_col in new_data.columns:
            y = new_data.get_column(y_col).cast(pl.Float64).to_numpy()
            residuals = y - yhat

        return GLMPred(
            yhat=yhat,
            se=se,
            lci=lci,
            uci=uci,
            df=df,
            residuals=residuals,
            alpha=alpha,
        )

    def margins(
        self,
        at: dict[str, list] | None = None,
        variables: list[str] | None = None,
        alpha: float = 0.05,
    ) -> "GLMMargins | list[GLMMargins]":
        from svy.regression.margins import margins as compute_margins

        return compute_margins(self, at=at, variables=variables, alpha=alpha)

    def _build_prediction_matrix(
        self,
        new_data: pl.DataFrame,
        fit: GLMFit,
    ) -> np.ndarray:
        """Build design matrix for prediction."""
        n = new_data.height
        terms = [c.term for c in fit.coefs]
        term_info = fit.term_info or {}
        k = len(terms)

        X = np.zeros((n, k))

        for j, term in enumerate(terms):
            X[:, j] = self._resolve_pred_term(term, new_data, term_info)

        return X

    def _resolve_pred_term(
        self,
        term: str,
        data: pl.DataFrame,
        term_info: dict,
    ) -> np.ndarray:
        """Resolve a single term for prediction."""
        n = data.height

        # Intercept
        if term == "_intercept_":
            return np.ones(n)

        # Interaction
        if ":" in term:
            parts = term.split(":")
            result = np.ones(n)
            for part in parts:
                result = result * self._resolve_pred_term(part, data, term_info)
            return result

        # Categorical dummy
        for var_name, info in term_info.items():
            if info.get("type") == "categorical":
                prefix = f"{var_name}_"
                if term.startswith(prefix):
                    level = term[len(prefix) :]
                    if level in info["levels"]:
                        col = data.get_column(var_name)
                        return self._compare_level(col, level)

        # Continuous
        if term in data.columns:
            return data.get_column(term).cast(pl.Float64).to_numpy()

        # Fallback: var_level pattern
        if "_" in term:
            for i in range(len(term) - 1, 0, -1):
                if term[i] == "_":
                    var = term[:i]
                    level = term[i + 1 :]
                    if var in data.columns:
                        col = data.get_column(var)
                        return self._compare_level(col, level)

        raise KeyError(f"Cannot resolve term '{term}'")

    def _compare_level(self, col: pl.Series, level) -> np.ndarray:
        """Compare column to level, handling type coercion."""
        col_np = col.to_numpy()

        try:
            return (col_np == level).astype(np.float64)
        except (TypeError, ValueError):
            pass

        try:
            level_num = float(level) if "." in str(level) else int(level)
            return (col_np == level_num).astype(np.float64)
        except (TypeError, ValueError):
            pass

        return (col_np.astype(str) == str(level)).astype(np.float64)

    # =========================================================================
    # Helpers
    # =========================================================================

    def _validate_response(self, y_data: pl.Series, family: str) -> None:
        """Validate response variable for family."""
        if family == "binomial":
            unique = set(y_data.unique().to_list())
            if not unique <= {0, 1, 0.0, 1.0}:
                raise ModelError.domain_violation(
                    where="GLM.fit",
                    family=family,
                    violation="Non-binary target",
                    hint="Must be 0/1",
                )
        elif family in ("gamma", "inversegaussian") and y_data.min() <= 0:  # type: ignore[operator]
            raise ModelError.domain_violation(
                where="GLM.fit", family=family, violation="Non-positive values"
            )
        elif family == "poisson" and y_data.min() < 0:  # type: ignore[operator]
            raise ModelError.domain_violation(
                where="GLM.fit", family=family, violation="Negative values"
            )

    def _compute_design_df(
        self, df: pl.DataFrame, s_col: str | None, p_col: str | None, n: int
    ) -> int:
        """Compute design-based degrees of freedom."""
        if p_col and s_col:
            psu_per_stratum = df.group_by(s_col).agg(pl.col(p_col).n_unique())
            total_psu = psu_per_stratum.get_column(p_col).sum()
            n_strata = df.get_column(s_col).n_unique()
            return int(max(1, total_psu - n_strata))
        elif p_col:
            return max(1, df.get_column(p_col).n_unique() - 1)
        elif s_col:
            return max(1, n - df.get_column(s_col).n_unique())
        else:
            return max(1, n - 1)

    def _build_stats(
        self,
        n: int,
        k: int,
        scale: float,
        dev: float,
        null_dev: float,
        iters: int,
        df_design: int,
        intercept: bool,
    ) -> GLMStats:
        """Build model statistics."""
        aic = dev + 2 * k
        bic = aic - 2 * k + k * math.log(n) if n > 0 else None

        r2 = 1.0 - dev / null_dev if null_dev > 1e-12 else None
        r2_adj = None
        if r2 is not None:
            p = k - (1 if intercept else 0)
            if n > p + 1:
                r2_adj = 1.0 - (1.0 - r2) * (n - 1) / (n - p - 1)

        return GLMStats(
            n=n,
            wald=FDist(value=0.0, df_num=0.0, df_den=df_design, p_value=1.0),
            wald_adj=FDist(value=0.0, df_num=0.0, df_den=df_design, p_value=1.0),
            scale=scale,
            deviance=dev,
            aic=aic,
            bic=bic,
            r_squared=r2,
            r_squared_adj=r2_adj,
            iterations=iters,
        )
