"""Tests for the LoRA mixed-type encoder."""

from __future__ import annotations

import numpy as np
import pytest

from ahdcma.search_space.encoder import (
    DIM,
    R_OPTIONS,
    SCHEDULE_OPTIONS,
    TARGET_MODULE_OPTIONS,
    decode,
    encode,
)

VALID_CONFIGS = [
    {
        "r": 1,
        "alpha": 1.0,
        "dropout": 0.0,
        "lr": 1e-5,
        "schedule": "cosine",
        "target_modules": ["q"],
    },
    {
        "r": 8,
        "alpha": 16.0,
        "dropout": 0.1,
        "lr": 3e-4,
        "schedule": "linear",
        "target_modules": ["q", "v"],
    },
    {
        "r": 64,
        "alpha": 256.0,
        "dropout": 0.5,
        "lr": 1e-2,
        "schedule": "polynomial",
        "target_modules": ["q", "k", "v", "o", "fc1", "fc2"],
    },
    {
        "r": 16,
        "alpha": 32.0,
        "dropout": 0.25,
        "lr": 1e-3,
        "schedule": "constant",
        "target_modules": ["k", "fc2"],
    },
]


@pytest.mark.parametrize("cfg", VALID_CONFIGS)
def test_round_trip_decode_encode(cfg: dict) -> None:
    x = encode(cfg)
    assert x.shape == (DIM,)
    assert np.all((x >= 0.0) & (x <= 1.0))
    cfg2 = decode(x)
    assert cfg2["r"] == cfg["r"]
    assert cfg2["schedule"] == cfg["schedule"]
    assert sorted(cfg2["target_modules"]) == sorted(cfg["target_modules"])
    assert cfg2["alpha"] == pytest.approx(cfg["alpha"], rel=1e-9)
    assert cfg2["dropout"] == pytest.approx(cfg["dropout"], rel=1e-9, abs=1e-12)
    assert cfg2["lr"] == pytest.approx(cfg["lr"], rel=1e-9)


def test_decode_clips_out_of_range() -> None:
    # Each coordinate is outside [0, 1]; verify it clips to the nearest bound
    # and decodes to the value at that bound.
    x = np.array([1.5, 2.0, -0.2, -1.0, 1.5, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    cfg = decode(x)
    assert cfg["r"] == 64
    assert cfg["alpha"] == pytest.approx(256.0)
    assert cfg["dropout"] == pytest.approx(0.0)
    assert cfg["lr"] == pytest.approx(1e-5)
    assert cfg["schedule"] == SCHEDULE_OPTIONS[-1]


def test_decode_empty_target_modules_falls_back() -> None:
    x = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    cfg = decode(x)
    assert cfg["target_modules"] == ["q", "v"]


def test_decode_wrong_shape_raises() -> None:
    with pytest.raises(ValueError):
        decode(np.zeros(10))


@pytest.mark.parametrize(
    "bad_field, value",
    [
        ("r", 3),
        ("schedule", "exotic"),
        ("alpha", 0.5),
        ("dropout", 0.6),
        ("lr", 5e-1),
    ],
)
def test_encode_rejects_bad_values(bad_field: str, value: object) -> None:
    cfg = dict(VALID_CONFIGS[1])
    cfg[bad_field] = value
    with pytest.raises(ValueError):
        encode(cfg)


def test_encode_rejects_unknown_target_modules() -> None:
    cfg = dict(VALID_CONFIGS[1])
    cfg["target_modules"] = ["q", "exotic"]
    with pytest.raises(ValueError):
        encode(cfg)


def test_all_r_options_round_trip() -> None:
    for r in R_OPTIONS:
        cfg = {
            "r": r,
            "alpha": 16.0,
            "dropout": 0.1,
            "lr": 3e-4,
            "schedule": "cosine",
            "target_modules": list(TARGET_MODULE_OPTIONS),
        }
        assert decode(encode(cfg))["r"] == r