"""
Two-timescale probabilistic affect estimation and behaviour generation
for humanoid robots -- synthetic Monte Carlo study.

ALL DATA ARE SYNTHETIC. Ground-truth affect trajectories and sensor
observations are generated from the assumed models documented below.
No human-subject data are used. Results characterise the estimator
under stated assumptions; they are not measurements of human emotion.
"""
import json, numpy as np
from scipy.stats import chi2

DT = 0.1            # s, filter/sample period (10 Hz)            [ASSUMED]
T_END = 120.0       # s, trial length                              [ASSUMED]
N = int(T_END / DT)
THETA = 0.5         # 1/s, emotion->mood reversion rate (tau=2 s)  [ASSUMED]
SIG_W = 0.15        # emotion diffusion intensity                  [ASSUMED]
EPISODE_S = (10.0, 20.0)  # episode length range, s                [ASSUMED]
Q_MU = 0.02         # filter random-walk std of mood per step      [ASSUMED]

# Episode target "mood" points in PAD space (illustrative, not literature values) [ASSUMED]
PROTOS = np.array([
    [ 0.6,  0.4,  0.3],   # elated / engaged
    [ 0.5, -0.4,  0.2],   # content / relaxed
    [-0.5,  0.5, -0.3],   # anxious
    [-0.5,  0.5,  0.3],   # irritated
    [-0.5, -0.4, -0.3],   # dejected
    [ 0.0,  0.0,  0.0],   # neutral
])
PROTO_NAMES = ["engaged", "relaxed", "anxious", "irritated", "dejected", "neutral"]

# Modality observation models y = H x + v  (x = [V, A, D])          [ASSUMED]
MODS = {
    "physio": dict(H=np.array([[0.0, 0.9, 0.0]]), sig=np.array([0.30])),
    "face":   dict(H=np.array([[1.0, 0.0, 0.0], [0.0, 0.4, 0.0]]), sig=np.array([0.25, 0.35])),
    "voice":  dict(H=np.array([[0.0, 0.6, 0.0], [0.0, 0.0, 0.7]]), sig=np.array([0.35, 0.40])),
}
MOD_ORDER = ["physio", "face", "voice"]
BURST = 10.0        # mean dropout burst length, frames            [ASSUMED]
OUT_SIG = 1.5       # outlier noise std                            [ASSUMED]


def gen_truth(rng, matched=False):
    """Fast emotion x mean-reverts to slow mood mu. Truth: mu jumps between
    episode prototypes. matched=True instead draws mu as the filter's own
    random walk (used only for NEES verification)."""
    x = np.zeros((N, 3)); mu = np.zeros((N, 3))
    if matched:
        m = rng.normal(0, 0.3, 3)
        for k in range(N):
            m = m + rng.normal(0, Q_MU, 3); mu[k] = m
    else:
        k = 0
        while k < N:
            L = int(rng.uniform(*EPISODE_S) / DT)
            mu[k:k + L] = PROTOS[rng.integers(len(PROTOS))]; k += L
    xk = mu[0].copy()
    a = 1 - THETA * DT
    for k in range(N):
        xk = a * xk + (1 - a) * mu[k] + np.sqrt(DT) * SIG_W * rng.normal(size=3)
        x[k] = xk
    return x, mu


def gen_obs(rng, x, noise_scale=1.0, p_drop=0.2, p_out=0.0):
    """Returns dict modality -> (Y [N,m], avail [N] bool). Dropout follows a
    two-state Markov chain with stationary missing probability p_drop."""
    obs = {}
    for name in MOD_ORDER:
        H, sig = MODS[name]["H"], MODS[name]["sig"] * noise_scale
        Y = x @ H.T + rng.normal(size=(N, len(sig))) * sig
        if p_out > 0:
            o = rng.random(N) < p_out
            Y[o] += rng.normal(size=(o.sum(), len(sig))) * OUT_SIG
        avail = np.ones(N, bool)
        if p_drop > 0:
            q_end = 1.0 / BURST                       # leave-missing prob
            q_start = q_end * p_drop / (1 - p_drop)   # enter-missing prob
            miss = rng.random() < p_drop
            for k in range(N):
                miss = (rng.random() > q_end) if miss else (rng.random() < q_start)
                avail[k] = not miss
        obs[name] = (Y, avail)
    return obs


def build_model(noise_scale=1.0):
    a = 1 - THETA * DT
    F = np.block([[a * np.eye(3), (1 - a) * np.eye(3)], [np.zeros((3, 3)), np.eye(3)]])
    Q = np.diag([DT * SIG_W ** 2] * 3 + [Q_MU ** 2] * 3)
    R = {n: np.diag((MODS[n]["sig"] * noise_scale) ** 2) for n in MOD_ORDER}
    Hs = {n: np.hstack([MODS[n]["H"], np.zeros((MODS[n]["H"].shape[0], 3))]) for n in MOD_ORDER}
    return F, Q, R, Hs


def kalman(obs, mods=MOD_ORDER, gate=False, noise_scale=1.0, alpha=0.01, return_P=False):
    F, Q, R, Hs = build_model(noise_scale)
    s = np.zeros(6); P = np.diag([0.3] * 6)
    est = np.zeros((N, 6)); Ps = np.zeros((N, 6, 6)); rej = 0; tot = 0
    thr = {n: chi2.ppf(1 - alpha, Hs[n].shape[0]) for n in mods}
    for k in range(N):
        s = F @ s; P = F @ P @ F.T + Q
        for n in mods:                       # sequential per-modality update
            Y, av = obs[n]
            if not av[k]:
                continue
            H = Hs[n]; r = Y[k] - H @ s; S = H @ P @ H.T + R[n]
            tot += 1
            if gate and r @ np.linalg.solve(S, r) > thr[n]:
                rej += 1; continue
            K = P @ H.T @ np.linalg.inv(S)
            s = s + K @ r
            P = (np.eye(6) - K @ H) @ P @ (np.eye(6) - K @ H).T + K @ R[n] @ K.T  # Joseph form
        est[k] = s; Ps[k] = P
    out = dict(x=est[:, :3], mu=est[:, 3:], rej_rate=rej / max(tot, 1))
    if return_P:
        out["P"] = Ps
    return out


def late_fusion(obs, ema=None):
    """Per-frame inversion of each available modality reading, averaged per
    dimension; last value held when nothing observes a dimension."""
    est = np.zeros((N, 3)); last = np.zeros(3); sm = np.zeros(3)
    for k in range(N):
        acc = [[] for _ in range(3)]
        for n in MOD_ORDER:
            Y, av = obs[n]
            if not av[k]:
                continue
            H = MODS[n]["H"]
            for i, row in enumerate(H):
                j = int(np.argmax(np.abs(row)))
                acc[j].append(Y[k, i] / row[j])
        cur = np.array([np.mean(a) if a else last[j] for j, a in enumerate(acc)])
        last = cur
        sm = cur if ema is None else (sm + ema * (cur - sm) if k else cur)
        est[k] = sm
    return est


def rmse(e, x, burn=50):
    return np.sqrt(np.mean((e[burn:] - x[burn:]) ** 2, axis=0))


def octant_agree(e, x, burn=50):
    """Circumplex-quadrant agreement: sign match on valence AND arousal."""
    return np.mean(np.all(np.sign(e[burn:, :2]) == np.sign(x[burn:, :2]), axis=1))


METHODS = ["Face-only KF", "Late fusion", "Late fusion+EMA", "Fused KF", "Gated fused KF"]


def run_methods(rng, noise_scale=1.0, p_drop=0.2, p_out=0.0, which=METHODS):
    x, mu = gen_truth(rng)
    obs = gen_obs(rng, x, noise_scale, p_drop, p_out)
    res = {}
    for m in which:
        if m == "Face-only KF":
            e = kalman(obs, mods=["face"], noise_scale=noise_scale)["x"]
        elif m == "Late fusion":
            e = late_fusion(obs)
        elif m == "Late fusion+EMA":
            e = late_fusion(obs, ema=0.15)
        elif m == "Fused KF":
            e = kalman(obs, noise_scale=noise_scale)["x"]
        else:
            e = kalman(obs, gate=True, noise_scale=noise_scale)["x"]
        r = rmse(e, x)
        res[m] = dict(rmse=r, overall=float(np.sqrt(np.mean(r ** 2))), oct=float(octant_agree(e, x)))
    return res


# ---------------- behaviour generation --------------------------------
W = np.diag([0.8, -0.5, -0.3])   # congruent valence, calming arousal, complementary dominance [ASSUMED]
R_MAX, DR_MAX = 0.8, 0.05        # display bounds and per-step rate limit [ASSUMED]


def behaviour(xhat, P, x_true, lam, tau=None, rate=True):
    r = np.zeros(3); out = np.zeros((N, 3))
    for k in range(N):
        c = 1.0 if tau is None else np.exp(-np.trace(P[k][:3, :3]) / tau)
        rstar = c * (W @ xhat[k])
        rn = (rstar + lam * r) / (1 + lam)
        if rate:
            rn = r + np.clip(rn - r, -DR_MAX, DR_MAX)
        r = np.clip(rn, -R_MAX, R_MAX); out[k] = r
    ideal = np.clip(x_true @ W.T, -R_MAX, R_MAX)
    b = 50
    track = float(np.sqrt(np.mean((out[b:] - ideal[b:]) ** 2)))
    jitter = float(np.sqrt(np.mean(np.diff(out[b:], 2, axis=0) ** 2)) / DT ** 2)
    wrong = float(np.mean((np.sign(out[b:, 0]) != np.sign(x_true[b:, 0])) & (np.abs(out[b:, 0]) > 0.2)))
    return out, dict(track=track, jitter=jitter, wrong=wrong)