from typing import Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import reduce
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler

import diffusion_policy.model.vision.crop_randomizer as dmvc
import robomimic.models.base_nets as rmbn
import robomimic.utils.obs_utils as ObsUtils
from diffusion_policy.common.pytorch_util import dict_apply, replace_submodules
from diffusion_policy.common.robomimic_config_util import get_robomimic_config
from diffusion_policy.model.common.normalizer import LinearNormalizer
from diffusion_policy.model.diffusion.mask_generator import LowdimMaskGenerator
from diffusion_policy.model.diffusion.pmf_transformer_for_diffusion import (
    PMFTransformerForDiffusion,
)
from diffusion_policy.policy.base_image_policy import BaseImagePolicy
from robomimic.algo import algo_factory
from robomimic.algo.algo import PolicyAlgo


class PMFTransformerHybridImagePolicy(BaseImagePolicy):
    def __init__(
        self,
        shape_meta: dict,
        noise_scheduler: DDPMScheduler,
        horizon,
        n_action_steps,
        n_obs_steps,
        num_inference_steps=None,
        crop_shape=(76, 76),
        obs_encoder_group_norm=False,
        eval_fixed_crop=False,
        n_layer=8,
        n_cond_layers=0,
        n_head=4,
        n_emb=256,
        p_drop_emb=0.0,
        p_drop_attn=0.0,
        causal_attn=True,
        obs_as_cond=True,
        pred_action_steps_only=False,
        n_time_tokens=4,
        n_head_layers=4,
        min_time=0.05,
        du_dt_epsilon=1.0e-3,
        pmf_u_loss_weight=1.0,
        pmf_v_loss_weight=1.0,
        noise_scale=1.0,
        adatloss_eps=0.01,
        p_mean=-0.4,
        p_std=1.0,
        tr_uniform=True,
        tr_uniform_prob=0.1,
        data_proportion=0.5,
        **kwargs,
    ):
        super().__init__()

        action_shape = shape_meta["action"]["shape"]
        assert len(action_shape) == 1
        action_dim = action_shape[0]
        obs_shape_meta = shape_meta["obs"]
        obs_config = {
            "low_dim": [],
            "rgb": [],
            "depth": [],
            "scan": [],
        }
        obs_key_shapes = dict()
        for key, attr in obs_shape_meta.items():
            shape = attr["shape"]
            obs_key_shapes[key] = list(shape)

            obs_type = attr.get("type", "low_dim")
            if obs_type == "rgb":
                obs_config["rgb"].append(key)
            elif obs_type == "low_dim":
                obs_config["low_dim"].append(key)
            else:
                raise RuntimeError(f"Unsupported obs type: {obs_type}")

        config = get_robomimic_config(
            algo_name="bc_rnn",
            hdf5_type="image",
            task_name="square",
            dataset_type="ph",
        )

        with config.unlocked():
            config.observation.modalities.obs = obs_config

            if crop_shape is None:
                for _, modality in config.observation.encoder.items():
                    if modality.obs_randomizer_class == "CropRandomizer":
                        modality["obs_randomizer_class"] = None
            else:
                crop_h, crop_w = crop_shape
                for _, modality in config.observation.encoder.items():
                    if modality.obs_randomizer_class == "CropRandomizer":
                        modality.obs_randomizer_kwargs.crop_height = crop_h
                        modality.obs_randomizer_kwargs.crop_width = crop_w

        ObsUtils.initialize_obs_utils_with_config(config)

        policy: PolicyAlgo = algo_factory(
            algo_name=config.algo_name,
            config=config,
            obs_key_shapes=obs_key_shapes,
            ac_dim=action_dim,
            device="cpu",
        )

        obs_encoder = policy.nets["policy"].nets["encoder"].nets["obs"]
        if obs_encoder_group_norm:
            replace_submodules(
                root_module=obs_encoder,
                predicate=lambda x: isinstance(x, nn.BatchNorm2d),
                func=lambda x: nn.GroupNorm(
                    num_groups=x.num_features // 16,
                    num_channels=x.num_features,
                ),
            )

        if eval_fixed_crop:
            replace_submodules(
                root_module=obs_encoder,
                predicate=lambda x: isinstance(x, rmbn.CropRandomizer),
                func=lambda x: dmvc.CropRandomizer(
                    input_shape=x.input_shape,
                    crop_height=x.crop_height,
                    crop_width=x.crop_width,
                    num_crops=x.num_crops,
                    pos_enc=x.pos_enc,
                ),
            )

        obs_feature_dim = obs_encoder.output_shape()[0]
        input_dim = action_dim if obs_as_cond else (obs_feature_dim + action_dim)
        cond_dim = obs_feature_dim if obs_as_cond else 0

        self.obs_encoder = obs_encoder
        self.model = PMFTransformerForDiffusion(
            input_dim=input_dim,
            output_dim=input_dim,
            horizon=horizon if not pred_action_steps_only else n_action_steps,
            n_obs_steps=n_obs_steps,
            cond_dim=cond_dim,
            n_layer=n_layer,
            n_head=n_head,
            n_emb=n_emb,
            p_drop_emb=p_drop_emb,
            p_drop_attn=p_drop_attn,
            causal_attn=causal_attn,
            obs_as_cond=obs_as_cond,
            n_cond_layers=n_cond_layers,
            n_time_tokens=n_time_tokens,
            n_head_layers=n_head_layers,
        )
        self.noise_scheduler = noise_scheduler
        self.mask_generator = LowdimMaskGenerator(
            action_dim=action_dim,
            obs_dim=0 if obs_as_cond else obs_feature_dim,
            max_n_obs_steps=n_obs_steps,
            fix_obs_steps=True,
            action_visible=False,
        )
        self.normalizer = LinearNormalizer()
        self.horizon = horizon
        self.obs_feature_dim = obs_feature_dim
        self.action_dim = action_dim
        self.n_action_steps = n_action_steps
        self.n_obs_steps = n_obs_steps
        self.obs_as_cond = obs_as_cond
        self.pred_action_steps_only = pred_action_steps_only
        self.min_time = min_time
        self.du_dt_epsilon = du_dt_epsilon
        self.pmf_u_loss_weight = pmf_u_loss_weight
        self.pmf_v_loss_weight = pmf_v_loss_weight
        self.noise_scale = noise_scale
        self.adatloss_eps = adatloss_eps
        self.p_mean = p_mean
        self.p_std = p_std
        self.tr_uniform = tr_uniform
        self.tr_uniform_prob = tr_uniform_prob
        self.data_proportion = data_proportion
        self.kwargs = kwargs

        if num_inference_steps is None:
            num_inference_steps = noise_scheduler.config.num_train_timesteps
        self.num_inference_steps = num_inference_steps

    def _encode_obs(self, nobs: Dict[str, torch.Tensor], n_steps: int) -> torch.Tensor:
        flat_nobs = dict_apply(nobs, lambda x: x[:, :n_steps, ...].reshape(-1, *x.shape[2:]))
        nobs_features = self.obs_encoder(flat_nobs)
        return nobs_features.reshape(next(iter(nobs.values())).shape[0], n_steps, -1)

    def _time_view(self, value: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
        return value.reshape(value.shape[0], *([1] * (ref.ndim - 1)))

    def _adatloss(self, loss: torch.Tensor) -> torch.Tensor:
        denom = loss.detach() + self.adatloss_eps
        return loss / denom

    def _sample_logit_normal(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        normal = torch.randn(batch_size, device=device, dtype=dtype)
        return torch.sigmoid(normal * self.p_std + self.p_mean)

    def _sample_tr(self, batch_size: int, device: torch.device, dtype: torch.dtype):
        t = self._sample_logit_normal(batch_size, device, dtype)
        r = self._sample_logit_normal(batch_size, device, dtype)

        if self.tr_uniform:
            uniform_mask = torch.rand(batch_size, device=device) < self.tr_uniform_prob
            uniform_t = torch.rand(batch_size, device=device, dtype=dtype)
            uniform_r = torch.rand(batch_size, device=device, dtype=dtype)
            t = torch.where(uniform_mask, uniform_t, t)
            r = torch.where(uniform_mask, uniform_r, r)

        data_size = int(batch_size * self.data_proportion)
        fm_mask = torch.arange(batch_size, device=device) < data_size
        r = torch.where(fm_mask, t, r)

        t_final = torch.maximum(t, r)
        r_final = torch.minimum(t, r)
        return t_final, r_final

    def _trajectory_inputs(
        self,
        nobs: Dict[str, torch.Tensor],
        nactions: torch.Tensor,
    ):
        batch_size = nactions.shape[0]
        horizon = nactions.shape[1]
        cond = None
        trajectory = nactions
        if self.obs_as_cond:
            cond = self._encode_obs(nobs, self.n_obs_steps)
            if self.pred_action_steps_only:
                start = self.n_obs_steps - 1
                end = start + self.n_action_steps
                trajectory = nactions[:, start:end]
        else:
            nobs_features = self._encode_obs(nobs, horizon)
            trajectory = torch.cat([nactions, nobs_features], dim=-1).detach()

        if self.pred_action_steps_only:
            condition_mask = torch.zeros_like(trajectory, dtype=torch.bool)
        else:
            condition_mask = self.mask_generator(trajectory.shape)

        return batch_size, trajectory, cond, condition_mask

    def _apply_conditioning(
        self,
        sample: torch.Tensor,
        condition_data: torch.Tensor,
        condition_mask: torch.Tensor,
    ) -> torch.Tensor:
        if not condition_mask.any():
            return sample
        return torch.where(condition_mask, condition_data, sample)

    def _compute_u_v(
        self,
        sample: torch.Tensor,
        t: torch.Tensor,
        r: torch.Tensor,
        cond: torch.Tensor,
    ):
        x_hat_u, x_hat_v = self.model(sample, t, r, cond)
        denom = self._time_view(t, sample)
        u = (sample - x_hat_u) / denom
        v = (sample - x_hat_v) / denom
        return u, v

    def _compute_du_dt(
        self,
        sample: torch.Tensor,
        t: torch.Tensor,
        r: torch.Tensor,
        cond: torch.Tensor,
        condition_data: torch.Tensor,
        condition_mask: torch.Tensor,
        tangent_v: torch.Tensor,
    ) -> torch.Tensor:
        tangent_sample = tangent_v.detach()
        tangent_r = torch.zeros_like(r)
        tangent_t = torch.ones_like(t)

        def u_fn(sample_input, r_input, t_input):
            conditioned_sample = self._apply_conditioning(
                sample_input, condition_data, condition_mask
            )
            u_value, _ = self._compute_u_v(conditioned_sample, t_input, r_input, cond)
            return u_value

        primals = (sample, r, t)
        tangents = (tangent_sample, tangent_r, tangent_t)
        try:
            _, du_dt = torch.func.jvp(u_fn, primals, tangents)
        except (AttributeError, NotImplementedError, RuntimeError):
            _, du_dt = torch.autograd.functional.jvp(
                u_fn,
                primals,
                tangents,
                create_graph=False,
                strict=False,
            )

        return du_dt

    # ========= inference ============
    def conditional_sample(
        self,
        condition_data,
        condition_mask,
        cond=None,
        generator=None,
        **kwargs,
    ):
        del kwargs

        trajectory = torch.randn(
            size=condition_data.shape,
            dtype=condition_data.dtype,
            device=condition_data.device,
            generator=generator,
        ) * self.noise_scale

        time_steps = torch.linspace(
            1.0,
            0.0,
            self.num_inference_steps + 1,
            dtype=trajectory.dtype,
            device=trajectory.device,
        )

        for step_idx in range(self.num_inference_steps):
            trajectory = self._apply_conditioning(trajectory, condition_data, condition_mask)
            t = time_steps[step_idx].expand(trajectory.shape[0])
            r = time_steps[step_idx + 1].expand(trajectory.shape[0])
            u, _ = self._compute_u_v(trajectory, t, r, cond)
            delta = self._time_view(t - r, trajectory)
            trajectory = trajectory - delta * u

        trajectory = self._apply_conditioning(trajectory, condition_data, condition_mask)
        return trajectory

    def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        assert "past_action" not in obs_dict
        nobs = self.normalizer.normalize(obs_dict)
        value = next(iter(nobs.values()))
        batch_size, to_steps = value.shape[:2]
        horizon = self.horizon
        action_dim = self.action_dim

        device = self.device
        dtype = self.dtype
        cond = None
        if self.obs_as_cond:
            cond = self._encode_obs(nobs, self.n_obs_steps)
            shape = (batch_size, horizon, action_dim)
            if self.pred_action_steps_only:
                shape = (batch_size, self.n_action_steps, action_dim)
            cond_data = torch.zeros(size=shape, device=device, dtype=dtype)
            cond_mask = torch.zeros_like(cond_data, dtype=torch.bool)
        else:
            nobs_features = self._encode_obs(nobs, self.n_obs_steps)
            shape = (batch_size, horizon, action_dim + self.obs_feature_dim)
            cond_data = torch.zeros(size=shape, device=device, dtype=dtype)
            cond_mask = torch.zeros_like(cond_data, dtype=torch.bool)
            cond_data[:, : self.n_obs_steps, action_dim:] = nobs_features
            cond_mask[:, : self.n_obs_steps, action_dim:] = True

        nsample = self.conditional_sample(
            cond_data,
            cond_mask,
            cond=cond,
            **self.kwargs,
        )

        naction_pred = nsample[..., :action_dim]
        action_pred = self.normalizer["action"].unnormalize(naction_pred)
        if self.pred_action_steps_only:
            action = action_pred
        else:
            start = to_steps - 1
            end = start + self.n_action_steps
            action = action_pred[:, start:end]
        return {
            "action": action,
            "action_pred": action_pred,
        }

    # ========= training ============
    def set_normalizer(self, normalizer: LinearNormalizer):
        self.normalizer.load_state_dict(normalizer.state_dict())

    def get_optimizer(
        self,
        transformer_weight_decay: float,
        obs_encoder_weight_decay: float,
        learning_rate: float,
        betas: Tuple[float, float],
    ) -> torch.optim.Optimizer:
        optim_groups = self.model.get_optim_groups(weight_decay=transformer_weight_decay)
        optim_groups.append(
            {
                "params": self.obs_encoder.parameters(),
                "weight_decay": obs_encoder_weight_decay,
            }
        )
        return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)

    def compute_loss(self, batch):
        assert "valid_mask" not in batch
        nobs = self.normalizer.normalize(batch["obs"])
        nactions = self.normalizer["action"].normalize(batch["action"])

        _, trajectory, cond, condition_mask = self._trajectory_inputs(nobs, nactions)
        noise = torch.randn_like(trajectory) * self.noise_scale
        batch_size = trajectory.shape[0]

        t, r = self._sample_tr(
            batch_size, device=trajectory.device, dtype=trajectory.dtype
        )
        z_t = (1 - self._time_view(t, trajectory)) * trajectory + self._time_view(t, trajectory) * noise
        z_t = self._apply_conditioning(z_t, trajectory, condition_mask)

        loss_mask = ~condition_mask
        target_v = noise - trajectory

        u, v = self._compute_u_v(z_t, t, r, cond)
        du_dt = self._compute_du_dt(
            sample=z_t,
            t=t,
            r=r,
            cond=cond,
            condition_data=trajectory,
            condition_mask=condition_mask,
            tangent_v=v,
        )
        pmf_velocity = u + self._time_view(t - r, trajectory) * du_dt.detach()

        loss_u = F.mse_loss(pmf_velocity, target_v, reduction="none")
        loss_v = F.mse_loss(v, target_v, reduction="none")
        loss_u = loss_u * loss_mask.type(loss_u.dtype)
        loss_v = loss_v * loss_mask.type(loss_v.dtype)
        loss_u = reduce(loss_u, "b ... -> b (...)", "mean").mean()
        loss_v = reduce(loss_v, "b ... -> b (...)", "mean").mean()
        loss_u = self._adatloss(loss_u)
        loss_v = self._adatloss(loss_v)
        return self.pmf_u_loss_weight * loss_u + self.pmf_v_loss_weight * loss_v
