import time
from dataclasses import asdict, dataclass
from functools import partial
from typing import Any

import optax
import tqdm
from flax.metrics.tensorboard import SummaryWriter
from flax.training.train_state import TrainState

from MLPs.mlps import (
    GenericDenseLayersWithActivation,
    Actor,
    OneDenseLayerMLP,
    AgentParams,
    Storage,
)
from brittle_star_project.dataclasses import PPOArgs, EpisodeStatistics
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
import jax
import jax.numpy as jnp
import numpy as np

from ppo import PPO


@jax.jit
def linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
    frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
    return learning_rate * frac


@jax.jit
def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
        obs_dict
    )


@jax.jit
def get_action_and_value_noise(
    sensor: GenericDenseLayersWithActivation,
    feature_extractor: GenericDenseLayersWithActivation,
    actor: Actor,
    critic: OneDenseLayerMLP,
    agent_state: TrainState,
    next_obs: jnp.ndarray,
    key: jax.random.PRNGKey,
):
    hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
    hidden_critic = feature_extractor.apply(
        agent_state.params["feature_extractor_params"], next_obs
    )

    # Continuous actions: sample from a Gaussian parameterized by the actor
    mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
    key, subkey = jax.random.split(key)
    noise = jax.random.normal(subkey, shape=mean.shape)
    std = jnp.exp(log_std)
    action = mean + noise * std
    logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
    value = critic.apply(agent_state.params["critic_params"], hidden_critic)
    return action, logprob, value.squeeze(-1), key


@jax.jit
def _step_once(
    carry,
    _,
    env_step_fn,
    sensor: GenericDenseLayersWithActivation,
    feature_extractor: GenericDenseLayersWithActivation,
    actor: Actor,
    critic: OneDenseLayerMLP,
):
    agent_state, episode_stats, obs, done, key, env_state = carry
    action, logprob, value, key = get_action_and_value_noise(
        sensor, feature_extractor, actor, critic, agent_state, obs, key
    )

    episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
        episode_stats, env_state, action
    )

    storage = Storage(
        obs=obs,
        actions=action,
        logprobs=logprob,
        dones=done,
        values=value,
        rewards=reward,
        returns=jnp.zeros_like(reward),
        advantages=jnp.zeros_like(reward),
    )
    return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage


@jax.jit
def _rollout_jit(
    agent_state,
    episode_stats,
    env_state,
    next_obs,
    next_done,
    key,
    args,
    env_step_fn,
    sensor: GenericDenseLayersWithActivation,
    feature_extractor: GenericDenseLayersWithActivation,
    actor: Actor,
    critic: OneDenseLayerMLP,
):
    (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
        partial(
            _step_once,
            sensor=sensor,
            feature_extractor=feature_extractor,
            actor=actor,
            critic=critic,
            env_step_fn=partial(_step_env_wrapped, env_step_fn=env_step_fn),
        ),
        (agent_state, episode_stats, next_obs, next_done, key, env_state),
        (),
        args.max_steps,
    )
    return agent_state, episode_stats, next_obs, next_done, storage, key, env_state


@jax.jit
def _step_env_wrapped(env_step_fn, env_state, action, episode_stats):
    next_env_state = env_step_fn(env_state, action)

    # Extract per-environment signals from the state object
    reward = next_env_state.reward  # (num_envs,)
    terminated = next_env_state.terminated  # (num_envs,)
    truncated = next_env_state.truncated  # (num_envs,)
    done = terminated | truncated  # (num_envs,)

    new_episode_return = episode_stats.episode_returns + reward
    new_episode_length = episode_stats.episode_lengths + 1

    episode_stats = episode_stats.replace(
        episode_returns=new_episode_return * (1 - done),
        episode_lengths=new_episode_length * (1 - done),
        returned_episode_returns=jnp.where(
            done, new_episode_return, episode_stats.returned_episode_returns
        ),
        returned_episode_lengths=jnp.where(
            done, new_episode_length, episode_stats.returned_episode_lengths
        ),
    )
    return (
        episode_stats,
        next_env_state,
        (convert_obs_dict_to_array(next_env_state.observations), reward, done),
    )


@jax.jit
def compute_gae_once(carry, inp, gamma, gae_lambda):
    advantages = carry
    nextdone, nextvalues, curvalues, reward = inp
    nextnonterminal = 1.0 - nextdone
    delta = reward + gamma * nextvalues * nextnonterminal - curvalues
    advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
    return advantages, advantages


@jax.jit
def compute_gae_jit(agent_state, storage, next_obs, next_done, sensor, critic, args):
    next_value = critic.apply(
        agent_state.params["critic_params"],
        sensor.apply(agent_state.params["sensor_params"], next_obs),
    ).squeeze(-1)

    advantages = jnp.zeros((args.num_envs,))
    dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
    values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
    _, advantages = jax.lax.scan(
        partial(compute_gae_once, gamma=args.gamma, gae_lambda=args.gae_lambda),
        advantages,
        (dones[1:], values[1:], values[:-1], storage.rewards),
        reverse=True,
    )
    return storage.replace(advantages=advantages, returns=advantages + storage.values)


@dataclass
class LossInfo:
    # todo: better typing
    loss: Any
    pg_loss: Any
    v_loss: Any
    entropy_loss: Any
    approx_kl: Any
    avg_episodic_return: Any


class PPOTrainer:
    def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_name: str):
        self.args = args
        self.env = env
        self.writer = SummaryWriter(f"runs/{run_name}")

        self.key = jax.random.PRNGKey(args.seed)

        self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
        self.sensor.apply = jax.jit(self.sensor.apply)
        self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
        self.actor.apply = jax.jit(self.actor.apply)
        self.critic.apply = jax.jit(self.critic.apply)

        self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)

        self.agent_state = self._init_agent_state()

        self.episode_stats = self._init_episode_stats()

    def _init_agent(self):
        sensor = GenericDenseLayersWithActivation()
        feature_extractor = GenericDenseLayersWithActivation()
        actor = Actor(
            action_dim=self.env.single_action_space.shape[0]
        )  # continuous actions for MJX
        critic = OneDenseLayerMLP()
        # messenger = OneDenseLayerMLP()
        return sensor, feature_extractor, actor, critic

    def _init_agent_state(self) -> TrainState:
        self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
            self.key, 5
        )

        sample_obs = jnp.concatenate(
            [
                v.flatten()
                for v in self.env.single_observation_space.sample(
                    rng=jax.random.PRNGKey(0)
                ).values()
                if v.size > 0
            ]
        )
        sensor_params = self.sensor.init(sensor_key, sample_obs)
        feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
        actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
        critic_params = self.critic.init(
            critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
        )

        return TrainState.create(
            apply_fn=None,
            params=asdict(
                AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
            ),
            tx=optax.chain(
                optax.clip_by_global_norm(self.args.max_grad_norm),
                optax.inject_hyperparams(optax.adam)(
                    learning_rate=linear_schedule
                    if self.args.anneal_lr
                    else self.args.learning_rate,
                    eps=1e-5,
                ),
            ),
        )

    def _init_episode_stats(self) -> EpisodeStatistics:
        return EpisodeStatistics(
            episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
            episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
            returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
            returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
        )

    def _rollout(self, env_state, next_obs, next_done) -> tuple[Storage, ...]:
        return _rollout_jit(
            self.agent_state,
            self.episode_stats,
            env_state,
            next_obs,
            next_done,
            self.key,
            self.args,
            self.env.step,
            self.sensor,
            self.feature_extractor,
            self.actor,
            self.critic,
        )

    def _compute_gae(self, storage, next_obs, next_done) -> Storage:
        return compute_gae_jit(
            self.agent_state, storage, next_obs, next_done, self.sensor, self.critic, self.args
        )

    def _log(
        self,
        global_step,
        episode_stats,
        avg_episodic_return,
        start_time,
        iteration_time_start,
        loss_info,
    ):
        self.writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
        self.writer.add_scalar(
            "charts/avg_episodic_length",
            np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
            global_step,
        )
        self.writer.add_scalar(
            "charts/learning_rate",
            self.agent_state.opt_state[1].hyperparams["learning_rate"].item(),
            global_step,
        )
        self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step)
        self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step)
        self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step)
        self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step)
        self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step)

        # iters_bar.set_postfix_str(f"SPS: {int(global_step / (time.time() - start_time))}")

        self.writer.add_scalar(
            "charts/SPS", int(global_step / (time.time() - start_time)), global_step
        )
        self.writer.add_scalar(
            "charts/SPS_update",
            int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)),
            global_step,
        )

    def _step(self, env_state, next_obs, next_done) -> tuple:
        storage, next_obs, next_done, env_state = self._rollout(env_state, next_obs, next_done)
        storage = self._compute_gae(storage, next_obs, next_done)
        self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = (
            self._ppo.update_ppo(self.agent_state, storage, self.key)
        )

        avg_episodic_return = jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns))

        return (
            env_state,
            next_obs,
            next_done,
            LossInfo(
                loss=loss,
                pg_loss=pg_loss,
                v_loss=v_loss,
                entropy_loss=entropy_loss,
                approx_kl=approx_kl,
                avg_episodic_return=avg_episodic_return,
            ),
        )

    def close(self):
        self.env.close()
        self.writer.close()

    def train(self):
        """
        Train the PPO agent for a specified number of iterations
        (passed through PPOArgs in constructor).
        Closes the environment at the end of training.
        """
        env_state = self.env.reset(seed=self.args.seed)
        next_obs = convert_obs_dict_to_array(env_state.observations)
        next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
        global_step = 0
        start_time = time.time()

        for _ in tqdm.tqdm(range(self.args.num_iterations)):
            iteration_time_start = time.time()

            env_state, next_obs, next_done, loss_info = self._step(env_state, next_obs, next_done)
            global_step += self.args.num_steps * self.args.num_envs
            self._log(
                global_step, self.episode_stats, 0, start_time, iteration_time_start, loss_info
            )

        if self.args.save_model:
            self._save_model(...)

        self.close()