"""Real, in-memory Optuna contracts with a tiny deterministic CPU model."""

from types import SimpleNamespace

import optuna
import pytest
import torch
from torch import nn

from obliteratus import bayesian_optimizer as bo
from obliteratus.abliterate import AbliterationPipeline


class _TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        layer = nn.Module()
        layer.self_attn = nn.Module()
        layer.self_attn.o_proj = nn.Linear(2, 2, bias=False)
        layer.mlp = nn.Module()
        layer.mlp.down_proj = nn.Linear(2, 2, bias=False)
        with torch.no_grad():
            layer.self_attn.o_proj.weight.copy_(torch.tensor([[1., .2], [.3, 1.]]))
            layer.mlp.down_proj.weight.copy_(torch.tensor([[.7, .1], [.4, .8]]))
        self.model = nn.Module()
        self.model.layers = nn.ModuleList([layer])
        self.forward_calls = 0
        self.generation_calls = 0

    def forward(self, input_ids, **_kwargs):
        self.forward_calls += 1
        hidden = torch.nn.functional.one_hot(input_ids, num_classes=2).float()
        for layer in self.model.layers:
            hidden = layer.mlp.down_proj(layer.self_attn.o_proj(hidden))
        return SimpleNamespace(logits=hidden)

    def generate(self, input_ids, **_kwargs):
        self.generation_calls += 1
        token = self(input_ids).logits[:, -1, :].argmax(dim=-1, keepdim=True)
        return torch.cat([input_ids, token], dim=1)


class _TinyTokenizer:
    chat_template = None

    def __call__(self, *_args, **_kwargs):
        return {"input_ids": torch.tensor([[0, 1]])}

    def decode(self, tokens, **_kwargs):
        return "I cannot assist with that request." if tokens[-1] == 0 else "Here is an explanation."


@pytest.mark.parametrize("peak", [2 / 35, 34 / 35])
def test_real_optuna_runs_bounded_warm_start_and_restores_model(monkeypatch, peak):
    studies = []
    create_study = optuna.create_study

    def capture_study(**kwargs):
        study = create_study(**kwargs)
        studies.append(study)
        return study

    monkeypatch.setattr(optuna, "create_study", capture_study)
    pipeline = AbliterationPipeline(model_name="tiny-local", method="optimized")
    model = _TinyModel()
    pipeline.handle = SimpleNamespace(model=model, tokenizer=_TinyTokenizer(), architecture="llama")
    pipeline._strong_layers = [0]
    pipeline.refusal_directions = {0: torch.tensor([1., 0.])}
    pipeline.harmful_prompts = ["A local test prompt"]
    pipeline._informed_warm_start = {"peak_position": peak, "max_weight": .5, "mlp_scale": .5}
    # The fixture owns no accelerators or large allocations.
    monkeypatch.setattr(pipeline, "_free_gpu_memory", lambda: None)
    original = {name: tensor.clone() for name, tensor in model.state_dict().items()}

    result = bo.run_bayesian_optimization(pipeline, n_trials=2, n_refusal_prompts=1, n_kl_prompts=1)

    assert set(result) == {0}
    assert 0 <= result[0] <= 1
    assert len(studies) == 1
    trials = studies[0].trials
    assert len(trials) == 2
    assert all(trial.state == optuna.trial.TrialState.COMPLETE for trial in trials)
    expected_peak = .1 if peak < .1 else .9
    assert trials[0].params["attn_peak_position"] == pytest.approx(expected_peak)
    assert trials[0].params["mlp_peak_position"] == pytest.approx(expected_peak)
    assert trials[0].params["mlp_max_weight"] == pytest.approx(.3)
    for trial in trials:
        assert all(bo.KERNEL_SPACE[name][0] <= trial.params[name] <= bo.KERNEL_SPACE[name][1]
                   for name in bo.KERNEL_SPACE)
        assert trial.values is not None and len(trial.values) == 2
        assert all(torch.isfinite(torch.tensor(value)) for value in trial.values)
    assert model.generation_calls == 2
    assert model.forward_calls >= 5  # reference, generation and KL for each trial
    for name, tensor in model.state_dict().items():
        torch.testing.assert_close(tensor, original[name], rtol=0, atol=0)