import json, random
from pathlib import Path
import numpy as np
import torch

SEED = 17
np.random.seed(SEED)
random.seed(SEED)
torch.manual_seed(SEED)


def weighted_qr(A, masses):
    """Return Q with Q.T M Q=I, preserving the span of A."""
    sqrt_m = np.sqrt(masses)
    q, _ = np.linalg.qr(sqrt_m[:, None] * A)
    return q / sqrt_m[:, None]


def projector(Q, masses):
    return Q @ Q.T @ np.diag(masses)


def mass_opnorm(A, masses_in, masses_out):
    """Induced M_in -> M_out operator norm."""
    scaled = np.diag(np.sqrt(masses_out)) @ A @ np.diag(1.0 / np.sqrt(masses_in))
    return float(np.linalg.svd(scaled, compute_uv=False)[0])


def build_complex():
    # Oriented filled triangle. d0 is edge-minus-vertex incidence and d1 is
    # the oriented face boundary, so d1 d0=0.
    d0 = np.array([[-1, 1, 0], [-1, 0, 1], [0, -1, 1]], dtype=float)
    d1 = np.array([[1, -1, 1]], dtype=float)
    return d0, d1


def math_check():
    d0, d1 = build_complex()
    m0 = np.array([1.0, 2.0, 1.0])
    m1 = np.array([1.0, 1.5, 2.0])
    # P0=I and P1 projects onto exact 1-cochains im(d0). This removes the
    # one-dimensional cycle component while preserving d(P0 z)=P1(dz).
    Q1 = weighted_qr(d0[:, :2], m1)
    P0 = np.eye(3)
    P1 = projector(Q1, m1)
    # Same-rank independent control (P0=I, random-looking edge subspace).
    R1 = weighted_qr(np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]), m1)
    C1 = projector(R1, m1)

    rows = []
    for name, p1 in [("compatible", P1), ("independent", C1)]:
        idem0 = mass_opnorm(P0 @ P0 - P0, m0, m0)
        idem1 = mass_opnorm(p1 @ p1 - p1, m1, m1)
        defect01 = mass_opnorm(d0 @ P0 - p1 @ d0, m0, m1)
        defect12 = mass_opnorm(d1 @ p1, m1, np.ones(1))
        bound0 = mass_opnorm(P0, m0, m0)
        bound1 = mass_opnorm(p1, m1, m1)
        rows.append(dict(name=name, idempotence_k0=idem0, idempotence_k1=idem1,
                         commutation_defect_d0=defect01,
                         commutation_defect_d1=defect12,
                         bound_k0=bound0, bound_k1=bound1))
    assert np.max(np.abs(d1 @ d0)) < 1e-12
    return d0, m0, m1, P0, P1, rows


def train_experiment(P0, P1, d0, steps=250, seed=SEED):
    """Tiny fixed-seed regression: noisy node/edge features -> exact cochains."""
    torch.manual_seed(seed)
    n, channels = 96, 3
    clean0 = torch.randn(n, 3, channels)
    clean1 = torch.einsum("ei,nic->nec", torch.tensor(d0, dtype=torch.float32), clean0)
    noisy0 = clean0 + 0.55 * torch.randn_like(clean0)
    noisy1 = clean1 + 0.55 * torch.randn_like(clean1)
    w0 = torch.nn.Parameter(torch.randn(channels, channels) * 0.2)
    w1 = torch.nn.Parameter(torch.randn(channels, channels) * 0.2)
    opt = torch.optim.Adam([w0, w1], lr=0.04)
    t0 = torch.tensor(P0, dtype=torch.float32)
    t1 = torch.tensor(P1, dtype=torch.float32)
    for _ in range(steps):
        y0, y1 = noisy0 @ w0, noisy1 @ w1
        z0 = torch.einsum("ij,njc->nic", t0, y0)
        z1 = torch.einsum("ij,njc->nic", t1, y1)
        loss = ((z0 - clean0) ** 2).mean() + ((z1 - clean1) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return float(((z0 - clean0) ** 2).mean() + ((z1 - clean1) ** 2).mean())


def main():
    d0, m0, m1, P0, P1, checks = math_check()
    seeds = [17, 23, 41, 59, 71]
    baseline_vals = [train_experiment(np.eye(3), np.eye(3), d0, seed=s) for s in seeds]
    idea_vals = [train_experiment(P0, P1, d0, seed=s) for s in seeds]
    baseline, idea = baseline_vals[0], idea_vals[0]
    result = {
        "seed": SEED,
        "repeat_seeds": seeds,
        "incidence_d1_d0_max": 0.0,
        "checks": checks,
        "denoising_mse": {"baseline_no_projection": baseline,
                          "bounded_commuting_layer": idea},
        "relative_mse_reduction": (baseline - idea) / baseline,
        "repeat_mse_mean_std": {
            "baseline_no_projection": [float(np.mean(baseline_vals)), float(np.std(baseline_vals))],
            "bounded_commuting_layer": [float(np.mean(idea_vals)), float(np.std(idea_vals))],
            "relative_reduction_mean": float(np.mean((np.array(baseline_vals)-np.array(idea_vals))/np.array(baseline_vals)))}
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()