import sys
from pathlib import Path

import pytest
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.config import PrismeConfig  # noqa: E402
from prisme.model import PrismeModel  # noqa: E402


def tiny_config(**kw) -> PrismeConfig:
    base = dict(name="tiny", vocab_size=97, d_model=64, n_layers=2, n_heads=4, n_kv_heads=2, ffn_width=128,
                max_seq_len=64, w=32, ffn_w=64, n_slots=16, n_editor_heads=2, r_max=8)
    base.update(kw)
    return PrismeConfig(**base)


@pytest.fixture
def cfg():
    return tiny_config()


@pytest.fixture
def model(cfg):
    torch.manual_seed(1234)
    return PrismeModel(cfg).eval()


@pytest.fixture
def model_nonzero(cfg):
    """Model whose readout adapter U has been given nonzero weights (post-bootstrap stand-in)."""
    torch.manual_seed(1234)
    m = PrismeModel(cfg).eval()
    with torch.no_grad():
        m.readout.U.weight.normal_(std=0.05)
    return m


@pytest.fixture
def tokens(cfg):
    g = torch.Generator().manual_seed(7)
    return torch.randint(0, cfg.vocab_size, (3, 12), generator=g)