"""Braid-monodromy set-state implementation.

This module is independent of the missing shared bench checkout.  It contains
matched sequence classifiers for local smoke testing: both consume unordered
object sets, while BraidClassifier replaces the GRU recurrence by a skew
Cayley transport plus innovation.
"""
import json
from pathlib import Path
import torch
from torch import nn


def cayley(A, dt=1.0):
    k = A.shape[-1]
    I = torch.eye(k, device=A.device, dtype=A.dtype).expand_as(A)
    return torch.linalg.solve(I - .5 * dt * A, I + .5 * dt * A)


class SetEncoder(nn.Module):
    def __init__(self, d=4, width=24):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d, width), nn.Tanh(),
                                 nn.Linear(width, width), nn.Tanh())
    def forward(self, x):
        return self.net(x).mean(dim=2)


class DeepSetGRU(nn.Module):
    def __init__(self, d=4, width=24, hidden=32, classes=2):
        super().__init__()
        self.enc = SetEncoder(d, width)
        self.rnn = nn.GRU(width, hidden, batch_first=True)
        self.out = nn.Linear(hidden, classes)
    def forward(self, x):
        h, _ = self.rnn(self.enc(x))
        return self.out(h[:, -1])


class BraidClassifier(nn.Module):
    def __init__(self, d=4, width=24, k=8, hidden=32, classes=2):
        super().__init__()
        self.k = k
        self.width = width
        self.item = SetEncoder(d, width).net
        self.pair = nn.Sequential(nn.Linear(2*width + d, 32), nn.Tanh(),
                                  nn.Linear(32, 32), nn.Tanh())
        self.gen = nn.Linear(32, k*k)
        self.weight = nn.Linear(32, 1)
        self.zproj = nn.Linear(width, hidden)
        self.innov = nn.Sequential(nn.Linear(hidden+k, hidden), nn.Tanh(),
                                   nn.Linear(hidden, k))
        self.out = nn.Sequential(nn.Linear(hidden+k, 32), nn.Tanh(),
                                 nn.Linear(32, classes))
    def forward(self, x, return_state=False):
        B, T, N, D = x.shape
        h = x.new_zeros(B, self.k)
        z = x.new_zeros(B, self.width)
        for t in range(T):
            xt = x[:, t]
            e = self.item(xt)
            z = e.mean(1)
            A = x.new_zeros(B, self.k, self.k)
            for i in range(N):
                for j in range(i+1, N):
                    q = self.pair(torch.cat((e[:, i] + e[:, j],
                                             (e[:, i] - e[:, j]).abs(),
                                             (xt[:, i] - xt[:, j]).abs()), dim=-1))
                    raw = self.gen(q).view(B, self.k, self.k)
                    A = A + torch.sigmoid(self.weight(q)).view(B, 1, 1) * (raw - raw.transpose(1, 2))
            h = (cayley(A) @ h.unsqueeze(-1)).squeeze(-1)
            h = h + self.innov(torch.cat((torch.tanh(self.zproj(z)), h), dim=-1))
        logits = self.out(torch.cat((torch.tanh(self.zproj(z)), h), dim=-1))
        return (logits, h, A) if return_state else logits


def math_check(seed=0, trials=32, k=8):
    torch.manual_seed(seed)
    skew_err, cay_err, euler_err = [], [], []
    for _ in range(trials):
        B = torch.randn(k, k); A = B - B.T; v = torch.randn(k)
        R = cayley(A, .2)
        skew_err.append((A + A.T).norm().item())
        cay_err.append(abs((R @ v).norm() - v.norm()).item())
        euler_err.append(abs(((torch.eye(k) + .2*A) @ v).norm() - v.norm()).item())
    return {'max_skew_residual': max(skew_err),
            'max_cayley_norm_error': max(cay_err),
            'mean_euler_control_error': sum(euler_err)/len(euler_err)}


BENCH_REPORT = {
  'status': 'unavailable',
  'reason': 'Required /home/maxwelhelp/all/math2nn/bench and bench/README.md are absent; no fixed-harness Stage-2 run was possible.',
  'track': 'dynamics (structurally appropriate, but not executed)',
  'baseline_sweep': None, 'idea_per_seed': None, 'paired_delta': None,
  'permutation_p_value': None,
  'mechanism_signature': {
      'predicted': 'skew Cayley transport preserves fiber norm absent innovation',
      'observed': None, 'confirmed': False
  }
}

if __name__ == '__main__':
    print(json.dumps({'math_check': math_check(), 'bench_report': BENCH_REPORT}, indent=2))