import time
from collections import deque

import erdos

import numpy as np

import pylot.prediction.utils
from pylot.prediction.messages import PredictionMessage
from pylot.prediction.obstacle_prediction import ObstaclePrediction
from pylot.utils import Location, Transform, time_epoch_ms

import torch

try:
    from pylot.prediction.prediction.r2p2.r2p2_model import R2P2
except ImportError:
    raise Exception('Error importing R2P2.')


class R2P2PredictorOperator(erdos.Operator):
    """Wrapper operator for R2P2 ego-vehicle prediction module.

    Args:
        point_cloud_stream (:py:class:`erdos.ReadStream`, optional): Stream on
            which point cloud messages are received.
        tracking_stream (:py:class:`erdos.ReadStream`):
            Stream on which
            :py:class:`~pylot.perception.messages.ObstacleTrajectoriesMessage`
            are received.
        prediction_stream (:py:class:`erdos.ReadStream`): Stream on which
            :py:class:`~pylot.prediction.messages.PredictionMessage`
            messages are published.
        lidar_setup (:py:class:`pylot.drivers.sensor_setup.LidarSetup`): Setup
            of the lidar. This setup is used to get the maximum range of the
            lidar.
    """
    def __init__(self, point_cloud_stream, tracking_stream, prediction_stream,
                 flags, lidar_setup):
        print("WARNING: R2P2 predicts only vehicle trajectories")
        self._logger = erdos.utils.setup_logging(self.config.name,
                                                 self.config.log_file_name)
        self._csv_logger = erdos.utils.setup_csv_logging(
            self.config.name + '-csv', self.config.csv_log_file_name)
        self._flags = flags

        self._device = torch.device('cuda')
        self._r2p2_model = R2P2().to(self._device)
        state_dict = torch.load(flags.r2p2_model_path)
        self._r2p2_model.load_state_dict(state_dict)

        point_cloud_stream.add_callback(self.on_point_cloud_update)
        tracking_stream.add_callback(self.on_trajectory_update)
        erdos.add_watermark_callback([point_cloud_stream, tracking_stream],
                                     [prediction_stream], self.on_watermark)

        self._lidar_setup = lidar_setup

        self._point_cloud_msgs = deque()
        self._tracking_msgs = deque()

    @staticmethod
    def connect(point_cloud_stream, tracking_stream):
        prediction_stream = erdos.WriteStream()
        return [prediction_stream]

    def destroy(self):
        self._logger.warn('destroying {}'.format(self.config.name))

    def on_watermark(self, timestamp, prediction_stream):
        self._logger.debug('@{}: received watermark'.format(timestamp))
        if timestamp.is_top:
            return
        point_cloud_msg = self._point_cloud_msgs.popleft()
        tracking_msg = self._tracking_msgs.popleft()

        start_time = time.time()
        nearby_trajectories, nearby_vehicle_ego_transforms, \
            nearby_trajectories_tensor, binned_lidars_tensor = \
            self._preprocess_input(tracking_msg, point_cloud_msg)

        num_predictions = len(nearby_trajectories)
        self._logger.info(
            '@{}: Getting R2P2 predictions for {} vehicles'.format(
                timestamp, num_predictions))

        if num_predictions == 0:
            prediction_stream.send(PredictionMessage(timestamp, []))
            return

        # Run the forward pass.
        z = torch.tensor(
            np.random.normal(size=(num_predictions,
                                   self._flags.prediction_num_future_steps,
                                   2))).to(torch.float32).to(self._device)
        model_start_time = time.time()
        prediction_array, _ = self._r2p2_model.forward(
            z, nearby_trajectories_tensor, binned_lidars_tensor)
        model_runtime = (time.time() - model_start_time) * 1000
        self._csv_logger.debug("{},{},{},{:.4f}".format(
            time_epoch_ms(), timestamp.coordinates[0],
            'r2p2-modelonly-runtime', model_runtime))
        prediction_array = prediction_array.cpu().detach().numpy()

        obstacle_predictions_list = self._postprocess_predictions(
            prediction_array, nearby_trajectories,
            nearby_vehicle_ego_transforms)
        runtime = (time.time() - start_time) * 1000
        self._csv_logger.debug("{},{},{},{:.4f}".format(
            time_epoch_ms(), timestamp.coordinates[0], 'r2p2-runtime',
            runtime))
        prediction_stream.send(
            PredictionMessage(timestamp, obstacle_predictions_list))

    def _preprocess_input(self, tracking_msg, point_cloud_msg):

        nearby_vehicle_trajectories, nearby_vehicle_ego_transforms = \
            tracking_msg.get_nearby_obstacles_info(
                self._flags.prediction_radius,
                lambda t: t.obstacle.is_vehicle())
        point_cloud = point_cloud_msg.point_cloud.points
        num_nearby_vehicles = len(nearby_vehicle_trajectories)
        if num_nearby_vehicles == 0:
            return [], [], [], []

        # Pad and rotate the trajectory of each nearby vehicle to its
        # coordinate frame. Also, remove the z-coordinate of the trajectory.
        nearby_trajectories_tensor = []  # Pytorch tensor for network input.

        for i in range(num_nearby_vehicles):
            cur_trajectory = nearby_vehicle_trajectories[
                i].get_last_n_transforms(self._flags.prediction_num_past_steps)
            cur_trajectory = np.stack(
                [[point.location.x, point.location.y, point.location.z]
                 for point in cur_trajectory])

            rotated_trajectory = nearby_vehicle_ego_transforms[
                i].inverse_transform_points(cur_trajectory)[:, :2]

            nearby_trajectories_tensor.append(rotated_trajectory)

        nearby_trajectories_tensor = np.stack(nearby_trajectories_tensor)
        nearby_trajectories_tensor = torch.tensor(
            nearby_trajectories_tensor).to(torch.float32).to(self._device)

        # For each vehicle, transform the lidar point cloud to that vehicle's
        # coordinate frame for purposes of prediction.
        binned_lidars = []
        for i in range(num_nearby_vehicles):
            rotated_point_cloud = nearby_vehicle_ego_transforms[
                i].inverse_transform_points(point_cloud)
            binned_lidars.append(
                pylot.prediction.utils.get_occupancy_grid(
                    rotated_point_cloud,
                    self._lidar_setup.transform.location.z,
                    int(self._lidar_setup.get_range_in_meters())))
        binned_lidars = np.concatenate(binned_lidars)
        binned_lidars_tensor = torch.tensor(binned_lidars).to(
            torch.float32).to(self._device)

        return nearby_vehicle_trajectories, nearby_vehicle_ego_transforms, \
            nearby_trajectories_tensor, binned_lidars_tensor

    def _postprocess_predictions(self, prediction_array, vehicle_trajectories,
                                 vehicle_ego_transforms):
        # The prediction_array consists of predictions with respect to each
        # vehicle. Transform each predicted trajectory to be in relation to the
        # ego-vehicle, then convert into an ObstaclePrediction.
        obstacle_predictions_list = []
        num_predictions = len(vehicle_trajectories)

        for idx in range(num_predictions):
            cur_prediction = prediction_array[idx]

            obstacle_transform = vehicle_trajectories[idx].obstacle.transform
            predictions = []
            # Because R2P2 only predicts (x,y) coordinates, we assume the
            # vehicle stays at the same height as its last location.
            for t in range(self._flags.prediction_num_future_steps):
                cur_point = vehicle_ego_transforms[idx].transform_points(
                    np.array([[
                        cur_prediction[t][0], cur_prediction[t][1],
                        vehicle_ego_transforms[idx].location.z
                    ]]))[0]
                # R2P2 does not predict vehicle orientation, so we use our
                # estimated orientation of the vehicle at its latest location.
                predictions.append(
                    Transform(location=Location(cur_point[0], cur_point[1],
                                                cur_point[2]),
                              rotation=vehicle_ego_transforms[idx].rotation))

            # Probability; currently a filler value because we are taking
            # just one sample from distribution
            obstacle_predictions_list.append(
                ObstaclePrediction(vehicle_trajectories[idx],
                                   obstacle_transform, 1.0, predictions))
        return obstacle_predictions_list

    def on_point_cloud_update(self, msg):
        self._logger.debug('@{}: received point cloud message'.format(
            msg.timestamp))
        self._point_cloud_msgs.append(msg)

    def on_trajectory_update(self, msg):
        self._logger.debug('@{}: received trajectories message'.format(
            msg.timestamp))
        self._tracking_msgs.append(msg)
