import os
import json
from PIL import Image

import numpy as np
from numpy.core.numeric import full
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
import sys

from pcl_generator import PointCloudGenerator
from lidar_generator import PseudoLidarGenerator


def load_pose(extrinsics: dict) -> np.array:
    """
    Load pose as 4x4 matrix from (x, y, z, yaw)

    :param extrinsics: Dictionary with extrinsic parameters x, y, z and yaw angle.
    :return T: 4x4 transformation matrix as numpy array
    """
    sinval = np.sin(extrinsics['yaw'] * np.pi / 180.)
    cosval = np.cos(extrinsics['yaw'] * np.pi / 180.)
    Rz = np.array([[cosval, -sinval, 0],
                   [sinval, cosval, 0],
                   [0, 0, 1]])
    t = np.array([extrinsics['x'], extrinsics['y'], extrinsics['z']])

    T = np.eye(4)
    T[:3, :3] = Rz
    T[:3, 3] = Rz @ t
    return T


class PseudoLidarData(Dataset):
    """
    Adapted from https://github.com/autonomousvision/transfuser/blob/main/transfuser/data.py
    """

    def __init__(self, root, config):

        self.seq_len = config.seq_len
        self.pred_len = config.pred_len
        self.ignore_sides = config.ignore_sides
        self.ignore_rear = config.ignore_rear

        self.input_resolution = config.input_resolution
        self.scale = config.scale

        with open('pseudolidar/extrinsics.json', mode='r') as f:
            extrinsics = json.load(f)
            self.front_pose = load_pose(extrinsics['cam_front'])
            self.left_pose = load_pose(extrinsics['cam_left'])
            self.right_pose = load_pose(extrinsics['cam_right'])
            self.rear_pose = load_pose(extrinsics['cam_rear'])

        self.img_dims = config.image_dim
        self.img_fov = config.fov

        self.pc_generator = PointCloudGenerator(self.img_fov, self.img_dims)
        self.lidar_generator = PseudoLidarGenerator([self.front_pose,
                                                     self.left_pose,
                                                     self.right_pose,
                                                     self.rear_pose])

        self.lidar = []
        self.front = []
        self.left = []
        self.right = []
        self.rear = []
        self.depth_front = []
        self.depth_left = []
        self.depth_right = []
        self.depth_rear = []
        self.x = []
        self.y = []
        self.x_command = []
        self.y_command = []
        self.theta = []
        self.steer = []
        self.throttle = []
        self.brake = []
        self.command = []
        self.velocity = []

        for sub_root in tqdm(root, file=sys.stdout):
            preload_file = os.path.join(sub_root, 'pseudo_lidar_diag_pl_' +
                                        str(self.seq_len)+'_'+str(self.pred_len)+'.npy')

            # dump to npy if no preload
            if not os.path.exists(preload_file):
                preload_front = []
                preload_left = []
                preload_right = []
                preload_rear = []
                preload_lidar = []
                preload_depth_front = []
                preload_depth_left = []
                preload_depth_right = []
                preload_depth_rear = []
                preload_x = []
                preload_y = []
                preload_x_command = []
                preload_y_command = []
                preload_theta = []
                preload_steer = []
                preload_throttle = []
                preload_brake = []
                preload_command = []
                preload_velocity = []

                # list sub-directories in root
                root_files = os.listdir(sub_root)
                routes = [folder for folder in root_files if not os.path.isfile(os.path.join(sub_root, folder))]
                for route in routes:
                    route_dir = os.path.join(sub_root, route)

                    # subtract final frames (pred_len) since there are no future waypoints
                    # first frame of sequence not used

                    num_seq = (len(os.listdir(route_dir+"/rgb_front/"))-self.pred_len-2)//self.seq_len

                    for seq in range(num_seq):
                        fronts = []
                        lefts = []
                        rights = []
                        rears = []
                        lidars = []
                        depth_fronts = []
                        depth_lefts = []
                        depth_rights = []
                        depth_rears = []
                        xs = []
                        ys = []
                        thetas = []

                        # read files sequentially (past and current frames)
                        for i in range(self.seq_len):
                            # images
                            filename = f"{str(seq*self.seq_len+1+i).zfill(4)}.png"
                            fronts.append(route_dir+"/rgb_front/"+filename)
                            lefts.append(route_dir+"/rgb_left/"+filename)
                            rights.append(route_dir+"/rgb_right/"+filename)
                            rears.append(route_dir+"/rgb_rear/"+filename)

                            # point cloud
                            lidars.append(route_dir + f"/lidar/{str(seq*self.seq_len+1+i).zfill(4)}.npy")
                            depth_fronts.append(route_dir + f"/depth_front/{str(seq*self.seq_len+1+i).zfill(4)}.png")
                            depth_lefts.append(route_dir + f"/depth_left/{str(seq*self.seq_len+1+i).zfill(4)}.png")
                            depth_rights.append(route_dir + f"/depth_right/{str(seq*self.seq_len+1+i).zfill(4)}.png")
                            depth_rears.append(route_dir + f"/depth_rear/{str(seq*self.seq_len+1+i).zfill(4)}.png")

                            # position
                            with open(route_dir + f"/measurements/{str(seq*self.seq_len+1+i).zfill(4)}.json", "r") as read_file:
                                data = json.load(read_file)
                            xs.append(data['x'])
                            ys.append(data['y'])
                            thetas.append(data['theta'])

                        # get control value of final frame in sequence
                        preload_x_command.append(data['x_command'])
                        preload_y_command.append(data['y_command'])
                        preload_steer.append(data['steer'])
                        preload_throttle.append(data['throttle'])
                        preload_brake.append(data['brake'])
                        preload_command.append(data['command'])
                        preload_velocity.append(data['speed'])

                        # read files sequentially (future frames)
                        for i in range(self.seq_len, self.seq_len + self.pred_len):
                            # point cloud
                            lidars.append(route_dir + f"/lidar/{str(seq*self.seq_len+1+i).zfill(4)}.npy")
                            depth_fronts.append(route_dir + f"/depth_front/{str(seq*self.seq_len+1+i).zfill(4)}.npy")
                            depth_lefts.append(route_dir + f"/depth_left/{str(seq*self.seq_len+1+i).zfill(4)}.npy")
                            depth_rights.append(route_dir + f"/depth_right/{str(seq*self.seq_len+1+i).zfill(4)}.npy")
                            depth_rears.append(route_dir + f"/depth_rear/{str(seq*self.seq_len+1+i).zfill(4)}.npy")

                            # position
                            with open(route_dir + f"/measurements/{str(seq*self.seq_len+1+i).zfill(4)}.json", "r") as read_file:
                                data = json.load(read_file)
                            xs.append(data['x'])
                            ys.append(data['y'])

                            # fix for theta=nan in some measurements
                            if np.isnan(data['theta']):
                                thetas.append(0)
                            else:
                                thetas.append(data['theta'])

                        preload_front.append(fronts)
                        preload_left.append(lefts)
                        preload_right.append(rights)
                        preload_rear.append(rears)
                        preload_lidar.append(lidars)
                        preload_depth_front.append(depth_fronts)
                        preload_depth_left.append(depth_lefts)
                        preload_depth_right.append(depth_rights)
                        preload_depth_rear.append(depth_rears)
                        preload_x.append(xs)
                        preload_y.append(ys)
                        preload_theta.append(thetas)

                # dump to npy
                preload_dict = {}
                preload_dict['front'] = preload_front
                preload_dict['left'] = preload_left
                preload_dict['right'] = preload_right
                preload_dict['rear'] = preload_rear
                preload_dict['lidar'] = preload_lidar
                preload_dict['depth_front'] = preload_depth_front
                preload_dict['depth_left'] = preload_depth_left
                preload_dict['depth_right'] = preload_depth_right
                preload_dict['depth_rear'] = preload_depth_rear
                preload_dict['x'] = preload_x
                preload_dict['y'] = preload_y
                preload_dict['x_command'] = preload_x_command
                preload_dict['y_command'] = preload_y_command
                preload_dict['theta'] = preload_theta
                preload_dict['steer'] = preload_steer
                preload_dict['throttle'] = preload_throttle
                preload_dict['brake'] = preload_brake
                preload_dict['command'] = preload_command
                preload_dict['velocity'] = preload_velocity
                np.save(preload_file, preload_dict)

            # load from npy if available
            preload_dict = np.load(preload_file, allow_pickle=True)
            self.front += preload_dict.item()['front']
            self.left += preload_dict.item()['left']
            self.right += preload_dict.item()['right']
            self.rear += preload_dict.item()['rear']
            self.lidar += preload_dict.item()['lidar']
            self.depth_front += preload_dict.item()['depth_front']
            self.depth_left += preload_dict.item()['depth_left']
            self.depth_right += preload_dict.item()['depth_right']
            self.depth_rear += preload_dict.item()['depth_rear']
            self.x += preload_dict.item()['x']
            self.y += preload_dict.item()['y']
            self.x_command += preload_dict.item()['x_command']
            self.y_command += preload_dict.item()['y_command']
            self.theta += preload_dict.item()['theta']
            self.steer += preload_dict.item()['steer']
            self.throttle += preload_dict.item()['throttle']
            self.brake += preload_dict.item()['brake']
            self.command += preload_dict.item()['command']
            self.velocity += preload_dict.item()['velocity']
            print("Preloading " + str(len(preload_dict.item()['front'])) + " sequences from " + preload_file)

    def __len__(self):
        """Returns the length of the dataset. """
        return len(self.front)

    def __getitem__(self, index):
        """Returns the item at index idx. """
        data = dict()
        data['fronts'] = []
        data['lefts'] = []
        data['rights'] = []
        data['rears'] = []
        data['lidars'] = []

        seq_fronts = self.front[index]
        seq_lefts = self.left[index]
        seq_rights = self.right[index]
        seq_rears = self.rear[index]
        seq_lidars = self.lidar[index]
        seq_depth_fronts = self.depth_front[index]
        seq_depth_lefts = self.depth_left[index]
        seq_depth_rights = self.depth_right[index]
        seq_depth_rears = self.depth_rear[index]
        seq_x = self.x[index]
        seq_y = self.y[index]
        seq_theta = self.theta[index]

        full_lidar = []
        pos = []
        neg = []
        for i in range(self.seq_len):
            data['fronts'].append(torch.from_numpy(np.array(
                scale_and_crop_image(Image.open(seq_fronts[i]), scale=self.scale, crop=self.input_resolution))))
            if not self.ignore_sides:
                data['lefts'].append(torch.from_numpy(np.array(
                    scale_and_crop_image(Image.open(seq_lefts[i]), scale=self.scale, crop=self.input_resolution))))
                data['rights'].append(torch.from_numpy(np.array(
                    scale_and_crop_image(Image.open(seq_rights[i]), scale=self.scale, crop=self.input_resolution))))
            if not self.ignore_rear:
                data['rears'].append(torch.from_numpy(np.array(
                    scale_and_crop_image(Image.open(seq_rears[i]), scale=self.scale, crop=self.input_resolution))))
            """
            lidar_unprocessed = np.load(seq_lidars[i])[..., :3]  # lidar: XYZI
            full_lidar.append(lidar_unprocessed)
            """

            pc_front = self.pc_generator.generate(seq_depth_fronts[i], max_depth=0.03)
            pc_left = self.pc_generator.generate(seq_depth_lefts[i], max_depth=0.03)
            pc_right = self.pc_generator.generate(seq_depth_rights[i], max_depth=0.03)
            pc_rear = self.pc_generator.generate(seq_depth_rears[i], max_depth=0.03)

            pseudolidar = self.lidar_generator.generate([pc_front, pc_left, pc_right, pc_rear])
            full_lidar.append(pseudolidar)

            # fix for theta=nan in some measurements
            if np.isnan(seq_theta[i]):
                seq_theta[i] = 0.

        ego_x = seq_x[i]
        ego_y = seq_y[i]
        ego_theta = seq_theta[i]

        # future frames
        for i in range(self.seq_len, self.seq_len + self.pred_len):
            lidar_unprocessed = np.load(seq_lidars[i])
            full_lidar.append(lidar_unprocessed)

        # lidar and waypoint processing to local coordinates
        waypoints = []
        for i in range(self.seq_len + self.pred_len):
            # waypoint is the transformed version of the origin in local coordinates
            # we use 90-theta instead of theta
            # LBC code uses 90+theta, but x is to the right and y is downwards here
            local_waypoint = transform_2d_points(np.zeros((1, 3)),
                                                 np.pi/2-seq_theta[i], -seq_x[i], -seq_y[i], np.pi/2-ego_theta, -ego_x, -ego_y)
            waypoints.append(tuple(local_waypoint[0, :2]))

            # process only past lidar point clouds
            if i < self.seq_len:
                # convert coordinate frame of point cloud
                """
                full_lidar[i][:, 1] *= -1  # inverts x, y
                full_lidar[i] = transform_2d_points(full_lidar[i],
                                                    np.pi/2-seq_theta[i], -seq_x[i], -seq_y[i], np.pi/2-ego_theta, -ego_x, -ego_y)
                lidar_processed = lidar_to_histogram_features(full_lidar[i], crop=self.input_resolution)
                """
                lidar_processed = PseudoLidarGenerator.post_process(full_lidar[i])
                lidar_processed = PseudoLidarGenerator.project_to_bev(lidar_processed)
                data['lidars'].append(lidar_processed.copy())

        data['waypoints'] = waypoints

        # convert x_command, y_command to local coordinates
        # taken from LBC code (uses 90+theta instead of theta)
        R = np.array([
            [np.cos(np.pi/2+ego_theta), -np.sin(np.pi/2+ego_theta)],
            [np.sin(np.pi/2+ego_theta),  np.cos(np.pi/2+ego_theta)]
        ])
        local_command_point = np.array([self.x_command[index]-ego_x, self.y_command[index]-ego_y])
        local_command_point = R.T.dot(local_command_point)
        data['target_point'] = tuple(local_command_point)

        data['steer'] = self.steer[index]
        data['throttle'] = self.throttle[index]
        data['brake'] = self.brake[index]
        data['command'] = self.command[index]
        data['velocity'] = self.velocity[index]

        return data


def lidar_to_histogram_features(lidar, crop=256):
    """
    Convert LiDAR point cloud into 2-bin histogram over 256x256 grid
    """
    def splat_points(point_cloud):
        # 256 x 256 grid
        pixels_per_meter = 16
        hist_max_per_pixel = 5
        x_meters_max = 8
        y_meters_max = 16
        xbins = np.linspace(-2*x_meters_max, 2*x_meters_max+1, 2*x_meters_max*pixels_per_meter+1)
        ybins = np.linspace(-y_meters_max, 0, y_meters_max*pixels_per_meter+1)
        hist = np.histogramdd(point_cloud[..., :2], bins=(xbins, ybins))[0]
        hist[hist > hist_max_per_pixel] = hist_max_per_pixel
        overhead_splat = hist/hist_max_per_pixel
        return overhead_splat

    below = lidar[lidar[..., 2] <= -2.0]
    above = lidar[lidar[..., 2] > -2.0]
    below_features = splat_points(below)
    above_features = splat_points(above)
    features = np.stack([below_features, above_features], axis=-1)
    features = np.transpose(features, (2, 0, 1)).astype(np.float32)
    return features


def scale_and_crop_image(image, scale=1, crop=256):
    """
    Scale and crop a PIL image, returning a channels-first numpy array.
    """
    # image = Image.open(filename)
    (width, height) = (int(image.width // scale), int(image.height // scale))
    im_resized = image.resize((width, height))
    image = np.asarray(im_resized)
    start_x = height//2 - crop//2
    start_y = width//2 - crop//2
    cropped_image = image[start_x:start_x+crop, start_y:start_y+crop]
    cropped_image = np.transpose(cropped_image, (2, 0, 1))
    return cropped_image


def transform_2d_points(xyz, r1, t1_x, t1_y, r2, t2_x, t2_y):
    """
    Build a rotation matrix and take the dot product.
    """
    # z value to 1 for rotation
    xy1 = xyz.copy()
    xy1[:, 2] = 1

    c, s = np.cos(r1), np.sin(r1)
    r1_to_world = np.matrix([[c, s, t1_x], [-s, c, t1_y], [0, 0, 1]])

    # np.dot converts to a matrix, so we explicitly change it back to an array
    world = np.asarray(r1_to_world @ xy1.T)

    c, s = np.cos(r2), np.sin(r2)
    r2_to_world = np.matrix([[c, s, t2_x], [-s, c, t2_y], [0, 0, 1]])
    world_to_r2 = np.linalg.inv(r2_to_world)

    out = np.asarray(world_to_r2 @ world).T

    # reset z-coordinate
    out[:, 2] = xyz[:, 2]

    return out
