""" DiffBC 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

from diffgro.utils.utils import print_b
from diffgro.common.buffers import TrajectoryBuffer
from diffgro.common.models.utils import fn_foward_noise
from diffgro.diffbc.policies import DiffBCPlannerPolicy


class DiffBCPlanner(OfflineAlgorithm):
    def __init__(
        self,
        policy: Union[str, Type[DiffBCPlannerPolicy]],
        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(DiffBCPlanner, 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.act = self.policy.act

    def train(self, gradient_steps: int, batch_size: int = 256) -> None:
        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=1) 
            task    = replay_data.tasks
            obs     = replay_data.observations[:,0,:]
            p_obs   = self.policy.preprocess(obs, training=True)
            act     = replay_data.actions[:,0,:]

            self.act.optim_state, self.act.params, loss, info = jit_optimize(
                self._loss,
                self.act.optim,
                self.act.optim_state,
                self.act.params,
                max_grad_norm =None,
                task=task,
                obs=p_obs,
                act=act,
                rng=next(self.policy.rng)
            )
            losses.append(loss)

        wandb_log = {"time/total_timesteps": self.num_timesteps}
        if self._n_updates % self.log_interval == 0:
            pred_act, _ = self.policy._predict(task, obs, deterministic=True)
            action_mse = jnp.mean(jnp.square(act - pred_act))
            self.logger.record("train/pi/mse", action_mse)
            wandb_log.update({"train/pi_mse": action_mse})

        self.logger.record("train/pi/loss", np.mean(losses))
        wandb_log.update({"train/pi/loss": np.mean(losses)})
        if self.wandb_log is not None:
            wandb.log(wandb_log)
            
    @partial(jax.jit, static_argnums=0)
    def _loss(
        self,
        pi_params: hk.Params,
        task: jax.Array,
        obs: jax.Array,
        act: 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)
        
        # input
        ts = jax.random.randint(rng_t, (batch_size, 1), minval=1, maxval=self.act.n_denoise+1)
        x_t, noise = fn_foward_noise(act, self.act.ddpm_dict.sqrtab[ts], self.act.ddpm_dict.sqrtmab[ts], rng_n)

        # inference
        batch_dict = {"task": task, "obs": obs}
        noise_pred, _ = self.act._pi(x_t, batch_dict, ts, False, False, pi_params, rng_p)

        # loss 
        if self.act.predict_epsilon: # prediction of noise
            loss = jnp.mean(jnp.square(noise_pred - noise))
        else: # prediction of original action
            loss = jnp.mean(jnp.square(noise_pred - 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 = "DiffBC",
        eval_log_path: Optional[str] = None,
        reset_num_timesteps: bool = True,
    ) -> "DiffBC":
        self.log_interval = log_interval

        # wandb configs
        if self.wandb_log is not None:
            self.wandb_config = dict(
                time=time.ctime(),
                algo='diffbc',
                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"[diffbc] : 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["pi_params"] = self.act.params
        return params_dict
    
    def _load_jax_params(self, params: Dict[str, hk.Params]) -> None:
        self.act._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)
