"""
Waymo dataset

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import numpy as np
import glob
import random

from .builder import DATASETS
from .defaults import DefaultDataset, DefaultImagePointDataset


def _append_timing_embedding(strength, timestamp):
    if strength.ndim == 1:
        strength = strength.reshape(-1, 1)
    timing = np.full((strength.shape[0], 1), timestamp, dtype=strength.dtype)
    return np.hstack((strength, timing))


@DATASETS.register_module()
class WaymoDataset(DefaultDataset):
    def __init__(
        self,
        timestamp=(0,),
        reference_label=True,
        timing_embedding=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        assert timestamp[0] == 0
        self.timestamp = timestamp
        self.reference_label = reference_label
        self.timing_embedding = timing_embedding
        self.data_list = sorted(self.data_list)
        _, self.sequence_offset, self.sequence_index = np.unique(
            [os.path.dirname(data) for data in self.data_list],
            return_index=True,
            return_inverse=True,
        )
        self.sequence_offset = np.append(self.sequence_offset, len(self.data_list))

    def get_data_list(self):
        if isinstance(self.split, str):
            self.split = [self.split]
        data_list = []
        for split in self.split:
            data_list += glob.glob(os.path.join(self.data_root, split, "*", "*"))
        return data_list

    @staticmethod
    def align_pose(coord, pose, target_pose):
        coord = np.hstack((coord, np.ones_like(coord[:, :1])))
        pose_align = np.matmul(np.linalg.inv(target_pose), pose)
        coord = (pose_align @ coord.T).T[:, :3]
        return coord

    @staticmethod
    def align_normal(normal, pose, target_pose):
        pose_align = np.matmul(np.linalg.inv(target_pose), pose)
        return (pose_align[:3, :3] @ normal.T).T

    def get_single_frame(self, idx):
        return super().get_data(idx)

    def get_data(self, idx):
        idx = idx % len(self.data_list)
        if self.timestamp == (0,):
            data_dict = self.get_single_frame(idx)
            if self.timing_embedding:
                data_dict["strength"] = _append_timing_embedding(
                    data_dict["strength"], 0
                )
            return data_dict

        sequence_index = self.sequence_index[idx]
        lower, upper = self.sequence_offset[[sequence_index, sequence_index + 1]]
        major_frame = self.get_single_frame(idx)
        name = major_frame.pop("name")
        target_pose = major_frame.pop("pose")
        if self.timing_embedding:
            major_frame["strength"] = _append_timing_embedding(
                major_frame["strength"], 0
            )
        for key in major_frame.keys():
            major_frame[key] = [major_frame[key]]

        for timestamp in self.timestamp[1:]:
            refer_idx = timestamp + idx
            if refer_idx < lower or upper <= refer_idx:
                continue
            refer_frame = self.get_single_frame(refer_idx)
            refer_frame.pop("name")
            pose = refer_frame.pop("pose")
            refer_frame["coord"] = self.align_pose(
                refer_frame["coord"], pose, target_pose
            )
            if "normal" in refer_frame:
                refer_frame["normal"] = self.align_normal(
                    refer_frame["normal"], pose, target_pose
                )
            if not self.reference_label:
                refer_frame["segment"] = (
                    np.ones_like(refer_frame["segment"]) * self.ignore_index
                )

            if self.timing_embedding:
                refer_frame["strength"] = _append_timing_embedding(
                    refer_frame["strength"], timestamp
                )

            for key in major_frame.keys():
                major_frame[key].append(refer_frame[key])
        for key in major_frame.keys():
            if isinstance(major_frame[key][0], str):
                continue
            major_frame[key] = np.concatenate(major_frame[key], axis=0)
        major_frame["name"] = name
        return major_frame

    def get_data_name(self, idx):
        file_path = self.data_list[idx % len(self.data_list)]
        sequence_path, frame_name = os.path.split(file_path)
        sequence_name = os.path.basename(sequence_path)
        data_name = f"{sequence_name}_{frame_name}"
        return data_name


@DATASETS.register_module()
class WaymoColorNormalDataset(WaymoDataset):
    @staticmethod
    def estimate_normals(points, center=np.array([0, 0, 0])):
        normals = points - center[None, :]
        norms = np.linalg.norm(normals, axis=1, keepdims=True)
        normals = normals / norms
        return normals

    def get_single_frame(self, idx):
        data_path = self.data_list[idx % len(self.data_list)]
        name = self.get_data_name(idx)
        split = self.get_split_name(idx)
        data_dict = {}
        assets = os.listdir(data_path)
        for asset in assets:
            if not asset.endswith(".npy"):
                continue
            if asset[:-4] not in self.VALID_ASSETS:
                continue
            data_dict[asset[:-4]] = np.load(os.path.join(data_path, asset))
        data_dict["name"] = name
        data_dict["split"] = split
        data_dict["color"] = np.zeros_like(data_dict["coord"])  # placeholder for color
        data_dict["normal"] = np.zeros_like(data_dict["coord"])

        if "coord" in data_dict.keys():
            data_dict["coord"] = data_dict["coord"].astype(np.float32)

        if "color" in data_dict.keys():
            data_dict["color"] = data_dict["color"].astype(np.float32)

        if "normal" in data_dict.keys():
            data_dict["normal"] = data_dict["normal"].astype(np.float32)

        if "segment" in data_dict.keys():
            data_dict["segment"] = data_dict["segment"].reshape([-1]).astype(np.int32)
        else:
            data_dict["segment"] = (
                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
            )

        if "instance" in data_dict.keys():
            data_dict["instance"] = data_dict["instance"].reshape([-1]).astype(np.int32)
        else:
            data_dict["instance"] = (
                np.ones(data_dict["coord"].shape[0], dtype=np.int32) * -1
            )
        return data_dict

    def get_data(self, idx):
        idx = idx % len(self.data_list)
        if self.timestamp == (0,):
            data_dict = self.get_single_frame(idx)
            if self.timing_embedding:
                data_dict["strength"] = _append_timing_embedding(
                    data_dict["strength"], 0
                )
            return data_dict

        sequence_index = self.sequence_index[idx]
        lower, upper = self.sequence_offset[[sequence_index, sequence_index + 1]]
        major_frame = self.get_single_frame(idx)
        name = major_frame.pop("name")
        target_pose = major_frame.pop("pose")
        if self.timing_embedding:
            major_frame["strength"] = _append_timing_embedding(
                major_frame["strength"], 0
            )
        for key in major_frame.keys():
            major_frame[key] = [major_frame[key]]

        for timestamp in self.timestamp[1:]:
            refer_idx = timestamp + idx
            if refer_idx < lower or upper <= refer_idx:
                continue
            refer_frame = self.get_single_frame(refer_idx)
            refer_frame.pop("name")
            pose = refer_frame.pop("pose")
            refer_frame["coord"] = self.align_pose(
                refer_frame["coord"], pose, target_pose
            )
            if "normal" in refer_frame:
                refer_frame["normal"] = self.align_normal(
                    refer_frame["normal"], pose, target_pose
                )
            if not self.reference_label:
                refer_frame["segment"] = (
                    np.ones_like(refer_frame["segment"]) * self.ignore_index
                )

            if self.timing_embedding:
                refer_frame["strength"] = _append_timing_embedding(
                    refer_frame["strength"], timestamp
                )

            for key in major_frame.keys():
                major_frame[key].append(refer_frame[key])
        for key in major_frame.keys():
            if isinstance(major_frame[key][0], str):
                major_frame[key] = major_frame[key][0]
                continue
            major_frame[key] = np.concatenate(major_frame[key], axis=0)
        major_frame["name"] = name
        return major_frame


@DATASETS.register_module()
class WaymoImagePointDataset(DefaultImagePointDataset):
    def __init__(
        self,
        if_img=False,
        if_sweep=False,
        reference_label=True,
        timing_embedding=False,
        sweeps=3,
        sweep_gap=5,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.if_img = if_img
        self.if_sweep = if_sweep
        if if_sweep:
            self.timestamp = tuple(range(0, sweeps * sweep_gap, sweep_gap))
        else:
            self.timestamp = (0,)
        self.reference_label = reference_label
        self.timing_embedding = timing_embedding
        self.sweeps = sweeps
        self.data_name = sorted(self.data_name)
        _, self.sequence_offset, self.sequence_index = np.unique(
            [data.split("_with_camera_labels_")[0] for data in self.data_name],
            return_index=True,
            return_inverse=True,
        )
        self.sequence_offset = np.append(self.sequence_offset, len(self.data_name))

    @staticmethod
    def align_pose(coord, pose, target_pose):
        coord = np.hstack((coord, np.ones_like(coord[:, :1])))
        pose_align = np.matmul(np.linalg.inv(target_pose), pose)
        coord = (pose_align @ coord.T).T[:, :3]
        return coord

    @staticmethod
    def align_normal(normal, pose, target_pose):
        pose_align = np.matmul(np.linalg.inv(target_pose), pose)
        return (pose_align[:3, :3] @ normal.T).T

    def get_single_frame(self, idx):
        return super().get_data(idx)

    def get_data(self, idx):
        idx = idx % len(self.data_list)
        if self.timestamp == (0,):
            data_dict = self.get_single_frame(idx)
            if self.timing_embedding:
                data_dict["strength"] = _append_timing_embedding(
                    data_dict["strength"], 0
                )
            return data_dict

        sequence_index = self.sequence_index[idx]
        lower, upper = self.sequence_offset[[sequence_index, sequence_index + 1]]
        valid_timestamps = [0] + [
            timestamp
            for timestamp in self.timestamp[1:]
            if lower <= timestamp + idx < upper
        ]
        imgs_idx = random.sample(valid_timestamps, 1)[0]
        major_frame = self.get_single_frame(idx)
        name = major_frame.pop("name")
        target_pose = major_frame.pop("pose")
        if self.timing_embedding:
            major_frame["strength"] = _append_timing_embedding(
                major_frame["strength"], 0
            )
        if not self.if_img:
            pass
        elif imgs_idx == 0:
            major_frame["correspondence"] = [major_frame["correspondence"]]
        else:
            major_frame.pop("img_num")
            major_frame.pop("images")
            major_frame["correspondence"] = [
                -np.ones_like(major_frame["correspondence"])
            ]
        for key in major_frame.keys():
            if key in self.PC_VALID_ASSETS:
                major_frame[key] = [major_frame[key]]

        for timestamp in valid_timestamps[1:]:
            refer_idx = timestamp + idx
            refer_frame = self.get_single_frame(refer_idx)
            refer_frame.pop("name")
            pose = refer_frame.pop("pose")
            refer_frame["coord"] = self.align_pose(
                refer_frame["coord"], pose, target_pose
            )
            if "normal" in refer_frame:
                refer_frame["normal"] = self.align_normal(
                    refer_frame["normal"], pose, target_pose
                )
            if not self.reference_label:
                refer_frame["segment"] = (
                    np.ones_like(refer_frame["segment"]) * self.ignore_index
                )

            if self.timing_embedding:
                refer_frame["strength"] = _append_timing_embedding(
                    refer_frame["strength"], timestamp
                )

            for key in major_frame.keys():
                if key in self.PC_VALID_ASSETS:
                    major_frame[key].append(refer_frame[key])

            if self.if_img:
                if imgs_idx == timestamp:
                    major_frame["img_num"] = refer_frame["img_num"]
                    major_frame["images"] = refer_frame["images"]
                    major_frame["correspondence"].append(refer_frame["correspondence"])
                else:
                    major_frame["correspondence"].append(
                        -np.ones_like(refer_frame["correspondence"])
                    )
        frame_pcd_offset = [
            frame_coord.shape[0] for frame_coord in major_frame["coord"]
        ]
        frame_pcd_offset = np.cumsum(frame_pcd_offset)

        for key in major_frame.keys():
            if key in self.PC_VALID_ASSETS + ["correspondence"]:
                major_frame[key] = np.concatenate(major_frame[key], axis=0)
        major_frame["name"] = name
        major_frame["frame_pcd_offset"] = frame_pcd_offset
        return major_frame