import json
from pathlib import Path
import numpy as np


SEED = 2161


def trim_mean(x, alpha):
    x = np.asarray(x)
    n = len(x)
    k = int(np.floor(alpha * n))
    if n - 2 * k <= 0:
        raise ValueError("alpha leaves no retained samples")
    s = np.sort(x, axis=0)
    return s[k:n-k].mean(axis=0)


def clustered_mixture(x, alpha, labels):
    x = np.asarray(x)
    labels = np.asarray(labels)
    comps, weights, retained = [], [], []
    for c in np.unique(labels):
        ids = np.flatnonzero(labels == c)
        z = x[ids]
        k = int(np.floor(alpha * len(z)))
        if len(z) - 2 * k <= 0:
            continue
        comps.append(trim_mean(z, alpha))
        weights.append(len(z) / len(x))
        # A conservative row-retention proxy for the covariance estimate.
        s = np.sort(z, axis=0)
        keep = np.all((z >= s[k]) & (z <= s[len(z)-k-1]), axis=1)
        retained.append(int(max(keep.sum(), 2)))
    weights = np.asarray(weights, dtype=float)
    weights /= weights.sum()
    return np.asarray(comps), weights, retained


def nearest_center_mse(centers, true_centers):
    return float(np.mean([np.min(np.sum((c - true_centers) ** 2, axis=1)) for c in centers]))


def mode_data(rng, separation, n_each=100, noise=.18, outlier_fraction=0.0, outlier_scale=4.0):
    true = np.array([[-separation / 2, 0.0], [separation / 2, 0.0]])
    x = np.vstack([true[0] + rng.normal(0, noise, (n_each, 2)),
                   true[1] + rng.normal(0, noise, (n_each, 2))])
    labels = np.repeat([0, 1], n_each)
    m = int(round(outlier_fraction * len(x)))
    if m:
        ids = rng.choice(len(x), m, replace=False)
        # One-sided adversarial/heavy-tailed contamination.
        x[ids] += np.array([outlier_scale, outlier_scale])
    return x, labels, true


def verify_math(rng):
    # Prediction 1: pi_j=n_j/N, hence weights sum to one for every partition.
    weight_errors = []
    for _ in range(30):
        n = rng.integers(2, 200, size=5)
        weight_errors.append(abs(np.sum(n / n.sum()) - 1.0))
    p1 = {"prediction": "mixture weights sum exactly to one", "max_abs_error": float(max(weight_errors))}

    # Prediction 2: with no contamination and symmetric Gaussian modes, local
    # trimmed centers remain unbiased (error should not grow with separation),
    # while global averaging collapses to the midpoint.
    sep_rows = []
    for sep in [1., 2., 4., 8.]:
        rr = []
        for rep in range(15):
            x, labels, true = mode_data(rng, sep, n_each=100, noise=.18)
            global_center = trim_mean(x, .1)
            labels2 = np.argmin(np.stack([np.abs(x[:, 0] + sep / 2), np.abs(x[:, 0] - sep / 2)]), axis=0)
            comps, w, _ = clustered_mixture(x, .1, labels2)
            rr.append((float(np.mean((global_center - np.mean(true, axis=0)) ** 2)),
                       nearest_center_mse(comps, true)))
        a = np.mean(rr, axis=0)
        sep_rows.append({"separation": sep, "global_midpoint_mse": float(a[0]),
                         "clustered_mode_mse": float(a[1])})
    p2 = {"prediction": "global midpoint error scales as separation^2; local error is separation-invariant",
          "rows": sep_rows}

    # Prediction 3: alpha trimming removes a one-sided contamination fraction q
    # only when q <= alpha (asymptotically); sweep q around alpha.
    trim_rows = []
    alpha = .10
    for q in [0., .05, .10, .15, .25]:
        errs_g, errs_c = [], []
        for _ in range(25):
            x, labels, true = mode_data(rng, 4., n_each=100, noise=.12,
                                         outlier_fraction=q, outlier_scale=5.)
            global_center = trim_mean(x, alpha)
            labels2 = np.argmin(np.stack([np.abs(x[:, 0] + 2.), np.abs(x[:, 0] - 2.)]), axis=0)
            comps, _, _ = clustered_mixture(x, alpha, labels2)
            errs_g.append(np.linalg.norm(global_center - np.mean(true, axis=0)))
            errs_c.append(nearest_center_mse(comps, true) ** .5)
        trim_rows.append({"contamination_q": q, "global_error": float(np.mean(errs_g)),
                          "clustered_error": float(np.mean(errs_c)),
                          "predicted_local_threshold": alpha})
    p3 = {"prediction": "one-sided trimming is effective through q approximately alpha, then residual outlier bias appears",
          "rows": trim_rows}
    return {"weight_normalization": p1, "separation_scaling": p2, "contamination_threshold": p3}


def compare_methods(rng):
    rows = []
    for q in [0., .05, .10, .15, .25]:
        g, t, c = [], [], []
        for _ in range(30):
            x, labels, true = mode_data(rng, 4., n_each=100, noise=.12,
                                         outlier_fraction=q, outlier_scale=5.)
            g.append(np.linalg.norm(x.mean(0) - np.mean(true, axis=0)))
            t.append(np.linalg.norm(trim_mean(x, .1) - np.mean(true, axis=0)))
            labels2 = np.argmin(np.stack([np.abs(x[:, 0] + 2.), np.abs(x[:, 0] - 2.)]), axis=0)
            comps, weights, _ = clustered_mixture(x, .1, labels2)
            # Report both mode error and mode coverage; the mixture preserves two modes.
            c.append((nearest_center_mse(comps, true) ** .5, len(comps), float(weights.min())))
        rows.append({"q": q, "ordinary_mean_error": float(np.mean(g)),
                     "global_trim_error": float(np.mean(t)),
                     "clustered_center_error": float(np.mean([z[0] for z in c])),
                     "cluster_count": float(np.mean([z[1] for z in c])),
                     "smallest_weight": float(np.mean([z[2] for z in c]))})
    return rows


def main():
    rng = np.random.default_rng(SEED)
    result = {"seed": SEED, "verification": verify_math(rng), "comparison": compare_methods(rng)}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()