"""Shared test fixtures for ai_mouse."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest
import torch

from ai_mouse.scroll.models import ScrollCVAE
from tools.models import TrajectoryFlowModel


@pytest.fixture
def model_dir(tmp_path: Path) -> Path:
    """Create a temporary directory with trained Flow model artifacts."""
    # Flow model
    model = TrajectoryFlowModel(seq_len=64)
    torch.save(model.state_dict(), tmp_path / "flow_model.pt")

    # Click distribution
    click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
    (tmp_path / "click_dist.json").write_text(json.dumps(click_dist))

    # Duration distribution
    dur_dist = {
        "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
        "params": [{"mu_log": 5.5, "sigma_log": 0.5}] * 8,
    }
    (tmp_path / "duration_dist.json").write_text(json.dumps(dur_dist))

    # Train config (architecture params)
    train_cfg = {
        "seq_len": 64,
        "d_model": 128,
        "nhead": 4,
        "num_layers": 4,
        "dim_feedforward": 256,
        "cond_dim": 3,
    }
    (tmp_path / "train_config.json").write_text(json.dumps(train_cfg))

    return tmp_path


@pytest.fixture
def scroll_model_dir(tmp_path: Path) -> Path:
    """Create a temporary directory with trained scroll model artifacts."""
    model = ScrollCVAE(seq_len=32)
    torch.save(model.state_dict(), tmp_path / "scroll_model.pt")

    scroll_cfg = {"seq_len": 32, "latent_dim": 16, "hidden": 64, "cond_dim": 7}
    (tmp_path / "scroll_config.json").write_text(json.dumps(scroll_cfg))

    return tmp_path