"""Method demonstration on SYNTHETIC data: choice-based conjoint for a fictional income-protection policy.
Simulates respondents from known preferences, then recovers them with (a) a pooled multinomial logit and
(b) a two-class latent-class logit fitted by EM. Nothing here describes a real product, client or market.
Run: python3 analysis/conjoint.py"""
import json, itertools
import numpy as np
from common import ROOT
rng = np.random.default_rng(20261114)
OUT = ROOT.parent / "src" / "static" / "data"

ATTRS = [  # first level of each attribute is the reference
    ("cover", "What is covered", ["Unemployment only", "Unemployment + sick leave", "Unemployment, sick leave + hospital stay"]),
    ("months", "Benefit paid for up to", ["6 months", "12 months", "24 months"]),
    ("wait", "Waiting period", ["90 days", "60 days", "30 days"]),
    ("seller", "Sold by", ["Your bank", "An insurer", "An online-only provider"]),
]
PRICES = [6, 9, 12, 15]                      # EUR per month, entered as a linear term
NAMES = [f"{a[0]}:{l}" for a in ATTRS for l in a[2][1:]] + ["price", "none"]
K = len(NAMES)
# True preferences of two simulated segments (part-worths vs reference level, price per EUR, utility of "none")
TRUE = {
  "Price-led":      np.array([0.35, 0.55,  0.20, 0.30,  0.15, 0.35,  0.10, -0.30,  -0.34, -2.2]),
  "Protection-led": np.array([0.90, 1.70,  0.55, 1.05,  0.25, 0.60, -0.15, -0.60,  -0.11,  0.2]),
}
SHARE = {"Price-led": 0.58, "Protection-led": 0.42}
N, T, J = 600, 10, 3                          # respondents, tasks, product alternatives (+ "none")

def encode(levels, price):
    x = np.zeros(K); i = 0
    for a, lv in zip(ATTRS, levels):
        if lv > 0: x[i + lv - 1] = 1
        i += len(a[2]) - 1
    x[K - 2] = price
    return x
NONE = np.zeros(K); NONE[K - 1] = 1

def simulate():
    X = np.zeros((N, T, J + 1, K)); Y = np.zeros((N, T), dtype=int); seg = []
    segs = list(TRUE)
    for n in range(N):
        s = segs[0] if rng.random() < SHARE[segs[0]] else segs[1]; seg.append(s)
        beta = TRUE[s] + rng.normal(0, 0.25, K) * (np.arange(K) < K - 2)   # individual taste noise on part-worths
        for t in range(T):
            for j in range(J):
                X[n, t, j] = encode([rng.integers(len(a[2])) for a in ATTRS], PRICES[rng.integers(4)])
            X[n, t, J] = NONE
            u = X[n, t] @ beta + rng.gumbel(size=J + 1)
            Y[n, t] = int(np.argmax(u))
    return X, Y, seg

def mnl(X, Y, w=None, beta0=None, iters=40):
    """Weighted MNL by Newton-Raphson. X: (obs, alts, K); Y: (obs,); w: (obs,)."""
    O = X.shape[0]; w = np.ones(O) if w is None else w
    b = np.zeros(K) if beta0 is None else beta0.copy()
    for _ in range(iters):
        u = X @ b; u -= u.max(1, keepdims=True); p = np.exp(u); p /= p.sum(1, keepdims=True)
        xbar = (p[:, :, None] * X).sum(1)
        g = (w[:, None] * (X[np.arange(O), Y] - xbar)).sum(0)
        H = -(np.einsum("o,oj,ojk,ojl->kl", w, p, X, X) - np.einsum("o,ok,ol->kl", w, xbar, xbar))
        step = np.linalg.solve(H, g); b = b - step
        if np.abs(step).max() < 1e-8: break
    u = X @ b; u -= u.max(1, keepdims=True); lp = u - np.log(np.exp(u).sum(1, keepdims=True))
    ll = (w * lp[np.arange(O), Y]).sum()
    return b, np.linalg.inv(-H), ll, lp

def latent_class(X, Y, C=2, starts=8):
    Xf = X.reshape(N * T, J + 1, K); Yf = Y.reshape(N * T); best = None
    for s in range(starts):
        r = rng.dirichlet(np.ones(C) * 0.6, size=N); ll_old = -1e18
        for it in range(200):
            fits = [mnl(Xf, Yf, w=np.repeat(r[:, c], T)) for c in range(C)]
            pi = r.mean(0)
            lpn = np.stack([f[3][np.arange(N * T), Yf].reshape(N, T).sum(1) for f in fits], 1) + np.log(pi)
            m = lpn.max(1, keepdims=True); ll = (m[:, 0] + np.log(np.exp(lpn - m).sum(1))).sum()
            r = np.exp(lpn - m); r /= r.sum(1, keepdims=True)
            if ll - ll_old < 1e-6: break
            ll_old = ll
        if best is None or ll > best[0]: best = (ll, [f[0] for f in fits], [f[1] for f in fits], pi, r)
    return best

def wtp(b, V, i):
    """Willingness to pay for level i (EUR/month) = -b_i / b_price, delta-method standard error."""
    p = K - 2; est = -b[i] / b[p]
    grad = np.zeros(K); grad[i] = -1 / b[p]; grad[p] = b[i] / b[p] ** 2
    return est, float(np.sqrt(grad @ V @ grad))

if __name__ == "__main__":
    X, Y, seg = simulate()
    Xf, Yf = X.reshape(N * T, J + 1, K), Y.reshape(N * T)
    b, V, ll, _ = mnl(Xf, Yf); se = np.sqrt(np.diag(V))
    ll0 = N * T * np.log(1 / (J + 1))
    llc, bc, Vc, pi, resp = latent_class(X, Y)
    order = np.argsort([x[K - 2] for x in bc])            # most price-sensitive class first
    labels = ["Price-led", "Protection-led"]
    assigned = np.array(labels)[np.argsort(order)[resp.argmax(1)]]
    acc = float((assigned == np.array(seg)).mean())
    npar1, npar2 = K, 2 * K + 1
    bic = lambda l, k: -2 * l + k * np.log(N * T)
    out = {
      "synthetic": True, "seed": 20261114, "design": {"respondents": N, "tasks": T, "alternatives": J, "choices": N * T},
      "attrs": [{"id": a[0], "label": a[1], "levels": a[2]} for a in ATTRS], "prices": PRICES, "names": NAMES,
      "pooled": {"beta": b.round(4).tolist(), "se": se.round(4).tolist(), "ll": round(ll, 1), "ll0": round(ll0, 1),
                 "rho2": round(1 - ll / ll0, 3), "bic": round(bic(ll, npar1), 1),
                 "wtp": [[round(v, 2) for v in wtp(b, V, i)] for i in range(K - 2)]},
      "classes": [{"label": labels[r_], "size": round(float(pi[c]), 3), "beta": bc[c].round(4).tolist(),
                   "se": np.sqrt(np.diag(Vc[c])).round(4).tolist(), "truth": TRUE[labels[r_]].tolist(), "true_size": SHARE[labels[r_]],
                   "wtp": [[round(v, 2) for v in wtp(bc[c], Vc[c], i)] for i in range(K - 2)]}
                  for r_, c in enumerate(order)],
      "lc": {"ll": round(llc, 1), "bic": round(bic(llc, npar2), 1), "rho2": round(1 - llc / ll0, 3), "assign_accuracy": round(acc, 3)},
      "none_share": round(float((Yf == J).mean()), 3),
    }
    (OUT / "conjoint.json").write_text(json.dumps(out, separators=(",", ":")), encoding="utf-8")
    print(json.dumps(out, indent=None)[:3000])