import json
import random
import time
from collections import Counter
from itertools import combinations

import numpy as np
import torch

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


def choose_device():
    if torch.cuda.is_available():
        try:
            torch.zeros(1, device="cuda")
            return "cuda"
        except Exception:
            pass
    return "cpu"


class Primitive:
    def __init__(self, name, support, kind, weight, bias, scale):
        self.name = name
        self.support = frozenset(support)
        self.kind = kind
        self.weight = weight
        self.bias = bias
        self.scale = scale

    def apply(self, x):
        # Every primitive reads/writes only its declared coordinates.
        idx = sorted(self.support)
        y = x.clone()
        z = x[..., idx]
        if self.kind == "affine":
            z = z @ self.weight + self.bias
        elif self.kind == "tanh_affine":
            z = torch.tanh(z @ self.weight + self.bias)
        elif self.kind == "scale":
            z = z * self.scale
        y[..., idx] = z
        return y


def build_library(dim=12, groups=4, per_group=2, device="cpu"):
    ops = {}
    # Three independently parameterized operators per group. Operators within
    # one group need not commute; operators from different groups do commute.
    for g in range(groups):
        support = list(range(g * per_group, (g + 1) * per_group))
        for j, kind in enumerate(("affine", "tanh_affine", "scale")):
            name = f"g{g}_{j}"
            if kind == "scale":
                weight = bias = None
                scale = torch.rand(per_group, device=device) * 0.8 + 0.6
            else:
                weight = torch.randn(per_group, per_group, device=device) * 0.35
                bias = torch.randn(per_group, device=device) * 0.1
                scale = None
            ops[name] = Primitive(name, support, kind, weight, bias, scale)
    return ops


def commute(a, b, ops):
    return bool(ops[a].support.isdisjoint(ops[b].support))


def canonical(word, order, ops):
    # Bubble-sort only inverted adjacent commuting letters. This is the
    # requested local rewrite system for a fixed total order.
    rank = {name: i for i, name in enumerate(order)}
    out = list(word)
    changed = True
    while changed:
        changed = False
        for i in range(len(out) - 1):
            a, b = out[i], out[i + 1]
            if rank[a] > rank[b] and commute(a, b, ops):
                out[i], out[i + 1] = b, a
                changed = True
    return tuple(out)


def run_word(word, ops, x):
    y = x
    for name in word:
        y = ops[name].apply(y)
    return y


def main():
    device = choose_device()
    ops = build_library(device=device)
    order = sorted(ops)
    names = list(ops)
    x = torch.randn(32, 12, device=device)

    # Core mathematical sanity check: every proposed adjacent commuting swap
    # should preserve the complete tensor output to numerical precision.
    swap_errors = []
    noncommuting_differences = []
    rng = random.Random(SEED + 1)
    for _ in range(300):
        a, b = rng.sample(names, 2)
        prefix = [rng.choice(names) for _ in range(3)]
        suffix = [rng.choice(names) for _ in range(3)]
        w1 = prefix + [a, b] + suffix
        w2 = prefix + [b, a] + suffix
        err = (run_word(w1, ops, x) - run_word(w2, ops, x)).abs().max().item()
        if commute(a, b, ops):
            swap_errors.append(err)
        else:
            noncommuting_differences.append(err)
    max_commuting_error = max(swap_errors)
    median_noncommuting_difference = float(np.median(noncommuting_differences))

    # Canonical words are functionally equal to their original words.
    canonical_errors = []
    for _ in range(100):
        word = tuple(rng.choice(names) for _ in range(12))
        cw = canonical(word, order, ops)
        canonical_errors.append((run_word(word, ops, x) - run_word(cw, ops, x)).abs().max().item())

    # Search/cache toy: fixed 500 proposed traces. Baseline evaluates all
    # strings; canonicalized search evaluates one forward per equivalence key.
    proposals = []
    for _ in range(50):
        base = [rng.choice(names) for _ in range(12)]
        for _ in range(10):
            v = list(base)
            for _ in range(20):
                i = rng.randrange(len(v) - 1)
                if commute(v[i], v[i + 1], ops):
                    v[i], v[i + 1] = v[i + 1], v[i]
            proposals.append(tuple(v))
    baseline_unique = len(set(proposals))
    canonical_keys = [canonical(w, order, ops) for w in proposals]
    idea_unique = len(set(canonical_keys))
    # Use a deterministic target and an actual scalar objective so cache hits
    # represent the same architecture/function, not merely string equality.
    target = torch.randn_like(x)
    t0 = time.perf_counter()
    baseline_losses = [torch.mean((run_word(w, ops, x) - target) ** 2).item() for w in proposals]
    baseline_time = time.perf_counter() - t0
    t0 = time.perf_counter()
    cache = {}
    idea_losses = []
    for w, key in zip(proposals, canonical_keys):
        if key not in cache:
            cache[key] = torch.mean((run_word(key, ops, x) - target) ** 2).item()
        idea_losses.append(cache[key])
    idea_time = time.perf_counter() - t0
    max_loss_disagreement = max(abs(a - b) for a, b in zip(baseline_losses, idea_losses))

    result = {
        "seed": SEED,
        "device": device,
        "alphabet_size": len(names),
        "word_length": 12,
        "random_proposals": len(proposals),
        "commuting_swap_trials": len(swap_errors),
        "max_commuting_swap_error": max_commuting_error,
        "median_noncommuting_swap_difference": median_noncommuting_difference,
        "max_word_vs_canonical_error": max(canonical_errors),
        "baseline_unique_strings": baseline_unique,
        "canonical_unique_classes": idea_unique,
        "deduplication_fraction": 1.0 - idea_unique / len(proposals),
        "baseline_mean_loss": float(np.mean(baseline_losses)),
        "idea_mean_loss": float(np.mean(idea_losses)),
        "max_baseline_vs_cached_loss_disagreement": max_loss_disagreement,
        "baseline_forward_seconds": baseline_time,
        "idea_forward_seconds": idea_time,
        "forward_reduction_fraction": 1.0 - idea_unique / len(proposals),
        "cache_speedup": baseline_time / idea_time if idea_time else None,
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        # CUDA can fail in a shared environment; retrying explicitly on CPU
        # keeps the experiment reproducible rather than silently fabricating it.
        if torch.cuda.is_available():
            print(json.dumps({"cuda_error": repr(exc), "retry": "cpu"}))
            torch.cuda.is_available = lambda: False
            main()
        else:
            raise