from typing import Any
import flax.linen as nn
import numpy as np
import optax
import jax
from flax.training.train_state import TrainState
import jax.numpy as jnp
from gymnasium import spaces
from stable_baselines3.common.env_util import make_vec_env

from sbx import PPO
from sbx.ppo.policies import PPOPolicy, Actor, Critic


class CustomPPO(PPOPolicy):
    def __init__(self, observation_space, action_space, lr_schedule, **kwargs):
        super().__init__(observation_space, action_space, lr_schedule, **kwargs)

    def build(self, key: jax.Array, lr_schedule, max_grad_norm) -> jax.Array:
        key, actor_key, vf_key = jax.random.split(key, 3)
        key, self.key = jax.random.split(key, 2)
        self.reset_noise()
        obs = jnp.array([self.observation_space.sample()])

        self.actor = self.actor_class(
            net_arch=self.net_arch_pi,
            log_std_init=self.log_std_init,
            activation_fn=self.activation_fn,
            ortho_init=self.ortho_init,
            action_dim=int(np.prod(self.action_space.shape)),
        )

        optimizer_class = optax.inject_hyperparams(self.optimizer_class)(
            learning_rate=lr_schedule(1), **self.optimizer_kwargs
        )

        # Optax chain with only one element
        self.actor_state = TrainState.create(
            apply_fn=self.actor.apply,
            params=self.actor.init(actor_key, obs),
            tx=optax.chain(
                optimizer_class,
            ),
        )

        self.vf = self.critic_class(net_arch=self.net_arch_vf, activation_fn=self.activation_fn)
        self.vf_state = TrainState.create(
            apply_fn=self.vf.apply,
            params=self.vf.init(vf_key, obs),
            tx=optax.chain(
                optimizer_class,
            ),
        )

        self.actor.apply = jax.jit(self.actor.apply)  # type: ignore
        self.vf.apply = jax.jit(self.vf.apply)        # type: ignore
        return key


def test_ppo() -> None:
    env = make_vec_env("Pendulum-v1")
    model = PPO(CustomPPO, env)
    model.learn(64, progress_bar=True)