import jax
import jax.numpy as jnp
from brittle_star_project.environment.env_config import MorphMode
from brittle_star_project.environment.padded_obs_wrapper import (
    compute_padding_masks,
)
from brittle_star_project.environment.obs_processing import create_obs_processor

# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py
from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP


def test_centralized_forward_pass_with_padding():
    batch_size = 2

    # 1. Simulate Amputated Observation [4, 0, 4, 2, 4] -> 14 segments total
    # 14 segments * 2 = 28 joints
    amputated_obs = {
        "joint_position": jnp.zeros((batch_size, 28)),
        "joint_velocity": jnp.zeros((batch_size, 28)),
        "segment_contact": jnp.zeros((batch_size, 14)),
    }

    segments_per_arm = jnp.array((4, 0, 4, 2, 4))
    num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item()

    # 2. Process and Pad Observation
    masks = compute_padding_masks(segments_per_arm=list(segments_per_arm))
    obs_processor = create_obs_processor(
        bounds_dict={},
        needed_copies=1,
        num_arms=num_arms,
        padding_masks=masks,
        morph_mode=MorphMode.CENTRALIZED,
        segments_per_arm=segments_per_arm,
    )
    global_state = obs_processor(amputated_obs)

    # joint_position: 5 arms × 8 joints (padded) = 40
    # joint_velocity:  5 arms × 8 joints (padded) = 40
    # segment_contact: 5 arms × 4 segs  (padded) = 20
    # Total = 100 (no disk or direction keys supplied)
    assert global_state.shape == (batch_size, 1, 100), (
        f"Expected global state shape (2, 1, 100), got {global_state.shape}"
    )

    actor = Actor(action_dim=40)
    critic = OneDenseLayerMLP()  # Acts as the centralized critic

    rng = jax.random.PRNGKey(0)
    rng_a, rng_c = jax.random.split(rng)

    # Initialize Flax variables
    actor_params = actor.init(rng_a, global_state)
    critic_params = critic.init(rng_c, global_state)

    # 5. Forward Pass Assertions
    action_mean, action_log_std = actor.apply(actor_params, global_state)
    value = critic.apply(critic_params, global_state)

    assert action_mean.shape == (batch_size, 1, 40), (
        f"Actor mean shape mismatch: {action_mean.shape}"
    )
    assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}"
    assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), (
        f"Critic value shape mismatch: {value.shape}"
    )


if __name__ == "__main__":
    test_centralized_forward_pass_with_padding()