import torch
import torch.nn as nn
import torch.optim as optim
import robomimic.utils.tensor_utils as TensorUtils
from copy import deepcopy

# do not import any additional diffusion model packages for this part of the assignment
# you are free to implement any packages you want for the experiments at the end of the assignment
from diffusers import DDPMScheduler
from film.util.helperModels import TransformerForDiffusion
from diffusers.training_utils import EMAModel


class DLM_Transformer(nn.Module):
    def __init__(
        self,
        output_dim,
        input_dim,
        denoising_steps,
        obs_horizon,
        prediction_horizon,
        action_horizon,
    ):

        super(DLM_Transformer, self).__init__()
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu")
        self.denoising_steps = denoising_steps
        self.output_dim = output_dim
        self.obs_horizon = obs_horizon
        self.prediction_horizon = prediction_horizon
        self.action_horizon = action_horizon
        self.input_dim = input_dim

        ######################################
        # Initialize the noise pred net      #
        # Change the input and cond dims     #
        ######################################
        # Input_dim is action dim
        # global con dim is flattened obs dim
        self.noise_pred_net = TransformerForDiffusion(
            input_dim=output_dim,
            output_dim=output_dim,
            horizon=prediction_horizon,
            n_obs_steps=obs_horizon,
            cond_dim=int(input_dim),
        )
        ######################################
        # End of code                        #
        ######################################

        self.scheduler = DDPMScheduler(
            num_train_timesteps=self.denoising_steps,
            # the choise of beta schedule has big impact on performance
            # we found squared cosine works the best
            beta_schedule='squaredcos_cap_v2',
            # clip output to [-1,1] to improve stability
            clip_sample=True,
            # our network predicts noise (instead of denoised action)
            prediction_type='epsilon'
        )

        self.noise_pred_net.to(self.device)
        self.optimizer = optim.Adam(
            self.noise_pred_net.parameters(), lr=1e-4, weight_decay=1e-3)
        self.ema = EMAModel(
            parameters=self.noise_pred_net.parameters(), power=0.75)
        self.ema.to(self.device)
        # This is just for saving the model, don't use this parameter
        self.epoch = 0

    def train_on_batch(self, batch, validate):
        """
        Args:
        - batch (dict): batch of data batch['obs'].shape = (B, T, Obs_dim) and batch['actions'].shape = (B, T, Action_dim)
        - validate (bool): whether batch is for validation
        Returns:
        - loss_val (float): value of loss
        """
        loss_val = None
        self.noise_pred_net.train()
        ######################################
        # Define the training loop here      #
        ######################################
        # make sure to flatten obs (B, T, D) -> (B, --)
        obs = batch['obs']
        B = obs.shape[0]

        start = self.obs_horizon-1
        end = start + self.action_horizon
        act = batch['actions'][:, start:end]

        # Get the noise and cond
        t = torch.randint(0, self.denoising_steps, (B,)).to(self.device)
        noise = torch.randn(act.shape).to(self.device)
        cond = obs[:, :self.obs_horizon, :]
        cond_mask = torch.zeros_like(noise, dtype=torch.bool)

        if validate:
            with torch.no_grad():
                # Add noise to action
                sample = self.scheduler.add_noise(act, noise, t)
                # sample[cond_mask] = act[cond_mask]

                pred = self.noise_pred_net(sample, t, cond).to(self.device)

                loss = torch.nn.MSELoss()(noise, pred)
        else:
            # Add noise to action
            sample = self.scheduler.add_noise(act, noise, t)
            pred = self.noise_pred_net(sample, t, cond).to(self.device)
            loss = torch.nn.MSELoss()(noise, pred)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

        loss_val = loss.item()

        ######################################
        # End of code                        #
        ######################################
        return loss_val

    def get_action(self, obs):
        """
        Args:
        - obs (dict): dictionary of observations of shape (B, prediction_horizon, Obs_dim)
        Returns:
        - actions (np.ndarray): action of shape (B, action_horizon, Action_dim)
        """
        # obs are (B, T, D) - make sure to flatten obs
        # also, remember obs_horizon is not the same as prediction_horizon
        # original_params = deepcopy(self.noise_pred_net.state_dict())
        # self.ema.copy_to(self.noise_pred_net.parameters())
        self.noise_pred_net.eval()

        ######################################################
        # Generate an action sample from the diffusion model #
        ######################################################
        self.scheduler.set_timesteps(self.denoising_steps)
        B = obs.shape[0]

        with torch.no_grad():
            input_shape = (B, self.prediction_horizon, self.output_dim)
            # cond_data = torch.zeros(size=input_shape, device=self.device)
            # cond_mask = torch.zeros_like(cond_data, dtype=torch.bool)
            noise = torch.randn(input_shape).to(self.device)

            cond = obs
            for t in self.scheduler.timesteps:
                # noise[cond_mask] = cond_data[cond_mask]
                pred = self.noise_pred_net(noise, t, cond)
                traj = self.scheduler.step(pred, t, noise)
                noise = traj.prev_sample
            # noise[cond_mask] = cond_data[cond_mask]

            start = self.obs_horizon-1
            end = start + self.action_horizon
            actions = noise[:, start:end]
        ####################################################################################
        #                                 END OF YOUR CODE                                 #
        ####################################################################################

        # restore original params
        # self.noise_pred_net.load_state_dict(original_params)

        return actions

    def reset(self):
        return

    def save(self, path):
        self.ema.copy_to(self.noise_pred_net.parameters())
        torch.save(self.state_dict(), path)
        # save the ema model, put ema before the .pth
        torch.save(self.ema.state_dict(), path.replace(".pth", "ema.pth"))

    def load(self, path):
        # get epoch from end of path
        self.load_state_dict(torch.load(path))
        self.ema.load_state_dict(torch.load(path.replace(".pth", "ema.pth")))

    def set_horizon(self, horizon):
        self.prediction_horizon = horizon
        self.action_horizon = horizon

    def set_eval(self):
        self.eval()
