"""Primitives of the WP19 search model: match-value distributions, seller-type grids, parameters.

All quantities are in units of the mean buyer valuation of the quality segment (q = 1), time in years.
"""
from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional

import numpy as np
import yaml
from scipy import integrate, special, stats

SQRT2PI = np.sqrt(2.0 * np.pi)


# ----------------------------------------------------------------------------------------------
# Match-value distributions G (valuation w of a buyer for a given house, PV units)
# ----------------------------------------------------------------------------------------------
class MatchDist:
    """Interface: cdf, pdf, H(v)=E[(w-v)^+], M(v)=E[w-v|w>=v], cond_mean, cond_var, mean, support."""

    lo: float
    hi: float
    name: str = "base"

    def cdf(self, v):
        raise NotImplementedError

    def pdf(self, v):
        raise NotImplementedError

    def H(self, v):
        raise NotImplementedError

    def surv(self, v):
        return 1.0 - self.cdf(v)

    def hazard(self, v):
        s = self.surv(v)
        return np.where(s > 0, self.pdf(v) / np.maximum(s, 1e-300), np.inf)

    def M(self, v):
        s = self.surv(v)
        return np.where(s > 0, self.H(v) / np.maximum(s, 1e-300), 0.0)

    def cond_mean(self, v):
        return v + self.M(v)

    def cond_var(self, v):
        """Var(w | w >= v) by quadrature (generic)."""
        v = float(v)
        s = float(self.surv(v))
        if s <= 0:
            return 0.0
        hi = self.hi if np.isfinite(self.hi) else v + 50.0 * max(1.0, self.M(v))
        m2 = integrate.quad(lambda w: w * w * float(self.pdf(w)), v, hi, limit=200)[0] / s
        m1 = float(self.cond_mean(v))
        return max(m2 - m1 * m1, 0.0)

    def mean(self):
        return float(self.H(self.lo) + self.lo) if np.isfinite(self.lo) else float(self.H(-1e6) - 1e6)

    def sample(self, rng: np.random.Generator, n: int):
        raise NotImplementedError

    def is_ifr(self) -> Optional[bool]:
        return None


@dataclass
class TruncNormal(MatchDist):
    mu: float = 1.0
    sigma: float = 0.10
    lo: float = 0.70
    hi: float = 1.30
    name: str = "truncnorm"

    def __post_init__(self):
        self.a = (self.lo - self.mu) / self.sigma
        self.b = (self.hi - self.mu) / self.sigma
        self.Z = stats.norm.cdf(self.b) - stats.norm.cdf(self.a)

    def _z(self, v):
        return (np.asarray(v, dtype=float) - self.mu) / self.sigma

    def cdf(self, v):
        z = np.clip(self._z(v), self.a, self.b)
        return (stats.norm.cdf(z) - stats.norm.cdf(self.a)) / self.Z

    def pdf(self, v):
        z = self._z(v)
        inside = (z >= self.a) & (z <= self.b)
        return np.where(inside, stats.norm.pdf(z) / (self.sigma * self.Z), 0.0)

    def H(self, v):
        """E[(w-v)^+] closed form for the truncated normal."""
        v = np.asarray(v, dtype=float)
        zv = np.clip(self._z(v), self.a, self.b)
        term = (self.sigma * (stats.norm.pdf(zv) - stats.norm.pdf(self.b))
                + (self.mu - np.maximum(v, self.lo)) * (stats.norm.cdf(self.b) - stats.norm.cdf(zv))) / self.Z
        # below the support: E[w] - v
        below = v < self.lo
        mean = self.mu + self.sigma * (stats.norm.pdf(self.a) - stats.norm.pdf(self.b)) / self.Z
        return np.where(below, mean - v, np.maximum(term, 0.0))

    def mean(self):
        return float(self.mu + self.sigma * (stats.norm.pdf(self.a) - stats.norm.pdf(self.b)) / self.Z)

    def sample(self, rng, n):
        return stats.truncnorm(self.a, self.b, loc=self.mu, scale=self.sigma).rvs(size=n, random_state=rng)

    def is_ifr(self):
        return True  # log-concave density


@dataclass
class Exponential(MatchDist):
    rate: float = 10.0
    lo: float = 0.70
    hi: float = np.inf
    name: str = "exponential"

    def cdf(self, v):
        v = np.asarray(v, dtype=float)
        return np.where(v < self.lo, 0.0, 1.0 - np.exp(-self.rate * (v - self.lo)))

    def pdf(self, v):
        v = np.asarray(v, dtype=float)
        return np.where(v < self.lo, 0.0, self.rate * np.exp(-self.rate * (v - self.lo)))

    def H(self, v):
        v = np.asarray(v, dtype=float)
        return np.where(v < self.lo, self.lo + 1.0 / self.rate - v, np.exp(-self.rate * (v - self.lo)) / self.rate)

    def mean(self):
        return self.lo + 1.0 / self.rate

    def sample(self, rng, n):
        return self.lo + rng.exponential(1.0 / self.rate, size=n)

    def is_ifr(self):
        return True  # constant failure rate (boundary case)


@dataclass
class Pareto(MatchDist):
    a: float = 3.0
    lo: float = 0.70
    hi: float = np.inf
    name: str = "pareto"

    def cdf(self, v):
        v = np.asarray(v, dtype=float)
        return np.where(v < self.lo, 0.0, 1.0 - (self.lo / np.maximum(v, self.lo)) ** self.a)

    def pdf(self, v):
        v = np.asarray(v, dtype=float)
        return np.where(v < self.lo, 0.0, self.a * self.lo ** self.a / np.maximum(v, self.lo) ** (self.a + 1))

    def H(self, v):
        v = np.asarray(v, dtype=float)
        mean = self.lo * self.a / (self.a - 1.0)
        above = (self.lo / np.maximum(v, self.lo)) ** self.a * np.maximum(v, self.lo) / (self.a - 1.0)
        return np.where(v < self.lo, mean - v, above)

    def mean(self):
        return self.lo * self.a / (self.a - 1.0)

    def sample(self, rng, n):
        return self.lo * (1.0 - rng.random(n)) ** (-1.0 / self.a)

    def is_ifr(self):
        return False  # decreasing failure rate


@dataclass
class Weibull(MatchDist):
    """w = lo + scale * X, X ~ Weibull(shape). IFR iff shape >= 1."""
    shape: float = 2.0
    scale: float = 0.2
    lo: float = 0.70
    hi: float = np.inf
    name: str = "weibull"

    def _x(self, v):
        return np.maximum((np.asarray(v, dtype=float) - self.lo) / self.scale, 0.0)

    def cdf(self, v):
        return 1.0 - np.exp(-self._x(v) ** self.shape)

    def pdf(self, v):
        x = self._x(v)
        return np.where(np.asarray(v) >= self.lo,
                        (self.shape / self.scale) * x ** (self.shape - 1.0) * np.exp(-x ** self.shape), 0.0)

    def H(self, v):
        v = np.asarray(v, dtype=float)
        x = self._x(v)
        # E[(X - x)^+] for Weibull: scale * Gamma(1+1/k) * Q(1+1/k, x^k) - x * exp(-x^k), Q = upper regularized gamma
        k = self.shape
        ex = self.scale * (special.gamma(1.0 + 1.0 / k) * special.gammaincc(1.0 + 1.0 / k, x ** k) - x * np.exp(-x ** k))
        mean = self.lo + self.scale * special.gamma(1.0 + 1.0 / k)
        return np.where(v < self.lo, mean - v, np.maximum(ex, 0.0))

    def mean(self):
        return self.lo + self.scale * special.gamma(1.0 + 1.0 / self.shape)

    def sample(self, rng, n):
        return self.lo + self.scale * rng.weibull(self.shape, size=n)

    def is_ifr(self):
        return self.shape >= 1.0


def make_dist(cfg: dict) -> MatchDist:
    d = dict(cfg)
    kind = d.pop("dist")
    if kind == "truncnorm":
        return TruncNormal(**d)
    if kind == "exponential":
        return Exponential(**d)
    if kind == "pareto":
        return Pareto(**d)
    if kind == "weibull":
        return Weibull(**d)
    raise ValueError(f"unknown match-value distribution {kind}")


# ----------------------------------------------------------------------------------------------
# Seller-type grid (discrete approximation of F)
# ----------------------------------------------------------------------------------------------
def seller_type_grid(c_min: float, c_max: float, n_types: int, beta_a: float, beta_b: float):
    """Midpoint discretisation of Beta(a,b) rescaled to [c_min, c_max]: returns (c, f) with f summing to one."""
    if n_types == 1:
        mean = c_min + (c_max - c_min) * beta_a / (beta_a + beta_b)
        return np.array([mean]), np.array([1.0])
    edges = np.linspace(0.0, 1.0, n_types + 1)
    mass = np.diff(stats.beta.cdf(edges, beta_a, beta_b))
    # conditional means within each bin (exact for Beta via incomplete beta ratio)
    cond = np.diff(stats.beta.cdf(edges, beta_a + 1.0, beta_b)) * (beta_a / (beta_a + beta_b)) / mass
    c = c_min + (c_max - c_min) * cond
    return c, mass / mass.sum()


# ----------------------------------------------------------------------------------------------
# Parameters
# ----------------------------------------------------------------------------------------------
@dataclass
class Params:
    r: float = 0.04
    s: float = 1.0
    delta: float = 0.5
    w0: float = 0.8
    k: float = 0.02
    kappa_b: float = 0.05
    phi: float = 0.85
    A: float = 4.5
    alpha: float = 0.5
    beta: float = 0.5
    c: np.ndarray = field(default_factory=lambda: np.array([0.05]))
    f: np.ndarray = field(default_factory=lambda: np.array([1.0]))
    G: MatchDist = field(default_factory=TruncNormal)
    # financing extension
    pi: float = 0.3
    nu: float = 12.0
    phi_m: float = 0.85
    # numerics
    theta_bounds: tuple = (0.02, 60.0)
    root_xtol: float = 1e-12
    grid_check_points: int = 60
    # dynamics
    T: float = 8.0
    dt: float = 0.02
    damping: float = 0.3
    dyn_tol: float = 1e-7
    dyn_max_iter: int = 400
    name: str = "baseline"
    seed: int = 20261020
    raw: dict = field(default_factory=dict)

    # --- derived -----------------------------------------------------------------------------
    def lam_s(self, theta):
        return self.A * np.asarray(theta, dtype=float) ** self.alpha

    def lam_b(self, theta):
        return self.A * np.asarray(theta, dtype=float) ** (self.alpha - 1.0)

    def copy(self, **changes) -> "Params":
        d = {k: v for k, v in self.__dict__.items()}
        d.update(changes)
        return Params(**d)

    def with_types(self, c_min, c_max, n_types, beta_a, beta_b) -> "Params":
        c, f = seller_type_grid(c_min, c_max, n_types, beta_a, beta_b)
        return self.copy(c=c, f=f)

    def config_hash(self) -> str:
        payload = {k: (v.tolist() if isinstance(v, np.ndarray) else (asdict(v) if hasattr(v, "__dataclass_fields__") else v))
                   for k, v in self.__dict__.items() if k != "raw"}
        payload["G"] = {"name": self.G.name, **{k: float(v) for k, v in self.G.__dict__.items()
                                                  if isinstance(v, (int, float)) and k not in ("Z", "a", "b")}}
        return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode()).hexdigest()[:16]


def load_params(path: str | Path) -> Params:
    cfg = yaml.safe_load(Path(path).read_text())
    c, f = seller_type_grid(cfg["sellers"]["c_min"], cfg["sellers"]["c_max"], cfg["sellers"]["n_types"],
                            cfg["sellers"]["beta_a"], cfg["sellers"]["beta_b"])
    fin = cfg.get("financing_extension", {})
    num = cfg.get("numerics", {})
    dyn = cfg.get("dynamics", {})
    return Params(
        r=cfg["preferences"]["r"], s=cfg["sellers"]["s"], delta=cfg["sellers"]["delta"], w0=cfg["sellers"]["w0"],
        k=cfg["buyers"]["k"], kappa_b=cfg["buyers"]["kappa_b"], phi=cfg["buyers"]["phi"],
        A=cfg["matching"]["A"], alpha=cfg["matching"]["alpha"], beta=cfg["bargaining"]["beta"],
        c=c, f=f, G=make_dist(cfg["match_values"]),
        pi=fin.get("pi", 0.3), nu=fin.get("nu", 12.0), phi_m=fin.get("phi", cfg["buyers"]["phi"]),
        theta_bounds=tuple(num.get("theta_bounds", (0.02, 60.0))), root_xtol=float(num.get("root_xtol", 1e-12)),
        grid_check_points=int(num.get("grid_check_points", 60)),
        T=float(dyn.get("T", 8.0)), dt=float(dyn.get("dt", 0.02)), damping=float(dyn.get("damping", 0.3)),
        dyn_tol=float(dyn.get("tol", 1e-7)), dyn_max_iter=int(dyn.get("max_iter", 400)),
        name=cfg.get("name", "baseline"), seed=int(cfg.get("seed", 0)), raw=cfg,
    )