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

SEED = 3145
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')

# C4 acts on both 2D input blocks (state and action) and on the 2D output.
def R(k, device=None):
    t = float(k % 4) * np.pi / 2
    c, s = np.cos(t), np.sin(t)
    return torch.tensor([[c, -s], [s, c]], dtype=torch.float32, device=device)

def transform(z, k):
    q = R(k, z.device)
    return torch.cat([z[:, :2] @ q.T, z[:, 2:] @ q.T], dim=1)

def transform_y(y, k):
    return y @ R(k, y.device).T

def true_base(z):
    x, a = z[:, :2], z[:, 2:]
    # Nonlinear but exactly O(2)-equivariant: scalar invariants multiply vectors.
    s = (x*x).sum(1, keepdim=True) + 0.35*(a*a).sum(1, keepdim=True)
    return (0.65 + 0.12*torch.tanh(s))*x + (0.42 - 0.08*torch.sigmoid(s))*a

def make_data(n, ks):
    z0 = torch.randn(n, 4) * torch.tensor([2.0, 0.35, 1.5, 0.25])
    # Every orientation is the same embedded local mechanism under C4.
    kvals = torch.tensor(np.random.choice(ks, n), dtype=torch.long)
    zs, ys = [], []
    for z, k in zip(z0, kvals):
        zz = transform(z[None], int(k)).squeeze(0)
        yy = transform_y(true_base(z[None]), int(k)).squeeze(0)
        zs.append(zz); ys.append(yy)
    return torch.stack(zs), torch.stack(ys)

class MLP(nn.Module):
    def __init__(self, hidden=48):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
    def forward(self, z): return self.net(z)

def eq_penalty(model, z):
    # Exact sampled analogue of E_h ||F(rho_in(h)z)-rho_out(h)F(z)||^2.
    y = model(z)
    vals = []
    for k in range(4):
        vals.append((model(transform(z, k)) - transform_y(y, k)).pow(2).mean())
    return torch.stack(vals).mean()

def fit(lam, train_z, train_y, epochs=450):
    model = MLP().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3, weight_decay=1e-5)
    z, y = train_z.to(device), train_y.to(device)
    for _ in range(epochs):
        opt.zero_grad()
        pred = model(z)
        loss = (pred-y).pow(2).mean() + lam*eq_penalty(model, z)
        loss.backward(); opt.step()
    return model

def mse(model, z, y):
    with torch.no_grad(): return float((model(z.to(device))-y.to(device)).pow(2).mean().cpu())

def main():
    # Structural sanity checks corresponding to the occurrence symmetry: C4 composition,
    # invertibility, and exact equivariance of the generated shared mechanism.
    z = torch.randn(64, 4)
    composition_err = max(float((transform(transform(z, a), b)-transform(z, (a+b)%4)).abs().max()) for a in range(4) for b in range(4))
    mechanism_err = max(float((true_base(transform(z, k))-transform_y(true_base(z), k)).abs().max()) for k in range(4))
    linear_good = nn.Linear(4, 2, bias=False)
    with torch.no_grad():
        linear_good.weight.copy_(torch.tensor([[1.,0.,0.5,0.],[0.,1.,0.,0.5]]))
    penalty_good = float(eq_penalty(linear_good, z))
    bad = nn.Linear(4, 2)
    with torch.no_grad(): bad.weight.copy_(torch.tensor([[1.7,0.2,-.4,.3],[.1,.3,.8,-.2]])); bad.bias.fill_(.4)
    penalty_bad = float(eq_penalty(bad, z))

    # Train on three orientations and test on a held-out fourth occurrence symmetry.
    train_z, train_y = make_data(768, [0,1,2])
    test_id_z, test_id_y = make_data(384, [0,1,2])
    test_ood_z, test_ood_y = make_data(384, [3])
    baseline = fit(0.0, train_z, train_y)
    equiv = fit(1.0, train_z, train_y)
    results = {
        'seed': SEED, 'device': str(device),
        'math_check': {'c4_composition_max_abs': composition_err,
                       'true_mechanism_equivariance_max_abs': mechanism_err,
                       'equivariance_penalty_exact_linear': penalty_good,
                       'equivariance_penalty_non_equivariant_linear': penalty_bad},
        'params_each_model': sum(p.numel() for p in baseline.parameters()),
        'baseline': {'train_mse': mse(baseline, train_z, train_y), 'in_distribution_mse': mse(baseline, test_id_z, test_id_y), 'heldout_rotation_mse': mse(baseline, test_ood_z, test_ood_y), 'heldout_equivariance_penalty': float(eq_penalty(baseline, test_ood_z.to(device)).cpu())},
        'equivariant_shared': {'train_mse': mse(equiv, train_z, train_y), 'in_distribution_mse': mse(equiv, test_id_z, test_id_y), 'heldout_rotation_mse': mse(equiv, test_ood_z, test_ood_y), 'heldout_equivariance_penalty': float(eq_penalty(equiv, test_ood_z.to(device)).cpu())}
    }
    Path('results.json').write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        if str(device) == 'cuda':
            print('CUDA failed; rerun on CPU:', repr(e))
            device = torch.device('cpu'); main()
        else: raise