"""Capacity, inherited behavior, and trainable new features."""
import random

import pytest
import torch

from lightnest.learned_patch import PatchNetwork, batch, load_predictor, make_example
from lightnest.patch_capacity import parameter_counts, widen_model


@pytest.mark.parametrize('width,total,heads', [
    (1, 107994, 714), (2, 405162, 1418), (3, 891514, 2122),
    (4, 1567050, 2826), (8, 6161034, 5642),
])
def test_declared_capacity_round_trip(tmp_path, width, total, heads):
    torch.set_num_threads(2)
    model = PatchNetwork(width=width)
    counts = parameter_counts(model)
    assert counts['total'] == counts['trainable'] == total
    assert counts['head'] == heads
    path = tmp_path / 'model.pt'
    torch.save({'state_dict': model.state_dict(), 'threshold': .55,
                'model_config': model.model_config}, path)
    restored = load_predictor(path).model
    assert restored.width == width
    assert all(torch.equal(restored.state_dict()[k], value)
               for k, value in model.state_dict().items())


@pytest.mark.parametrize('before,after', [(1, 1), (1, 2), (1, 4), (1, 8), (2, 5)])
def test_widening_preserves_outputs_and_parent(before, after):
    torch.set_num_threads(2)
    torch.manual_seed(96171)
    model = PatchNetwork(width=before).eval()
    old = {k: v.clone() for k, v in model.state_dict().items()}
    commands, binaries, _, _ = batch([
        make_example(random.Random(10291 + i), 'combined') for i in range(8)
    ])
    random_before = torch.random.get_rng_state().clone()
    expanded = widen_model(model, after, seed=21419)
    assert torch.equal(torch.random.get_rng_state(), random_before)
    with torch.no_grad():
        for expected, actual in zip(model(commands, binaries), expanded(commands, binaries)):
            torch.testing.assert_close(actual, expected, atol=3e-6, rtol=1e-6)
    assert all(torch.equal(model.state_dict()[k], v) for k, v in old.items())
    assert next(model.parameters()).data_ptr() != next(expanded.parameters()).data_ptr()


def test_new_features_receive_updates_after_connections_learn():
    torch.set_num_threads(2)
    torch.manual_seed(19019)
    expanded = widen_model(PatchNetwork(), 2, seed=72419)
    before = {k: v.clone() for k, v in expanded.state_dict().items()}
    commands, binaries, pointers, replacements = batch([
        make_example(random.Random(62291 + i)) for i in range(8)
    ])
    optimizer = torch.optim.AdamW(expanded.parameters(), lr=.003, weight_decay=0)
    for _ in range(4):
        logits, bits = expanded(commands, binaries)
        positive = pointers != 256
        loss = torch.nn.functional.cross_entropy(logits, pointers)
        loss += torch.nn.functional.binary_cross_entropy_with_logits(
            bits[positive, pointers[positive]], replacements[positive])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    # Without weight decay, new features need gradients to change.
    assert not torch.equal(expanded.command_embedding.weight[:, 24:],
                           before['command_embedding.weight'][:, 24:])
    assert not torch.equal(expanded.binary_embedding.weight[:, 24:],
                           before['binary_embedding.weight'][:, 24:])
    assert not torch.equal(expanded.binary_conv[2].weight[64:],
                           before['binary_conv.2.weight'][64:])
    assert torch.count_nonzero(expanded.pointer.weight[:, 64:])


@pytest.mark.parametrize('config', [None, {'width': True}, {'width': 0},
                                    {'width': 9}, {'width': 2.0}, {'depth': 2}])
def test_invalid_or_undeclared_dimensions_reject(tmp_path, config):
    path = tmp_path / 'bad.pt'
    torch.save({'state_dict': PatchNetwork().state_dict(), 'threshold': .55,
                'model_config': config}, path)
    with pytest.raises(ValueError):
        load_predictor(path)


def test_wider_tensors_require_matching_metadata(tmp_path):
    path = tmp_path / 'bad.pt'
    torch.save({'state_dict': PatchNetwork(width=2).state_dict(), 'threshold': .55}, path)
    with pytest.raises(ValueError, match='checkpoint tensor'):
        load_predictor(path)


def test_checkpoint_size_limit_precedes_deserialization(tmp_path, monkeypatch):
    path = tmp_path / 'oversized.pt'
    with path.open('wb') as stream:
        stream.truncate(32 * 1024**2 + 1)
    def unexpected_load(*args, **kwargs):
        raise AssertionError('oversized file must not reach torch.load')
    monkeypatch.setattr(torch, 'load', unexpected_load)
    with pytest.raises(ValueError, match='32 MiB'):
        load_predictor(path)