import json
import math
import os
import random
import numpy as np
from sklearn.metrics import f1_score

SEED = 2057
rng = np.random.default_rng(SEED)


def group_prox(B, lr_lam):
    # B shape [outputs, regulators]; one regulator is one regulator-to-state group.
    norms = np.linalg.norm(B, axis=0, keepdims=True)
    return B * np.maximum(0.0, 1.0 - lr_lam / (norms + 1e-12))


def fit_group(X, Y, lam=0.02, epochs=500, lr=None):
    # Minimize mean squared multivariate regression loss + group lasso.
    X = np.asarray(X); Y = np.asarray(Y)
    n, p = X.shape
    if lr is None:
        lr = 0.8 / (np.linalg.norm(X, 2) ** 2 / max(n, 1) + 1e-8)
    W = np.zeros((p, Y.shape[1]))
    # Work with rows as regulator groups.
    for _ in range(epochs):
        G = (X.T @ (X @ W - Y)) / n
        W -= lr * G
        W = group_prox(W.T, lr * lam).T
    return W


def make_data(n=5000, d=8, q=12, dt_choices=(0.05, 0.1, 0.2), sigma=0.05,
              rho=0.0, seed=0):
    rg = np.random.default_rng(seed)
    active = np.array([0, 3, 7])
    B = np.zeros((q, d))
    # correlated regulator features, with a controllable active Gram condition.
    z = rg.normal(size=(n, 1))
    R = rg.normal(size=(n, q))
    for j in active[1:]:
        R[:, j] = rho * R[:, 0] + math.sqrt(max(1-rho*rho, 0)) * R[:, j]
    B[active] = rg.normal(0, 0.65, size=(len(active), d))
    dt = rg.choice(dt_choices, size=n)
    # latent increments from the integral equation, noisy observations at endpoints.
    clean = (dt[:, None] * (R @ B))
    eps0 = rg.normal(0, sigma, size=(n, d))
    eps1 = rg.normal(0, sigma, size=(n, d))
    y_int = clean + eps1 - eps0
    # Derivative target is exactly the same noisy increment divided by dt.
    y_fd = y_int / dt[:, None]
    return R, dt, y_int, y_fd, B, active


def noise_scaling():
    # Directly test the two claims: Var(delta x noise)=2 sigma^2 and
    # Var((delta x)/dt noise)=2 sigma^2/dt^2.
    sigma = 0.07
    out = []
    for dt in [0.025, 0.05, 0.1, 0.2, 0.4]:
        rg = np.random.default_rng(1000 + int(dt * 10000))
        e = rg.normal(0, sigma, size=(250000, 1))
        f = rg.normal(0, sigma, size=(250000, 1))
        inc = (f-e).ravel()
        fd = inc / dt
        vi, vf = float(np.var(inc)), float(np.var(fd))
        out.append({"dt": dt, "integral_var": vi, "integral_pred": 2*sigma*sigma,
                    "fd_var": vf, "fd_pred": 2*sigma*sigma/(dt*dt),
                    "fd_over_integral": vf/vi, "ratio_pred": 1/(dt*dt)})
    return out


def support_sweep():
    # Prediction: as rho -> 1, lambda_min(active Gram) -> 0 and support recovery degrades.
    rows = []
    for rho in [0.0, 0.5, 0.8, 0.95, 0.99]:
        R, dt, yi, yf, B, active = make_data(n=7000, sigma=0.045, rho=rho, seed=20+int(rho*100))
        # Integral formulation uses integrated features dt*r and increments.
        Xi = R * dt[:, None]
        W = fit_group(Xi, yi, lam=0.018, epochs=350)
        # FD formulation uses r and derivative targets; same equal-weight objective.
        Wfd = fit_group(R, yf, lam=0.018, epochs=350)
        gram = (R[:, active].T @ R[:, active]) / len(R)
        eig = float(np.linalg.eigvalsh(gram).min())
        def score(w):
            pred = np.linalg.norm(w, axis=1) > 0.06
            truth = np.zeros(w.shape[0], dtype=bool); truth[active] = True
            return float(f1_score(truth, pred)), int(pred.sum())
        fi, ni = score(W); ff, nf = score(Wfd)
        rows.append({"rho": rho, "lambda_min": eig, "integral_f1": fi,
                     "fd_f1": ff, "integral_edges": ni, "fd_edges": nf})
    return rows


def method_comparison():
    # Irregular intervals: integrated loss naturally weights observations by physical
    # increment size, while FD noise is amplified on the short intervals.
    R, dt, yi, yf, B, active = make_data(n=12000, sigma=0.06,
                                          dt_choices=(0.025, 0.05, 0.1, 0.2), rho=0.5, seed=77)
    Xi = R * dt[:, None]
    Wi = fit_group(Xi, yi, lam=0.0008, epochs=450)
    Wf = fit_group(R, yf, lam=0.02, epochs=450)
    # Evaluate physical increment prediction on fresh noiseless transitions.
    R2, dt2, y2, yf2, _, _ = make_data(n=4000, sigma=0.0,
                                         dt_choices=(0.025, 0.05, 0.1, 0.2), rho=0.5, seed=91)
    true_inc = dt2[:, None] * (R2 @ B)
    pi = dt2[:, None] * (R2 @ Wi)
    pf = dt2[:, None] * (R2 @ Wf)
    def f1(w):
        pred = np.linalg.norm(w, axis=1) > 0.06
        truth = np.zeros(w.shape[0], dtype=bool); truth[active] = True
        return float(f1_score(truth, pred))
    return {"integral_clean_increment_mse": float(np.mean((pi-true_inc)**2)),
            "finite_difference_clean_increment_mse": float(np.mean((pf-true_inc)**2)),
            "integral_support_f1": f1(Wi), "finite_difference_support_f1": f1(Wf),
            "active_edges": active.tolist()}


if __name__ == "__main__":
    result = {"seed": SEED,
              "prediction_1_noise_scaling": noise_scaling(),
              "prediction_2_conditioning_support": support_sweep(),
              "mini_experiment": method_comparison()}
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))