"""Distributed estimators (in-memory, statistically faithful to the coordinator/worker protocol).

All routines take a batched dataset `db` = {'X': (K, m, p), 'y': (K, m)} (or {'x': (K, m)}) so that block k is
`db[...][k]`.  Worker-side quantities are computed block by block; coordinator-side quantities are block
averages.  Nothing here measures communication; see `communication.py` (accounting) and `protocol.py` (real runs).
"""
from __future__ import annotations
import numpy as np
from .jackknife import vandermonde_weights, DEFAULT_SPLITS
from .likelihoods import STATUS_OK


# ----------------------------------------------------------------------------- local fits
def local_split_fits(model, db, splits=(1,), **fit_kw):
    """Fit the MLE on every sub-block of every block for each split multiplier s in `splits`.

    Returns dict s -> (theta (K, s, p), status (K, s), hessian (K, s, p, p)).
    """
    K = next(iter(db.values())).shape[0]
    m = model.n_obs(db)
    out = {}
    for s in splits:
        if m % s:
            raise ValueError(f"m={m} not divisible by s={s}")
        sub = {k: v.reshape((K * s, m // s) + v.shape[2:]) for k, v in db.items()}
        th, info = model.fit_batched(sub, **fit_kw)
        out[s] = (th.reshape(K, s, -1), info["status"].reshape(K, s), info["hessian"].reshape(K, s, th.shape[1], th.shape[1]))
    return out


def worker_corrected_estimates(fits, splits):
    """J_{r,k} = sum_j w_j * mean_l theta_hat(sub-block l of level j)  -> (K, p); also validity mask (K,)."""
    w = vandermonde_weights(splits)
    J = 0.0; ok = None
    for wj, s in zip(w, splits):
        th, st, _ = fits[s]
        J = J + wj * th.mean(axis=1)
        good = np.all(st == STATUS_OK, axis=1)
        ok = good if ok is None else (ok & good)
    return J, ok


def pilot_estimator(J, block_sizes=None, ok=None):
    """Coordinator average of the worker estimates (weights m_k/N if sizes given). Invalid workers can be dropped."""
    J = np.asarray(J)
    if block_sizes is None:
        wts = np.ones(J.shape[0])
    else:
        wts = np.asarray(block_sizes, float)
    if ok is not None:
        wts = wts * ok
    wts = wts / wts.sum()
    return np.einsum('k,kp->p', wts, J)


# ----------------------------------------------------------------------------- aggregated derivatives
def aggregated_score_hessian(model, db, theta, block_sizes=None):
    """Workers evaluate (score, Hessian) of their block at the common theta; coordinator averages (weighted by m_k)."""
    K = next(iter(db.values())).shape[0]
    th = np.broadcast_to(np.asarray(theta, float), (K, len(theta)))
    g, H = model.score_hessian(th, db)
    if block_sizes is None:
        return g.mean(axis=0), H.mean(axis=0), g, H
    w = np.asarray(block_sizes, float); w = w / w.sum()
    return np.einsum('k,kp->p', w, g), np.einsum('k,kpq->pq', w, H), g, H


def newton_step(theta, g, H, damping=None):
    """theta - H^{-1} g via a linear solve (never an explicit inverse). Returns (new theta, diagnostics)."""
    try:
        step = np.linalg.solve(-H, g)
        ok = np.all(np.isfinite(step))
    except np.linalg.LinAlgError:
        step = np.zeros_like(g); ok = False
    diag = {"solve_ok": ok, "cond": float(np.linalg.cond(-H)) if ok else np.inf, "step_norm": float(np.linalg.norm(step))}
    if damping is not None and ok:
        step = step * damping
    return theta + step, diag


def global_newton(model, db, theta0, q=1, hessian_mode="exact", block_sizes=None, stale_hessian=None,
                  quantize_bits=None, single_machine=0, rng=None):
    """q globally aggregated Newton corrections from theta0.

    hessian_mode:
      'exact'   workers send score and Hessian at the broadcast iterate (p + p(p+1)/2 numbers each round);
      'stale'   Hessian fixed = `stale_hessian` (e.g. average of local Hessians at local MLEs, sent once with the pilot);
                workers send only scores each round (p numbers);
      'single'  Hessian of worker `single_machine` evaluated at the iterate (no Hessian uplink; CSL-type curvature);
      'quantized' exact aggregated Hessian rounded to `quantize_bits` significant bits per entry (uniform relative
                  quantization), scores at full precision.
    Returns list of iterates [theta0, theta1, ..., thetaq] and per-step diagnostics.
    """
    th = np.asarray(theta0, float).copy(); path = [th.copy()]; diags = []
    for _ in range(q):
        g, H, g_k, H_k = aggregated_score_hessian(model, db, th, block_sizes)
        if hessian_mode == "stale":
            H = np.asarray(stale_hessian)
        elif hessian_mode == "single":
            H = H_k[single_machine]
        elif hessian_mode == "quantized":
            H = quantize_relative(H, quantize_bits)
        elif hessian_mode != "exact":
            raise ValueError(hessian_mode)
        th, d = newton_step(th, g, H)
        d["grad_norm"] = float(np.linalg.norm(g))
        path.append(th.copy()); diags.append(d)
    return path, diags


def quantize_relative(A, bits):
    """Round each entry to `bits` significant binary digits (mantissa truncation), keeping the sign and exponent."""
    if bits is None:
        return A
    A = np.asarray(A, float)
    mant, exp = np.frexp(A)
    scale = 2.0 ** bits
    return np.ldexp(np.round(mant * scale) / scale, exp)


# ----------------------------------------------------------------------------- baselines
def aee_lin_xi(theta_local, hess_local):
    """Lin & Xi (2011) aggregated estimating equation: (sum_k A_k)^{-1} sum_k A_k theta_k with A_k = -H_k(theta_k)."""
    A = -np.asarray(hess_local)                       # (K,p,p)
    rhs = np.einsum('kpq,kq->p', A, theta_local)
    return np.linalg.solve(A.sum(axis=0), rhs)


def csl(model, db, theta_init, rounds=1, machine=0, block_sizes=None, fit_kw=None):
    """Jordan, Lee & Yang (2019) communication-efficient surrogate likelihood, maximized on `machine`.

    Round t: workers send gradients at theta_t (p numbers each); machine 0 maximizes
        L_tilde(theta) = L_0(theta) - <grad L_0(theta_t) - grad L_N(theta_t), theta>
    by Newton's method on its own block (local Hessian only).  Returns iterates.
    """
    fit_kw = fit_kw or {}
    th = np.asarray(theta_init, float).copy(); path = [th.copy()]
    d0 = {k: v[machine][None] for k, v in db.items()}
    for _ in range(rounds):
        gN, _, g_k, _ = aggregated_score_hessian(model, db, th, block_sizes)
        shift = g_k[machine] - gN                      # p-vector
        # maximize L_0(theta) - shift' theta  with Newton on machine 0
        x = th.copy()
        for it in range(100):
            g, H = model.score_hessian(x[None], d0)
            g = g[0] - shift; H = H[0]
            step = np.linalg.solve(-H, g)
            # backtracking on the surrogate objective
            f = lambda z: model.loglik(z[None], d0)[0] - shift @ z
            f0 = f(x); t = 1.0
            while f(x + t * step) < f0 - 1e-12 and t > 1e-8:
                t *= 0.5
            x = x + t * step
            if np.max(np.abs(g)) < 1e-10:
                break
        th = x; path.append(th.copy())
    return path


def analytic_bias_correction(model, db, theta_local):
    """Cox–Snell / Cordeiro–McCullagh O(1/m) bias correction of each local MLE (GLM canonical link, or Exponential):
    b_hat = (1/2) I^{-1} [n^{-1} sum_i nabla^2 psi_i (I^{-1})]  per block;  theta_corrected = theta - b_hat / m.
    (Derivation: theta_hat - theta0 = -H^{-1} Psi - H^{-1}(V delta) - (1/2) H^{-1} nabla^2 Psi[delta,delta] + ...; with H = -I,
    E[delta delta^T] = I^{-1}/n and the V-term vanishing conditionally on the design for canonical GLMs.)
    Equivalent to the general second-order bias of M-estimators with the first (score-covariance) term dropped,
    which is exact for canonical GLMs conditional on the design.
    """
    K, m = next(iter(db.values())).shape[:2]
    g, H = model.score_hessian(theta_local, db)
    I = -H
    Iinv = np.linalg.inv(I)                           # p small; explicit inverse acceptable for a diagnostic quantity
    T = model.third_contraction(theta_local, db, Iinv)   # (K,p) = n^{-1} sum nabla^2 psi_i [I^{-1}]
    b = 0.5 * np.einsum("kpq,kq->kp", Iinv, T)        # per-block O(1) bias coefficient b_1 = (1/2) I^{-1} T
    return theta_local - b / m, b


# ----------------------------------------------------------------------------- one call = all estimators on one sample
def estimate_all(model, db, splits_max=3, q_max=2, block_sizes=None, fit_kw=None, extras=True):
    """Compute, on the same partitioned sample, the full family of estimators used in the Monte Carlo.

    Returns dict name -> theta (p,), plus a dict of diagnostics (validity masks, Newton diagnostics, local fit stats).
    Names:  raw (r=1,q=0) ... 'r{r}_q{q}' for r=1..splits_max, q=0..q_max  with splits (1,2,4,..);
            'central' centralized MLE; 'aee' Lin–Xi; 'csl1','csl2' CSL rounds from the raw average;
            'cs_q{q}' analytic Cox–Snell corrected average + q Newton; 'stale_r{r}_q{q}' stale-Hessian variants.
    """
    fit_kw = fit_kw or {}
    splits = DEFAULT_SPLITS[splits_max]
    fits = local_split_fits(model, db, splits, **fit_kw)
    K, m = next(iter(db.values())).shape[:2]
    out, diag = {}, {}
    diag["local_status_counts"] = {s: np.bincount(fits[s][1].ravel(), minlength=4).tolist() for s in splits}
    # centralized MLE on the pooled data
    full = {k: v.reshape((1, K * m) + v.shape[2:]) for k, v in db.items()}
    th_c, info_c = model.fit_batched(full, tol=1e-12, maxit=100)
    out["central"] = th_c[0]; diag["central_status"] = int(info_c["status"][0])
    theta_full_local, status_full, H_local = fits[1]
    H_bar = H_local[:, 0].mean(axis=0)                 # average local Hessian at local MLEs (sent with the pilot)
    for r in range(1, splits_max + 1):
        sp = splits[:r]
        J, ok = worker_corrected_estimates({s: fits[s] for s in sp}, sp)
        diag[f"valid_r{r}"] = ok
        th0 = pilot_estimator(J, block_sizes, ok=None)   # keep all workers (failures tracked separately)
        path, nd = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)
        for q in range(q_max + 1):
            out[f"r{r}_q{q}"] = path[q]
        diag[f"newton_r{r}"] = nd
        if extras:
            path_s, _ = global_newton(model, db, th0, q=q_max, hessian_mode="stale", stale_hessian=H_bar, block_sizes=block_sizes)
            for q in range(1, q_max + 1):
                out[f"stale_r{r}_q{q}"] = path_s[q]
    if extras:
        out["aee"] = aee_lin_xi(theta_full_local[:, 0], H_local[:, 0])
        raw = out["r1_q0"]
        cpath = csl(model, db, raw, rounds=2, block_sizes=block_sizes)
        out["csl1"], out["csl2"] = cpath[1], cpath[2]
        try:
            th_cs, b = analytic_bias_correction(model, db, theta_full_local[:, 0])
            th0 = pilot_estimator(th_cs, block_sizes)
            pth, _ = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)
            for q in range(q_max + 1):
                out[f"cs_q{q}"] = pth[q]
        except NotImplementedError:
            pass
    return out, diag