from typing import Dict
import numba
import torch
import numpy as np
import copy
from diffusion_policy.common.pytorch_util import dict_apply
from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.common.sampler import (
    SequenceSampler,
    get_val_mask,
    downsample_mask,
)
from diffusion_policy.model.common.normalizer import LinearNormalizer
from diffusion_policy.dataset.base_dataset import BaseImageDataset
from diffusion_policy.common.normalize_util import get_image_range_normalizer
import pdb


class RobotImageDataset(BaseImageDataset):

    def __init__(
        self,
        zarr_path,
        horizon=1,
        pad_before=0,
        pad_after=0,
        seed=42,
        val_ratio=0.0,
        batch_size=128,
        max_train_episodes=None,
        normalizer_input_min=None,
        normalizer_input_max=None,
        action_indices=None,
        state_indices=None,
    ):

        super().__init__()
        self.replay_buffer = ReplayBuffer.copy_from_path(
            zarr_path,
            keys=["head_camera", "state", "action"],
        )

        val_mask = get_val_mask(n_episodes=self.replay_buffer.n_episodes, val_ratio=val_ratio, seed=seed)
        train_mask = ~val_mask
        train_mask = downsample_mask(mask=train_mask, max_n=max_train_episodes, seed=seed)

        self.sampler = SequenceSampler(
            replay_buffer=self.replay_buffer,
            sequence_length=horizon,
            pad_before=pad_before,
            pad_after=pad_after,
            episode_mask=train_mask,
        )
        self.train_mask = train_mask
        self.horizon = horizon
        self.pad_before = pad_before
        self.pad_after = pad_after
        self.normalizer_input_min = (
            None if normalizer_input_min is None else np.asarray(normalizer_input_min, dtype=np.float32)
        )
        self.normalizer_input_max = (
            None if normalizer_input_max is None else np.asarray(normalizer_input_max, dtype=np.float32)
        )
        self.action_indices = None if action_indices is None else np.asarray(action_indices, dtype=np.int64)
        self.state_indices = None if state_indices is None else np.asarray(state_indices, dtype=np.int64)
        if (self.normalizer_input_min is None) != (self.normalizer_input_max is None):
            raise ValueError("normalizer_input_min and normalizer_input_max must be set together")
        if self.normalizer_input_min is not None and self.normalizer_input_min.shape != self.normalizer_input_max.shape:
            raise ValueError("normalizer_input_min and normalizer_input_max must have the same shape")

        self.batch_size = batch_size
        sequence_length = self.sampler.sequence_length
        self.buffers = {
            k: np.zeros((batch_size, sequence_length, *v.shape[1:]), dtype=v.dtype)
            for k, v in self.sampler.replay_buffer.items()
        }
        self.buffers_torch = {k: torch.from_numpy(v) for k, v in self.buffers.items()}
        for v in self.buffers_torch.values():
            v.pin_memory()

    def get_validation_dataset(self):
        val_set = copy.copy(self)
        val_set.sampler = SequenceSampler(
            replay_buffer=self.replay_buffer,
            sequence_length=self.horizon,
            pad_before=self.pad_before,
            pad_after=self.pad_after,
            episode_mask=~self.train_mask,
        )
        val_set.train_mask = ~self.train_mask
        return val_set

    def _camera_to_chw_numpy(self, value):
        if value.ndim >= 3 and value.shape[-1] == 3 and value.shape[-3] != 3:
            return np.moveaxis(value, -1, -3)
        return value

    def _camera_to_chw_torch(self, value):
        if value.ndim >= 3 and value.shape[-1] == 3 and value.shape[-3] != 3:
            return torch.movedim(value, -1, -3)
        return value

    def get_normalizer(self, mode="limits", **kwargs):
        action = self.replay_buffer["action"]
        agent_pos = self.replay_buffer["state"]
        if self.action_indices is not None:
            action = action[..., self.action_indices]
        if self.state_indices is not None:
            agent_pos = agent_pos[..., self.state_indices]
        if self.normalizer_input_min is not None:
            expected_dim = int(action.shape[-1])
            if int(agent_pos.shape[-1]) != expected_dim:
                raise ValueError(
                    "shared normalizer bounds require state and action to have the same final dimension, "
                    f"got state={agent_pos.shape[-1]} action={action.shape[-1]}"
                )
            if int(self.normalizer_input_min.shape[-1]) != expected_dim:
                raise ValueError(
                    f"normalizer bounds dim {self.normalizer_input_min.shape[-1]} does not match action dim {expected_dim}"
                )
            bounds = np.stack((self.normalizer_input_min, self.normalizer_input_max), axis=0)
            action = bounds
            agent_pos = bounds
        data = {
            "action": action,
            "agent_pos": agent_pos,
        }
        normalizer = LinearNormalizer()
        normalizer.fit(data=data, last_n_dims=1, mode=mode, **kwargs)
        normalizer["head_cam"] = get_image_range_normalizer()
        return normalizer

    def __len__(self) -> int:
        return len(self.sampler)

    def _sample_to_data(self, sample):
        agent_pos = sample["state"].astype(np.float32)  # (agent_posx2, block_posex3)
        if self.state_indices is not None:
            agent_pos = agent_pos[..., self.state_indices]
        head_cam = self._camera_to_chw_numpy(sample["head_camera"]) / 255
        obs = {
            "head_cam": head_cam,
            "agent_pos": agent_pos,
        }

        data = {
            "obs": obs,
            "action": self._select_action(sample["action"]).astype(np.float32),  # T, D
        }
        return data

    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
        if isinstance(idx, slice):
            raise NotImplementedError  # Specialized
        elif isinstance(idx, int):
            sample = self.sampler.sample_sequence(idx)
            sample = dict_apply(sample, torch.from_numpy)
            return sample
        elif isinstance(idx, np.ndarray):
            assert len(idx) == self.batch_size
            for k, v in self.sampler.replay_buffer.items():
                batch_sample_sequence(
                    self.buffers[k],
                    v,
                    self.sampler.indices,
                    idx,
                    self.sampler.sequence_length,
                )
            return self.buffers_torch
        else:
            raise ValueError(idx)

    def postprocess(self, samples, device):
        agent_pos = samples["state"].to(device, non_blocking=True)
        if self.state_indices is not None:
            agent_pos = agent_pos[..., torch.as_tensor(self.state_indices, device=agent_pos.device)]
        head_cam = self._camera_to_chw_torch(samples["head_camera"].to(device, non_blocking=True)) / 255.0
        obs = {
            "head_cam": head_cam,
            "agent_pos": agent_pos,
        }
        action = samples["action"].to(device, non_blocking=True)
        if self.action_indices is not None:
            action = action[..., torch.as_tensor(self.action_indices, device=action.device)]
        return {
            "obs": obs,
            "action": action,  # B, T, D
        }

    def _select_action(self, action):
        if self.action_indices is None:
            return action
        return action[..., self.action_indices]


def _batch_sample_sequence(
    data: np.ndarray,
    input_arr: np.ndarray,
    indices: np.ndarray,
    idx: np.ndarray,
    sequence_length: int,
):
    for i in numba.prange(len(idx)):
        buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx = indices[idx[i]]
        data[i, sample_start_idx:sample_end_idx] = input_arr[buffer_start_idx:buffer_end_idx]
        if sample_start_idx > 0:
            data[i, :sample_start_idx] = data[i, sample_start_idx]
        if sample_end_idx < sequence_length:
            data[i, sample_end_idx:] = data[i, sample_end_idx - 1]


_batch_sample_sequence_sequential = numba.jit(_batch_sample_sequence, nopython=True, parallel=False)
_batch_sample_sequence_parallel = numba.jit(_batch_sample_sequence, nopython=True, parallel=True)


def batch_sample_sequence(
    data: np.ndarray,
    input_arr: np.ndarray,
    indices: np.ndarray,
    idx: np.ndarray,
    sequence_length: int,
):
    batch_size = len(idx)
    assert data.shape == (batch_size, sequence_length, *input_arr.shape[1:])
    if batch_size >= 16 and data.nbytes // batch_size >= 2**16:
        _batch_sample_sequence_parallel(data, input_arr, indices, idx, sequence_length)
    else:
        _batch_sample_sequence_sequential(data, input_arr, indices, idx, sequence_length)
