"""No base weights, network, CUDA or neural inference used by these tests."""
import importlib.util
import json
from pathlib import Path
import struct
import sys
import unittest

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / 'training'))
import infer
from adapters import attach, LoRALinear
from objectives import group_advantages, learning_rate, policy_sum
import torch
from torch import nn


class ReleaseTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.cfg = json.loads((ROOT / 'adapter_config.json').read_text())
        with (ROOT / 'adapter.safetensors').open('rb') as stream:
            size = struct.unpack('<Q', stream.read(8))[0]
            cls.header = json.loads(stream.read(size))

    def test_integrity(self):
        self.assertEqual(infer.sha256(ROOT / 'adapter.safetensors'), self.cfg['adapter_sha256'])

    def test_all_pairs_fp32_and_count(self):
        h = self.header
        self.assertEqual(len(h), 292)
        self.assertEqual(sum(entry['shape'][0] * entry['shape'][1] for entry in h.values()), 96337920)
        for key, entry in h.items():
            self.assertEqual(entry['dtype'], 'F32')
            if key.endswith('.a'):
                pair = h[key[:-1] + 'b']
                self.assertEqual(entry['shape'][0], 128)
                self.assertEqual(pair['shape'][1], 128)

    def test_custom_loader_linear_shapes(self):
        model = nn.Module()
        for key, entry in self.header.items():
            if not key.endswith('.a'):
                continue
            parts = key[:-2].split('.')
            node = model
            for part in parts[:-1]:
                if part not in node._modules:
                    node.add_module(part, nn.Module())
                node = node._modules[part]
            out_features = self.header[key[:-1] + 'b']['shape'][0]
            node.add_module(parts[-1], nn.Linear(entry['shape'][1], out_features, bias=False, device='meta'))
        targets = attach(model, 128)
        self.assertEqual(len(targets), 146)
        shapes = {name: list(p.shape) for name, p in model.named_parameters() if p.requires_grad}
        self.assertEqual(shapes, {name: value['shape'] for name, value in self.header.items()})

    def test_residual_zero_and_fp32(self):
        layer = LoRALinear(nn.Linear(3, 4), rank=2)
        x = torch.randn(5, 3)
        self.assertTrue(torch.equal(layer(x), layer.base(x)))
        self.assertEqual(layer.a.dtype, torch.float32)
        self.assertEqual(layer.b.dtype, torch.float32)

    def test_advantages_are_finite_and_masked(self):
        values = torch.tensor([[.6, .9, .7], [.7, .8, .8], [float('nan')]*3])
        advantage = group_advantages(values, torch.tensor([True, True, False]), [.33, .33, .34])
        self.assertTrue(torch.isfinite(advantage).all())
        self.assertEqual(float(advantage[-1]), 0.)

    def test_schedule_endpoints(self):
        self.assertLess(learning_rate(0, 1500, 5e-6, 5e-7), 5e-6)
        self.assertAlmostEqual(learning_rate(1499, 1500, 5e-6, 5e-7), 5e-7)

    def test_vespo_gradient(self):
        x = torch.tensor([-.3, -.7], requires_grad=True)
        loss = policy_sum(x, x.detach(), torch.tensor(1.))
        loss.backward()
        self.assertTrue(torch.isfinite(x.grad).all())


if __name__ == '__main__':
    unittest.main()