84class DistributionalGLM:
 85    """
 86    GAMLSS model: model ALL distribution parameters as functions of covariates.
 87
 88    Standard GLM models only E[Y|X]. This models the full conditional
 89    distribution p(Y|X) by allowing each parameter (mean, variance, shape)
 90    to depend on covariates.
 91
 92    Example
 93    -------
 94    >>> import polars as pl
 95    >>> from insurance_distributional_glm import DistributionalGLM
 96    >>> from insurance_distributional_glm.families import Gamma
 97    >>>
 98    >>> df = pl.DataFrame({"age": [25, 35, 45], "vehicle_value": [5000, 15000, 30000]})
 99    >>> y = np.array([500.0, 1200.0, 3500.0])
100    >>>
101    >>> model = DistributionalGLM(
102    ...     family=Gamma(),
103    ...     formulas={"mu": ["age", "vehicle_value"], "sigma": ["age"]},
104    ... )
105    >>> model.fit(df, y)
106    >>> model.predict(df, parameter="mu")
107    """
108
109    def __init__(
110        self,
111        family: GamlssFamily,
112        formulas: Optional[Dict[str, List[str]]] = None,
113    ):
114        """
115        Parameters
116        ----------
117        family : GamlssFamily
118            Distribution family. Determines which parameters exist and their
119            default link functions.
120        formulas : dict, optional
121            Maps parameter name -> list of column names to include.
122            Intercept is always added automatically.
123            If None: mu gets all columns, other parameters get intercept-only.
124            If {'mu': [], 'sigma': []}: intercept-only for all parameters.
125        """
126        self.family = family
127        self.formulas = formulas
128        self._betas: Dict[str, np.ndarray] = {}
129        self._params: Dict[str, np.ndarray] = {}
130        self._columns: List[str] = []
131        self._loglik: float = float("nan")
132        self._loglik_history: List[float] = []
133        self._converged: bool = False
134        self._n: int = 0
135        self._fitted = False
136
137    def fit(
138        self,
139        X,
140        y,
141        exposure: Optional[np.ndarray] = None,
142        weights: Optional[np.ndarray] = None,
143        max_iter: int = 100,
144        tol: float = 1e-6,
145        verbose: bool = False,
146    ) -> "DistributionalGLM":
147        """
148        Fit the model via the RS algorithm.
149
150        Parameters
151        ----------
152        X : polars/pandas DataFrame or numpy array
153            Covariate matrix. Column names must match the formulas dict.
154        y : array-like
155            Response observations. Must be positive for Gamma/LogNormal/IG.
156        exposure : array-like, optional
157            Exposure values (not log). Applied as log-offset on mu.
158        weights : array-like, optional
159            Observation weights.
160        max_iter : int
161            Maximum RS outer iterations (default 100).
162        tol : float
163            Convergence tolerance on log-likelihood change (default 1e-6).
164        verbose : bool
165            Print log-likelihood at each RS iteration.
166
167        Returns
168        -------
169        self
170        """
171        y = np.asarray(y, dtype=float)
172        self._n = len(y)
173        self._columns = _get_columns(X)
174
175        # Resolve formulas
176        formulas = self._resolve_formulas()
177
178        # Build design matrices
179        design_matrices = {}
180        for pname in self.family.param_names:
181            cols = formulas.get(pname, [])
182            design_matrices[pname] = _build_design_matrix(X, cols)
183
184        # Exposure offset
185        log_offset = None
186        if exposure is not None:
187            exposure = np.asarray(exposure, dtype=float)
188            log_offset = np.log(np.clip(exposure, 1e-10, None))
189
190        if weights is not None:
191            weights = np.asarray(weights, dtype=float)
192
193        # Run RS algorithm
194        betas, params, ll_history, converged = rs_fit(
195            family=self.family,
196            design_matrices=design_matrices,
197            y=y,
198            log_offset=log_offset,
199            weights=weights,
200            max_iter=max_iter,
201            tol=tol,
202            verbose=verbose,
203        )
204
205        self._betas = betas
206        self._params = params
207        self._loglik_history = ll_history
208        self._loglik = ll_history[-1] if ll_history else float("nan")
209        self._converged = converged
210        self._design_matrices = design_matrices
211        self._formulas = formulas
212        self._fitted = True
213
214        if not converged and verbose:
215            print(f"Warning: RS algorithm did not converge in {max_iter} iterations.")
216
217        return self
218
219    def _resolve_formulas(self) -> Dict[str, List[str]]:
220        """Apply default formula logic if formulas is None."""
221        param_names = self.family.param_names
222        if self.formulas is None:
223            result = {}
224            for i, pname in enumerate(param_names):
225                result[pname] = self._columns if i == 0 else []
226            return result
227        # Fill in missing params with intercept-only
228        result = {}
229        for pname in param_names:
230            result[pname] = self.formulas.get(pname, [])
231        return result
232
233    def predict(
234        self,
235        X,
236        parameter: str = "mu",
237        exposure: Optional[np.ndarray] = None,
238    ) -> np.ndarray:
239        """
240        Predict a distribution parameter on the response scale.
241
242        Parameters
243        ----------
244        X : DataFrame or array
245            Covariates.
246        parameter : str
247            Which parameter to predict (e.g. 'mu', 'sigma', 'pi').
248        exposure : array-like, optional
249            Exposure offset for mu predictions.
250
251        Returns
252        -------
253        ndarray, shape (n,)
254            Predicted parameter values on response scale.
255        """
256        self._check_fitted()
257        cols = self._formulas.get(parameter, [])
258        X_mat = _build_design_matrix(X, cols)
259        link = self.family.default_links[parameter]
260
261        eta = X_mat @ self._betas[parameter]
262        if parameter == "mu" and exposure is not None:
263            exposure = np.asarray(exposure, dtype=float)
264            eta = eta + np.log(np.clip(exposure, 1e-10, None))
265
266        return link.inverse(eta)
267
268    def predict_distribution(self, X, exposure: Optional[np.ndarray] = None):
269        """
270        Return a list of scipy.stats frozen distribution objects, one per row.
271
272        Parameters
273        ----------
274        X : DataFrame or array
275        exposure : array-like, optional
276
277        Returns
278        -------
279        list of scipy.stats frozen distributions (or None if unsupported)
280        """
281        self._check_fitted()
282        from scipy import stats
283
284        param_arrays = {}
285        for pname in self.family.param_names:
286            param_arrays[pname] = self.predict(X, parameter=pname, exposure=exposure)
287
288        n = len(param_arrays[self.family.param_names[0]])
289        family_name = self.family.__class__.__name__
290        dists = []
291
292        for i in range(n):
293            p = {k: v[i] for k, v in param_arrays.items()}
294            dist = _make_scipy_dist(family_name, p, self.family)
295            dists.append(dist)
296
297        return dists
298
299    def predict_mean(self, X, exposure: Optional[np.ndarray] = None) -> np.ndarray:
300        """E[Y|X] — convenience wrapper.
301
302        For ZIP, E[Y] = (1-pi)*mu, not mu alone.
303        """
304        mu = self.predict(X, parameter="mu", exposure=exposure)
305        if self.family.__class__.__name__ == "ZIP":
306            pi = self.predict(X, parameter="pi")
307            return (1.0 - pi) * mu
308        return mu
309
310    def predict_variance(self, X, exposure: Optional[np.ndarray] = None) -> np.ndarray:
311        """
312        Var[Y|X] from the fitted distribution parameters.
313
314        Uses the distribution's variance formula rather than empirical variance.
315        """
316        self._check_fitted()
317        family_name = self.family.__class__.__name__
318        param_arrays = {
319            pname: self.predict(X, parameter=pname, exposure=exposure)
320            for pname in self.family.param_names
321        }
322        return _compute_variance(family_name, param_arrays, self.family)
323
324    def volatility_score(self, X, exposure: Optional[np.ndarray] = None) -> np.ndarray:
325        """
326        Coefficient of variation CV = sqrt(Var[Y|X]) / E[Y|X] per risk.
327
328        Insurance interpretation: risks with CV > 1 have larger standard
329        deviation than mean — typical for heavy-tailed severity.
330        Useful for pricing loading decisions and risk selection.
331        """
332        mu = self.predict_mean(X, exposure=exposure)
333        var = self.predict_variance(X, exposure=exposure)
334        return np.sqrt(np.clip(var, 0, None)) / np.clip(mu, 1e-10, None)
335
336    def relativities(self, parameter: str = "mu"):
337        """
338        Return coefficient relativities for a given parameter.
339
340        For log-linked parameters (mu, sigma on most families):
341          - Intercept relativity = exp(beta_intercept) — the fitted baseline
342            on the response scale.
343          - All other term relativities = exp(beta_j) — the multiplicative
344            effect of a one-unit increase in that term on the response scale.
345
346        For logit-linked parameters (pi in ZIP):
347          Returns odds ratios: exp(beta_j).
348
349        Returns
350        -------
351        polars DataFrame with columns: param, term, coefficient, relativity, link
352        """
353        self._check_fitted()
354        if not _HAS_POLARS:
355            raise ImportError("polars required for relativities(). Install polars.")
356
357        betas = self._betas[parameter]
358        cols = self._formulas.get(parameter, [])
359        link = self.family.default_links[parameter]
360
361        terms = ["(Intercept)"] + cols
362        intercept = betas[0]
363
364        rows = []
365        for term, beta in zip(terms, betas):
366            if link.name == "log":
367                rel = float(np.exp(beta))
368                if term == "(Intercept)":
369                    base = float(np.exp(intercept))
370                    rows.append({"param": parameter, "term": term, "coefficient": float(beta),
371                                 "relativity": base, "link": "log"})
372                else:
373                    rows.append({"param": parameter, "term": term, "coefficient": float(beta),
374                                 "relativity": rel, "link": "log"})
375            elif link.name == "logit":
376                rows.append({"param": parameter, "term": term, "coefficient": float(beta),
377                             "relativity": float(np.exp(beta)), "link": "logit (odds ratio)"})
378            else:
379                rows.append({"param": parameter, "term": term, "coefficient": float(beta),
380                             "relativity": float(beta), "link": link.name})
381
382        return pl.DataFrame(rows)
383
384    def summary(self) -> None:
385        """Print coefficient summary for all parameters."""
386        self._check_fitted()
387        print(f"\nDistributionalGLM — {self.family.__class__.__name__}")
388        print(f"  n = {self._n}, loglik = {self._loglik:.4f}")
389        print(f"  Converged: {self._converged}")
390        print(f"  GAIC(2): {self.gaic(penalty=2):.4f}")
391        print()
392
393        for pname in self.family.param_names:
394            betas = self._betas[pname]
395            cols = self._formulas.get(pname, [])
396            terms = ["(Intercept)"] + cols
397            link = self.family.default_links[pname]
398            print(f"  Parameter: {pname}  (link: {link.name})")
399            print(f"  {'Term':<30} {'Coef':>12}")
400            print(f"  {'-'*44}")
401            for term, beta in zip(terms, betas):
402                print(f"  {term:<30} {beta:>12.5f}")
403            print()
404
405    def gaic(self, penalty: float = 2.0) -> float:
406        """
407        Generalised AIC with given penalty.
408
409        penalty=2   -> standard AIC
410        penalty=log(n)  -> BIC
411        """
412        self._check_fitted()
413        total_params = sum(len(b) for b in self._betas.values())
414        return _gaic(self._loglik, total_params, penalty=penalty)
415
416    def score(
417        self,
418        X,
419        y,
420        metric: str = "nll",
421        exposure: Optional[np.ndarray] = None,
422    ) -> float:
423        """
424        Evaluate model on held-out data.
425
426        Parameters
427        ----------
428        X : DataFrame or array
429        y : array-like
430        metric : str
431            'nll' (negative log-likelihood), 'deviance', or 'crps'.
432        exposure : array-like, optional
433
434        Returns
435        -------
436        float
437            Metric value. Lower is better for nll and deviance.
438        """
439        self._check_fitted()
440        y = np.asarray(y, dtype=float)
441
442        param_arrays = {
443            pname: self.predict(X, parameter=pname, exposure=exposure)
444            for pname in self.family.param_names
445        }
446        ll = self.family.log_likelihood(y, param_arrays)
447
448        if metric == "nll":
449            return float(-np.mean(ll[np.isfinite(ll)]))
450        elif metric == "deviance":
451            # Saturated model: each obs gets its own parameter
452            ll_sat = self._saturated_loglik(y, param_arrays)
453            return float(2.0 * np.sum(ll_sat - ll))
454        elif metric == "crps":
455            dists = self.predict_distribution(X, exposure=exposure)
456            crps_vals = []
457            for dist, yi in zip(dists, y):
458                if dist is not None:
459                    # Monte Carlo CRPS approximation using the standard estimator:
460                    #   CRPS = E|X - y| - 0.5 * E|X - X'|
461                    # where X, X' are independent draws from the predicted distribution.
462                    # We draw n_samples and use a random split for the second term —
463                    # this avoids the O(n^2) full pairwise calculation and does not
464                    # depend on drawing exactly 1000 samples.
465                    n_samples = 500
466                    s1 = dist.rvs(n_samples)
467                    s2 = dist.rvs(n_samples)
468                    crps_vals.append(
469                        np.mean(np.abs(s1 - yi)) - 0.5 * np.mean(np.abs(s1 - s2))
470                    )
471            return float(np.mean(crps_vals))
472        else:
473            raise ValueError(f"Unknown metric: {metric}. Use 'nll', 'deviance', or 'crps'.")
474
475    def _saturated_loglik(self, y, params):
476        """Saturated log-likelihood (each obs perfectly fitted to itself)."""
477        sat_params = {k: v.copy() for k, v in params.items()}
478        sat_params["mu"] = y.copy()
479        try:
480            ll = self.family.log_likelihood(y, sat_params)
481            return np.where(np.isfinite(ll), ll, 0.0)
482        except Exception:
483            return np.zeros(len(y))
484
485    def _check_fitted(self):
486        if not self._fitted:
487            raise RuntimeError("Model not fitted. Call fit() first.")
488
489    @property
490    def coefficients(self) -> Dict[str, np.ndarray]:
491        """Fitted coefficients for all parameters."""
492        self._check_fitted()
493        return {k: v.copy() for k, v in self._betas.items()}
494
495    @property
496    def n_observations(self) -> int:
497        return self._n
498
499    @property
500    def loglikelihood(self) -> float:
501        self._check_fitted()
502        return self._loglik
503
504    @property
505    def converged(self) -> bool:
506        self._check_fitted()
507        return self._converged
508
509    def __repr__(self) -> str:
510        status = "fitted" if self._fitted else "unfitted"
511        return (
512            f"DistributionalGLM({self.family.__class__.__name__}, "
513            f"n={self._n}, {status})"
514        )

 58def choose_distribution(
 59    X,
 60    y: np.ndarray,
 61    families: List[GamlssFamily],
 62    formulas: Optional[Dict] = None,
 63    exposure: Optional[np.ndarray] = None,
 64    weights: Optional[np.ndarray] = None,
 65    penalty: float = 2.0,
 66    max_iter: int = 100,
 67    tol: float = 1e-6,
 68    verbose: bool = False,
 69) -> List[SelectionResult]:
 70    """
 71    Fit multiple families and rank by GAIC.
 72
 73    Parameters
 74    ----------
 75    X : DataFrame (polars or pandas)
 76        Covariate matrix.
 77    y : ndarray
 78        Response observations.
 79    families : list of GamlssFamily
 80        Candidate distribution families to compare.
 81    formulas : dict, optional
 82        Passed to DistributionalGLM. If None, all columns for mu,
 83        intercept-only for other parameters.
 84    exposure : ndarray, optional
 85        Exposure offset (values, not log). Applied to mu.
 86    weights : ndarray, optional
 87        Observation weights.
 88    penalty : float
 89        GAIC penalty. Default 2 (AIC).
 90    max_iter : int
 91        Maximum RS iterations per family.
 92    tol : float
 93        Convergence tolerance.
 94    verbose : bool
 95        Print progress.
 96
 97    Returns
 98    -------
 99    list of SelectionResult
100        Sorted ascending by GAIC (best first).
101    """
102    # Import here to avoid circular dependency
103    from .model import DistributionalGLM
104
105    results = []
106    for fam in families:
107        fname = fam.__class__.__name__
108        if verbose:
109            print(f"Fitting {fname}...")
110        try:
111            mdl = DistributionalGLM(family=fam, formulas=formulas)
112            mdl.fit(
113                X, y,
114                exposure=exposure,
115                weights=weights,
116                max_iter=max_iter,
117                tol=tol,
118            )
119            ll = mdl._loglik
120            df = sum(len(b) for b in mdl._betas.values())
121            g = gaic(ll, df, penalty=penalty)
122            results.append(SelectionResult(
123                family=fam,
124                family_name=fname,
125                gaic=g,
126                loglik=ll,
127                df=df,
128                converged=mdl._converged,
129                model=mdl,
130            ))
131        except (ValueError, ArithmeticError, np.linalg.LinAlgError) as e:
132            if verbose:
133                print(f"  {fname} failed: {e}")
134
135    results.sort(key=lambda r: r.gaic)
136    return results

36def gaic(loglik: float, n_params: int, penalty: float = 2.0) -> float:
37    """
38    Generalised AIC.
39
40    Parameters
41    ----------
42    loglik : float
43        Total log-likelihood at convergence.
44    n_params : int
45        Total number of estimated parameters (sum of df across all
46        distribution components).
47    penalty : float
48        Penalty per parameter. Use 2 for AIC, log(n) for BIC.
49
50    Returns
51    -------
52    float
53        GAIC value. Lower is better.
54    """
55    return -2.0 * loglik + penalty * n_params