from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")

from dovla_cil.data.schema import ActionChunk
from dovla_cil.models import devectorize_toy_action, vectorize_toy_action
from dovla_cil.models.dovla import DoVLAConfig, DoVLAModel, load_model_state


def _model_and_inputs():
    config = DoVLAConfig(
        obs_dim=10,
        lang_dim=16,
        action_dim=8,
        hidden_dim=32,
        action_horizon=3,
        effect_dim=7,
        intervention_dim=24,
    )
    model = DoVLAModel(config)
    observation = torch.randn(4, config.obs_dim)
    instructions = [
        "pick the red mug",
        "put the mug in the bowl",
        "open the drawer",
        "push the cube",
    ]
    action = torch.randn(4, config.action_horizon, config.action_dim)
    return model, config, observation, instructions, action


def test_forward_policy_works_and_shapes() -> None:
    model, config, observation, instructions, _action = _model_and_inputs()
    predicted = model.forward_policy(observation, instructions)
    assert predicted.shape == (4, config.action_horizon, config.action_dim)


def test_forward_effect_works_and_shapes() -> None:
    model, config, observation, instructions, action = _model_and_inputs()
    output = model.forward_effect(observation, instructions, action)
    assert output["effect_vector"].shape == (4, config.effect_dim)
    assert output["success_logit"].shape == (4,)
    assert output["success"].shape == (4,)
    assert output["progress"].shape == (4,)


def test_forward_reward_works_and_shapes() -> None:
    model, _config, observation, instructions, action = _model_and_inputs()
    reward = model.forward_reward(observation, instructions, action)
    regret = model.forward_regret(observation, instructions, action)
    z = model.encode_intervention(observation, instructions, action)
    assert reward.shape == (4,)
    assert regret.shape == (4,)
    assert z.shape == (4, model.config.intervention_dim)


def test_gradients_flow() -> None:
    model, _config, observation, instructions, action = _model_and_inputs()
    policy = model.forward_policy(observation, instructions)
    effect = model.forward_effect(observation, instructions, action)
    reward = model.forward_reward(observation, instructions, action)
    loss = policy.square().mean() + effect["effect_vector"].square().mean() + reward.mean()
    loss.backward()
    grads = [
        parameter.grad
        for parameter in model.parameters()
        if parameter.requires_grad and parameter.grad is not None
    ]
    assert grads
    assert sum(float(grad.abs().sum()) for grad in grads) > 0.0


def test_rgb_observation_encoder_drives_policy_and_field_gradients() -> None:
    config = DoVLAConfig(
        obs_dim=10,
        lang_dim=16,
        action_dim=7,
        hidden_dim=32,
        action_horizon=2,
        effect_dim=6,
        intervention_dim=24,
        observation_mode="rgb",
    )
    model = DoVLAModel(config)
    images = torch.randint(0, 256, (3, 32, 40, 3), dtype=torch.uint8)
    instructions = ["push the cube", "stack the cubes", "lift the peg"]
    actions = torch.randn(3, config.action_horizon, config.action_dim)

    policy = model.forward_policy(images, instructions)
    field = model.forward_field(images, instructions, actions)
    (policy.square().mean() + field["potential"].mean()).backward()

    assert policy.shape == (3, 2, 7)
    assert field["effect_vector"].shape == (3, 6)
    assert any(
        parameter.grad is not None and float(parameter.grad.abs().sum()) > 0
        for parameter in model.observation_encoder.image_net.parameters()
    )


def test_toy_action_vectorize_and_devectorize() -> None:
    action = ActionChunk(
        representation="semantic",
        horizon=2,
        values=[
            {"command": "move_to", "object": "red_mug"},
            {"command": "push", "object": "red_mug", "dx": 0.1, "dy": -0.2},
        ],
        skill_type="push",
    )
    matrix = vectorize_toy_action(action, action_dim=8, action_horizon=3)
    restored = devectorize_toy_action(matrix, skill_type="push")
    assert len(matrix) == 3
    assert len(matrix[0]) == 8
    assert restored.representation == "semantic"
    assert restored.skill_type == "push"


def test_numeric_action_chunk_vectorization_preserves_simulator_controls() -> None:
    action = ActionChunk(
        representation="maniskill_pd_ee_delta_pose",
        horizon=2,
        values=[
            [0.1, -0.2, 0.3, 0.4, -0.5, 0.6, 0.7],
            [0.0, 0.1, -0.1, 0.2, -0.2, 0.3, -0.3],
        ],
        skill_type="pick_cube",
    )

    matrix = vectorize_toy_action(action, action_dim=8, action_horizon=3)

    assert matrix[0] == [0.1, -0.2, 0.3, 0.4, -0.5, 0.6, 0.7, 0.0]
    assert matrix[1] == [0.0, 0.1, -0.1, 0.2, -0.2, 0.3, -0.3, 0.0]
    assert matrix[2] == [0.0] * 8


def test_model_state_loader_allows_only_declared_omitted_prefixes() -> None:
    model, config, _observation, _instructions, _action = _model_and_inputs()
    state = model.state_dict()
    prefix = "observation_encoder.net.0."
    compact = {key: value for key, value in state.items() if not key.startswith(prefix)}
    restored = DoVLAModel(config)

    load_model_state(
        restored,
        {
            "model_state_dict": compact,
            "omitted_state_prefixes": [prefix],
        },
    )

    invalid = dict(compact)
    invalid.pop("language_encoder.net.0.weight")
    with pytest.raises(RuntimeError, match="invalid_missing"):
        load_model_state(
            DoVLAModel(config),
            {
                "model_state_dict": invalid,
                "omitted_state_prefixes": [prefix],
            },
        )