import json, os, random
import numpy as np
from scipy.optimize import linear_sum_assignment
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

SEED = 17
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def positive_axis_match(z_a, z_b, eps=1e-8):
    # Rows are examples, columns are axes. Cost is normalized residual after
    # independently fitting the allowed positive scale for each candidate pair.
    va = np.var(z_a, axis=0) + eps
    za2 = np.sum(z_a * z_a, axis=0) + eps
    cross = z_a.T @ z_b
    alpha = np.maximum(eps, cross / za2[:, None])
    residual = z_a[:, :, None] * alpha[None, :, :] - z_b[:, None, :]
    cost = np.mean(residual * residual, axis=0) / va[:, None]
    rows, cols = linear_sum_assignment(cost)
    return rows, cols, alpha[rows, cols], cost[rows, cols]


def affine_fit_residual(z_a, z_b):
    # Fit E(z_b)=z_a, including an intercept, and return RMS residual.
    xb = np.concatenate([z_b, np.ones((len(z_b), 1))], axis=1)
    coef, *_ = np.linalg.lstsq(xb, z_a, rcond=None)
    pred = xb @ coef
    return float(np.sqrt(np.mean((pred - z_a) ** 2)))


def math_sanity():
    rng = np.random.default_rng(SEED)
    n, d = 1200, 8
    za = rng.normal(size=(n, d))
    perm = rng.permutation(d)
    scales = rng.uniform(.35, 2.4, size=d)
    values = []
    for noise in [0.0, .02, .05, .10, .20, .40]:
        zb = za[:, perm] * scales[None, :] + noise * rng.normal(size=(n, d))
        rows, cols, alpha, costs = positive_axis_match(za, zb)
        recovered = np.mean(cols == np.argsort(perm))
        axis_quality = float(np.mean(np.exp(-costs)))
        weak = affine_fit_residual(za, zb)
        values.append({"noise": noise, "affine_rms": weak,
                       "axis_quality": axis_quality,
                       "permutation_recovery": float(recovered),
                       "median_normalized_cost": float(np.median(costs))})
    # A control with unrelated coordinates tests that matching is not automatic.
    zb = rng.normal(size=(n, d))
    _, _, _, costs = positive_axis_match(za, zb)
    control = {"random_axis_quality": float(np.mean(np.exp(-costs))),
               "random_median_cost": float(np.median(costs))}
    return {"sweep": values, "control": control}


class MLP(nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.fc1 = nn.Linear(64, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.out = nn.Linear(hidden, 10)
    def forward(self, x, return_z=False):
        z1 = self.fc1(x); h1 = torch.relu(z1)
        z2 = self.fc2(h1); h2 = torch.relu(z2)
        y = self.out(h2)
        return (y, [z1, z2]) if return_z else y


def torch_axis_loss(za, zb):
    # Differentiable once the permutation is stop-gradient/fixed for this batch.
    with torch.no_grad():
        _, cols, alpha, _ = positive_axis_match(za.detach().cpu().numpy(), zb.detach().cpu().numpy())
    loss = 0.0
    count = 0
    for j, i, a in zip(range(len(cols)), cols, alpha):
        # Match teacher axis j to assigned student axis i.
        denom = za[:, j].var(unbiased=False).detach() + 1e-5
        loss = loss + ((zb[:, i] - float(a) * za[:, j]) ** 2 / denom).mean()
        count += 1
    return loss / max(count, 1)


def train_teacher_and_students():
    d = load_digits()
    x = StandardScaler().fit_transform(d.data).astype("float32")
    y = d.target.astype("int64")
    xt, xv, yt, yv = train_test_split(x, y, test_size=.25, random_state=SEED, stratify=y)
    Xtr, Ytr = torch.tensor(xt, device=DEVICE), torch.tensor(yt, device=DEVICE)
    Xva, Yva = torch.tensor(xv, device=DEVICE), torch.tensor(yv, device=DEVICE)
    teacher = MLP().to(DEVICE)
    opt = torch.optim.Adam(teacher.parameters(), lr=2e-3)
    for step in range(350):
        idx = torch.randint(len(Xtr), (96,), device=DEVICE)
        loss = nn.functional.cross_entropy(teacher(Xtr[idx]), Ytr[idx])
        opt.zero_grad(); loss.backward(); opt.step()
    teacher.eval()
    with torch.no_grad():
        ty, tz = teacher(Xtr, True); tv, _ = teacher(Xva, True)
        tacc = float((tv.argmax(1) == Yva).float().mean())
    results = {}
    for kind in ["baseline", "axis"]:
        torch.manual_seed(SEED + (0 if kind == "baseline" else 1))
        student = MLP().to(DEVICE)
        opt = torch.optim.Adam(student.parameters(), lr=2e-3)
        last_axis = None
        for step in range(350):
            idx = torch.randint(len(Xtr), (96,), device=DEVICE)
            sx, sz = student(Xtr[idx], True)
            with torch.no_grad():
                target = ty[idx]
            task = nn.functional.cross_entropy(sx, Ytr[idx])
            kd = nn.functional.mse_loss(sx, target)
            total = task + .15 * kd
            if kind == "axis":
                total = total + .03 * (torch_axis_loss(tz[0][idx], sz[0]) + torch_axis_loss(tz[1][idx], sz[1]))
            opt.zero_grad(); total.backward(); opt.step()
            if step % 100 == 0:
                with torch.no_grad():
                    _, szfull = student(Xtr, True)
                    _, cols, _, costs = positive_axis_match(tz[0].cpu().numpy(), szfull[0].cpu().numpy())
                    last_axis = {"quality": float(np.mean(np.exp(-costs))),
                                 "matched_fraction_cost_lt_0.25": float(np.mean(costs < .25)),
                                 "permutation": cols.tolist()}
        student.eval()
        with torch.no_grad():
            sy, _ = student(Xva, True)
            val_loss = float(nn.functional.cross_entropy(sy, Yva))
            acc = float((sy.argmax(1) == Yva).float().mean())
            train_loss = float(nn.functional.cross_entropy(student(Xtr), Ytr))
        results[kind] = {"val_accuracy": acc, "val_loss": val_loss,
                         "train_ce": train_loss, "final_assignment": last_axis}
    return {"device": DEVICE, "teacher_accuracy": tacc, "students": results}


def main():
    out = {"seed": SEED, "math_sanity": math_sanity(), "distillation": train_teacher_and_students()}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, AssertionError) as e:
        if torch.cuda.is_available() and DEVICE == "cuda":
            print("CUDA failed; rerun with CPU:", repr(e))
            os.environ["CUDA_VISIBLE_DEVICES"] = ""
            # Explicit CPU fallback in a fresh process avoids partially allocated CUDA state.
            os.execv("/home/maxwelhelp/main/bin/python3", ["python3", __file__])
        raise