""" LCD Planner Implementation """
from functools import partial
from typing import Any, Tuple, List, Dict, Union, Type, Optional, Callable

import time
import gym
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import wandb

from sb3_jax.common.offline_algorithm import OfflineAlgorithm
from sb3_jax.common.buffers import BaseBuffer
from sb3_jax.common.type_aliases import GymEnv, MaybeCallback, Schedule
from sb3_jax.common.jax_utils import jax_print, jit_optimize, stop_grad

from diffgro.utils.utils import print_b
from diffgro.common.buffers import TrajectoryBuffer
from diffgro.common.models.utils import apply_cond
from diffgro.lcd.policies import LCDPlannerPolicy


class LCDPlanner(OfflineAlgorithm):
    def __init__(
        self,
        policy: Union[str, Type[LCDPlannerPolicy]],
        env: Union[GymEnv, str],
        replay_buffer: Type[BaseBuffer] = TrajectoryBuffer,
        learning_rate: float = 3e-4,
        batch_size: int = 256, 
        gamma: float = 0.99,
        gradient_steps: int = 1,
        wandb_log: Optional[str] = None,
        policy_kwargs: Optional[Dict[str, Any]] = None,
        verbose: int = 0,
        seed: Optional[int] = 1,
        _init_setup_model: bool = True,
    ):
        super(LCDPlanner, self).__init__(
            policy,
            env,
            replay_buffer=replay_buffer,
            learning_rate=learning_rate,
            batch_size=batch_size,
            gamma=gamma,
            gradient_steps=gradient_steps,
            tensorboard_log=None,
            wandb_log=wandb_log,
            policy_kwargs=policy_kwargs,
            verbose=verbose,
            create_eval_env=False,
            seed=seed,
            _init_setup_model=False,
            supported_action_spaces=(gym.spaces.Box),
            support_multi_env=False,
        )
        self.learning_rate = learning_rate

        if _init_setup_model:
            self._setup_model()

    def _setup_model(self) -> None:
        self.set_random_seed(self.seed)
        self.policy = self.policy_class(
            self.observation_space,
            self.action_space,
            self.learning_rate,
            seed=self.seed,
            **self.policy_kwargs,
        )
        self._create_aliases()

    def _create_aliases(self) -> None:
        self.hact = self.policy.hact
        self.lact = self.policy.lact

    def train(self, gradient_steps: int, batch_size: int = 256) -> None:
        hpi_losses, lpi_losses = [], []
        
        for gradient_step in range(gradient_steps):
            self._n_updates += 1
            batch_keys = ['tasks', 'observations', 'actions', 'skills']
            replay_data = self.replay_buffer.sample(batch_keys, batch_size, max_length=self.hact.horizon) 
            task    = replay_data.tasks
            obs     = replay_data.observations
            p_obs   = self.policy.preprocess(obs.reshape(-1, self.hact.obs_dim), training=True)
            p_obs   = p_obs.reshape(batch_size, -1, self.hact.obs_dim)
            act     = replay_data.actions
            skill   = replay_data.skills

            # 1. learning high-level planner
            self.hact.optim_state, self.hact.params, hpi_loss, hpi_info = jit_optimize(
                self._hpi_loss,
                self.hact.optim,
                self.hact.optim_state,
                self.hact.params,
                max_grad_norm=None,
                obs=p_obs,
                task=task, 
                skill=skill,
                rng=next(self.policy.rng)
            )
            hpi_losses.append(hpi_loss)

            # 2. learning low-level executor
            self.lact.optim_state, self.lact.params, lpi_loss, lpi_info = jit_optimize(
                self._lpi_loss,
                self.lact.optim,
                self.lact.optim_state,
                self.lact.params,
                max_grad_norm=None,
                obs=p_obs,
                act=act,
                skill=skill,
                rng=next(self.policy.rng)
            )
            lpi_losses.append(lpi_loss)

        wandb_log = {"time/total_timesteps": self.num_timesteps}
        if self._n_updates % self.log_interval == 0:
            # reconstruct hpi mse
            pred_obs, _ = self.hact._denoise(hpi_info['cond'], task, skill, deterministic=True)
            hpi_mse = jnp.mean(jnp.square(pred_obs - p_obs))
            # reconstruct lpi mse
            obs_0, obs_1 = pred_obs[:,0,:], pred_obs[:,1,:]
            pred_act = self.lact._predict(obs_0, obs_1, skill)
            lpi_mse = jnp.mean(jnp.square(pred_act - act[:,0,:]))
            self.logger.record("train/hpi/recon_mse", hpi_mse)
            self.logger.record("train/lpi/recon_mse", lpi_mse)
            wandb_log.update({"train/hpi/recon_mse": hpi_mse})
            wandb_log.update({"train/lpi/recon_mse": lpi_mse})


        self.logger.record("train/batch_size", batch_size)
        self.logger.record("train/hpi/loss", np.mean(hpi_losses))
        self.logger.record("train/lpi/loss", np.mean(lpi_losses))
        wandb_log.update({"train/hpi/loss": np.mean(hpi_losses)})
        wandb_log.update({"train/lpi/loss": np.mean(lpi_losses)})
        if self.wandb_log is not None:
            wandb.log(wandb_log)

    @partial(jax.jit, static_argnums=0)
    def _hpi_loss(
        self,
        hpi_params: hk.Params,
        obs: jax.Array,
        task: jax.Array,
        skill: jax.Array,
        rng=None,
    ) -> Tuple[jax.Array, Dict[str, jax.Array]]:
        batch_size = obs.shape[0]
        rng_t, rng_n, rng_p, rng = jax.random.split(rng, num=4)

        # 1. forward process
        horizon = self.hact.horizon
        noise = jax.random.normal(rng_n, shape=(batch_size, horizon, self.hact.out_dim))
        ts = jax.random.randint(rng_t, shape=(batch_size, 1), minval=1, maxval=self.hact.n_denoise+1)
        sqrtab = jnp.repeat(self.hact.ddpm_dict.sqrtab[ts], horizon, axis=0).reshape(batch_size, -1, 1)
        sqrtmab = jnp.repeat(self.hact.ddpm_dict.sqrtmab[ts], horizon, axis=0).reshape(batch_size, -1, 1)
        x_t = sqrtab * obs + sqrtmab * noise
        
        # 2. conditioning with inpainting
        cond = obs[:,0,:]
        x_t = apply_cond(x_t, cond)

        # 3. diffusion inference
        batch_dict = {"lang": task, "skill": skill}
        noise_pred, info = self.hact._hpi(x_t, batch_dict, ts, False, False, hpi_params, rng_p)

        # 4. diffusion loss
        if self.hact.predict_epsilon:
            noise = apply_cond(noise, cond)
            loss = jnp.mean(jnp.square(noise_pred - nosie))
        else:
            loss = jnp.mean(jnp.square(noise_pred - obs))
        return loss, {'cond': cond}

    @partial(jax.jit, static_argnums=0)
    def _lpi_loss(
        self, 
        lpi_params: hk.Params,
        obs: jax.Array,
        act: jax.Array,
        skill: jax.Array,
        rng=None,
    ) -> Tuple[jax.Array, Dict[str, jax.Array]]:
        obs_0 = obs[:,:-1,:].reshape(-1, self.lact.obs_dim) # 0 ~ horizon-1
        obs_1 = obs[:,1:,:].reshape(-1, self.lact.obs_dim)  # 1 ~ horizon
        act = act[:,:-1,:].reshape(-1, self.lact.act_dim)
        skill = np.repeat(skill, self.hact.horizon-1, axis=0)

        batch_dict = {"obs_0": obs_0, "obs_1": obs_1, "skill": skill}
        pred_act = self.lact._lpi(batch_dict, lpi_params, rng)

        loss = jnp.mean(jnp.square(pred_act - act))
        return loss, {}

    def learn(
        self,
        total_timesteps: Tuple[int, int],
        callback: MaybeCallback = None,
        log_interval: int = 1,
        eval_env: Optional[GymEnv] = None,
        eval_freq: int = -1,
        n_eval_episodes: int = 5,
        tb_log_name: str = "LCDPlanner",
        eval_log_path: Optional[str] = None,
        reset_num_timesteps: bool = True,
    ) -> "LCDPlanner":
        self.log_interval = log_interval

        # wandb configs
        if self.wandb_log is not None:
            self.wandb_config = dict(
                time=time.ctime(),
                algo='lcd/planner',
                tag=self.wandb_log['tag'],
                learning_rate=self.learning_rate,
                batch_size=self.batch_size,
                gamma=self.gamma,
                gradient_steps=self.gradient_steps,
                seed=self.seed,
            )
            self.wandb_config.update(self.policy._get_constructor_parameters())
        
        total_timesteps, callback = self._setup_learn(
            total_timesteps, eval_env, callback, eval_freq, n_eval_episodes, eval_log_path, reset_num_timesteps, tb_log_name
        )
        callback.on_training_start(locals(), globals())

        # 2. learn policy module
        start_time = time.time()
        num_timesteps = 0
        while num_timesteps < total_timesteps:
            self.train(gradient_steps=self.gradient_steps, batch_size=self.batch_size)

            self.num_timesteps += 1
            num_timesteps += 1
            if log_interval is not None and num_timesteps % log_interval == 0:
                fps = int(num_timesteps / (time.time() - start_time))
                self.logger.record("time/fps", fps)
                self.logger.record("time/time_elapsed", int(time.time() - start_time), exclude="tensorboard")
                self.logger.record("time/total_timesteps", num_timesteps, exclude="tensorboard")
                self.logger.dump(step=num_timesteps)
            
            callback.update_locals(locals())
            if callback.on_step() is False:
                return False
        
        callback.on_training_end()
        return self

    def load_params(self, path: str) -> None:
        print_b(f"[lcd] : loading params")
        data, params = load_from_zip_file(path, verbose=1)
        self._load_jax_params(params)
        self._load_norm_layer(path)

    def _save_jax_params(self) -> Dict[str, hk.Params]:
        params_dict = {} 
        params_dict["hpi_params"] = self.hact.params
        params_dict["lpi_params"] = self.lact.params
        return params_dict
    
    def _load_jax_params(self, params: Dict[str, hk.Params]) -> None:
        self.hact._load_jax_params(params)
        self.lact._load_jax_params(params)

    def _save_norm_layer(self, path: str) -> None:
        if self.policy.normalization_class is not None:
            self.policy.normalization_layer.save(path)

    def _load_norm_layer(self, path: str) -> None:
        if self.policy.normalization_class is not None:
            self.policy.normalization_layer = self.policy.normalization_layer.load(path)
