# coding=utf-8
# get prepare data from argoverse dataset into trajectory, person_box, other_box
import argparse
import cv2
import os
import operator
import json
import pickle
import numpy as np

from glob import glob
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument("datapath", help="under this would be the video_ids folder")
parser.add_argument("newframepath")
parser.add_argument("outpath")

class2classid = {
    "VEHICLE": 1,
    "PEDESTRIAN": 0,
    "ON_ROAD_OBSTACLE": 3,
    "LARGE_VEHICLE": 1,
    "BICYCLE": 8,
    "BICYCLIST": 8,
    "BUS": 1,
    "OTHER_MOVER": 3,
    "TRAILER": 1,
    "MOTORCYCLIST": 8,
    "MOPED": 8,
    "MOTORCYCLE": 8,
    #"STROLLER": 8,
    "EMERGENCY_VEHICLE": 1,
    #"ANIMAL": 14,
    #"WHEELCHAIR": 15,
    "SCHOOL_BUS": 1,
}

def get_center(bbox):
  x1, y1, x2, y2 = bbox
  return (x1 + x2)/2.0, (y1 + y2)/2.0

def find_closest(frame_file, label_files):
  label_int_idxs = [os.path.splitext(os.path.basename(one))[0]
                    for one in label_files]
  label_int_idxs = sorted([int(one.split("_")[-1]) for one in label_int_idxs])
  label_int_idxs = np.array(label_int_idxs)

  frame_int = int(
      os.path.splitext(os.path.basename(frame_file))[0].split("_")[-1])

  closest_ind = np.argmin(np.absolute(label_int_idxs - frame_int))
  closest_label_file = label_files[closest_ind]

  return closest_label_file


from scipy.spatial.transform import Rotation
# these are from argoverse dataset API
def get_2d_box_from_3d_box(label, camera_params):

  def transform_point_cloud(point_cloud, transform_matrix):
    num_pts = point_cloud.shape[0]
    homogeneous_pts = np.hstack([point_cloud, np.ones((num_pts, 1))])
    transformed_point_cloud = homogeneous_pts.dot(transform_matrix.T)
    return transformed_point_cloud[:, :3]

  def get_3d_points(label):
    tr_x = label["center"]["x"]
    tr_y = label["center"]["y"]
    tr_z = label["center"]["z"]
    translation = np.array([tr_x, tr_y, tr_z])

    rot_w = label["rotation"]["w"]
    rot_x = label["rotation"]["x"]
    rot_y = label["rotation"]["y"]
    rot_z = label["rotation"]["z"]
    quaternion = np.array([rot_w, rot_x, rot_y, rot_z])
    quaternion = quat2rotmat(quaternion)

    transform_matrix = get_transform_matrix(quaternion, translation)

    length = label["length"]
    width = label["width"]
    height = label["height"]

    x_corners = length / 2 * np.array([1, 1, 1, 1, -1, -1, -1, -1])
    y_corners = width / 2 * np.array([1, -1, -1, 1, 1, -1, -1, 1])
    z_corners = height / 2 * np.array([1, 1, -1, -1, 1, 1, -1, -1])
    corners_object_frame = np.vstack((x_corners, y_corners, z_corners)).T

    return transform_point_cloud(corners_object_frame, transform_matrix)


  # [8, 3]  # all 8 box corners
  points_3d = get_3d_points(label)
  num_pts = points_3d.shape[0]
  # [4, 8]
  points_3d_h = np.hstack([points_3d, np.ones((num_pts, 1))]).T

  R = camera_params.extrinsic[:3, :3]
  t = camera_params.extrinsic[:3, 3]
  transform_matrix = get_transform_matrix(R, t)

  points_egovehicle = points_3d_h.T[:, :3]
  # [8, 3] # coordinates in the camera's coordinate system
  uv_cam = transform_point_cloud(points_egovehicle, transform_matrix)

  # [8, 3]
  # project to the camera image plane
  # last dim negative means behind the camera
  uvh = proj_cam_to_uv(uv_cam, camera_params)

  # [4]
  box_2d = get_2d_bbox(uvh, camera_params.img_width, camera_params.img_height)

  return box_2d

def proj_cam_to_uv(uv_cam, camera_params):
  num_points = uv_cam.shape[0]
  uvh = np.zeros((num_points, 3))
  # (x_transformed_m, y_transformed_m, z_transformed_m)

  for idx in range(num_points):
    x_transformed_m = uv_cam[idx, 0]
    y_transformed_m = uv_cam[idx, 1]
    z_transformed_m = uv_cam[idx, 2]

    z_transformed_fixed_m = z_transformed_m

    # If we're behind the camera, z value (homogeneous coord w in image plane)
    # will be negative. If we are on the camera, there would be division by zero
    # later. To prevent that, move an epsilon away from zero.

    Z_EPSILON = 1.0e-4
    if np.absolute(z_transformed_m) <= Z_EPSILON:
        z_transformed_fixed_m = np.sign(z_transformed_m) * Z_EPSILON

    pinhole_x = x_transformed_m / z_transformed_fixed_m
    pinhole_y = y_transformed_m / z_transformed_fixed_m

    K = camera_params.intrinsic
    u_px = K[0, 0] * pinhole_x + K[0, 1] * pinhole_y + K[0, 2]

    v_px = K[1, 1] * pinhole_y + K[1, 2]

    uvh[idx] = np.array([u_px, v_px, z_transformed_m])

  #uv = uvh[:, :2]
  return uvh


def quat2rotmat(q):
  assert np.isclose(np.linalg.norm(q), 1.0, atol=1e-12)
  w, x, y, z = q
  q_scipy = np.array([x, y, z, w])
  return Rotation.from_quat(q_scipy).as_dcm()

def get_transform_matrix(rotation, translation):
    transform_matrix = np.eye(4)
    transform_matrix[:3, :3] = rotation
    transform_matrix[:3, 3] = translation
    return transform_matrix

def get_2d_bbox(bbox_3d, max_w, max_h):
  """Given the computed [8, 3] points with depth, get the one bbox."""
  if all(bbox_3d[:, 2] > 0):
    # make one 2d bbox from 8 points
    x1 = round(np.min(bbox_3d[:, 0]), 3)
    y1 = round(np.min(bbox_3d[:, 1]), 3)
    x2 = round(np.max(bbox_3d[:, 0]), 3)
    y2 = round(np.max(bbox_3d[:, 1]), 3)
    if x1 > max_w or y1 > max_h:
      return None
    if x2 < 0 or y2 < 0:
      return None
    if x1 < 0:
      x1 = 0
    if y1 < 0:
      y1 = 0
    if x2 > max_w:
      x2 = max_w
    if y2 > max_h:
      y2 = max_h
    return [x1, y1, x2, y2]
  else:
    return None

class Camera_params:
  def __init__(self, camera_config):
    all_camera_data = camera_config["camera_data_"]
    for camera_data in all_camera_data:
        if "image_raw_ring_front_center" in camera_data["key"]:
            camera_config = camera_data["value"]
            break
    vehicle_SE3_sensor = camera_config["vehicle_SE3_camera_"]
    egovehicle_t_camera = np.array(vehicle_SE3_sensor["translation"])
    egovehicle_q_camera = vehicle_SE3_sensor["rotation"]["coefficients"]
    egovehicle_R_camera = quat2rotmat(egovehicle_q_camera)

    extrinsic = get_transform_matrix(
        egovehicle_R_camera.T, egovehicle_R_camera.T.dot(-egovehicle_t_camera))

    intrinsic_matrix = np.zeros((3, 4))
    intrinsic_matrix[0, 0] = camera_config["focal_length_x_px_"]
    intrinsic_matrix[0, 1] = camera_config["skew_"]
    intrinsic_matrix[0, 2] = camera_config["focal_center_x_px_"]
    intrinsic_matrix[1, 1] = camera_config["focal_length_y_px_"]
    intrinsic_matrix[1, 2] = camera_config["focal_center_y_px_"]
    intrinsic_matrix[2, 2] = 1.0

    img_width = 1920.0
    img_height = 1200.0

    distortion_coef = camera_config["distortion_coefficients_"]

    self.extrinsic = extrinsic
    self.intrinsic = intrinsic_matrix
    self.img_height = img_height
    self.img_width = img_width
    self.distortion_coef = distortion_coef

if __name__ == "__main__":
  args = parser.parse_args()

  video_paths = glob(os.path.join(args.datapath, "*"))

  # we will clipping the original image from 1920x1200 to 1920x1080
  clip_height = 120.0
  drop_frame = 12  # original is 30 fps
  target_resolution = (1920.0, 1080.0)

  def clip_box(box):
    x1, y1, x2, y2 = box
    y1 -= clip_height
    y2 -= clip_height
    y1 = 0.0 if y1 < 0.0 else y1
    y2 = 0.0 if y2 < 0.0 else y2
    return [x1, y1, x2, y2]

  for video_path in tqdm(video_paths):
    video_id = video_path.strip("/").split("/")[-1]

    track_labels = os.path.join(video_path, "per_sweep_annotations_amodal", "*")
    track_labels = sorted(glob(track_labels))

    rgb_frames = os.path.join(video_path, "ring_front_center", "*")
    rgb_frames = sorted(glob(rgb_frames))

    camera_json = os.path.join(
        video_path, "vehicle_calibration_info.json")
    with open(camera_json) as f:
      camera_json = json.load(f)
    camera_params = Camera_params(camera_json)

    # 1. collect the tracklet info
    anno_data = []  # track_id, box, frame_id
    trackid_mapping = {}  # we will have new track_id

    for frame_idx, rgb_frame in enumerate(rgb_frames):
      track_label = find_closest(rgb_frame, track_labels)

      with open(track_label) as f:
        labels = json.load(f)
      for label in labels:
        classname = label["label_class"]

        track_uuid = label["track_label_uuid"]
        if track_uuid not in trackid_mapping:
          trackid_mapping[track_uuid] = len(trackid_mapping)
        track_id = trackid_mapping[track_uuid]

        if "occlusion" in label:
          occlusion = label["occlusion"]
        else:
          occlusion = 0
        if occlusion == 100:
          continue

        # [x1, y1, x2, y2]
        track_box = get_2d_box_from_3d_box(label, camera_params)
        if track_box is None:  # box is behind camera
          continue

        track_box = clip_box(track_box)

        anno_data.append([track_id, track_box, frame_idx, classname])

    # 2. get the needed frameIdx after dropping some (original 30 fps)
    frame_idxs = {}
    for one in anno_data:
      if one[-1] == "PEDESTRIAN":
        frame_idxs[int(one[-2])] = 1
    frame_idxs = sorted(frame_idxs.keys())
    needed_frame_idxs = frame_idxs[::drop_frame]
    if len(needed_frame_idxs) < 8 + 12:
      print("warning, %s video has only %s frames, skipped.." % (
          video_id, len(frame_idxs)))
      continue

    # 3. save all the data into # frame_idx -> data
    frame_data = {}
    for one in anno_data:
      track_id, (x1, y1, x2, y2), frame_idx, classname = one
      if frame_idx not in needed_frame_idxs:
        continue
      if classname not in class2classid:  # ignore some classes
        continue
      if frame_idx not in frame_data:
        frame_data[frame_idx] = []
      frame_data[frame_idx].append({
          "class_name": classname,
          "track_id": track_id,
          "bbox": [x1, y1, x2, y2]
      })

    # 4. gather data for each frame_idx, each person_idx
    traj_data = []  # [frame_idx, person_idx, x, y]
    person_box_data = {}  # (frame_idx, person_id) -> boxes
    other_box_data = {}  # (frame_idx, person_id) -> other boxes + boxclasids
    for frame_idx in needed_frame_idxs:
      box_list = frame_data[frame_idx]
      box_list.sort(key=operator.itemgetter("track_id"))
      for i, box in enumerate(box_list):
        class_name = box["class_name"]
        track_id = box["track_id"]
        bbox = box["bbox"]
        if class_name == "PEDESTRIAN":
          person_key = "%s_%d_%d" % (video_id, frame_idx, track_id)

          x, y = get_center(bbox)

          # ignore points outside of current resolution
          if (x > target_resolution[0]) or (y > target_resolution[1]):
            continue

          traj_data.append((frame_idx, float(track_id), x, y))

          person_box_data[person_key] = bbox

          all_other_boxes = [box_list[j]["bbox"]
                             for j in range(len(box_list)) if j != i]
          all_other_boxclassids = [class2classid[box_list[j]["class_name"]]
                                   for j in range(len(box_list)) if j != i]

          other_box_data[person_key] = (all_other_boxes,
                                        all_other_boxclassids)

    # save the data
    # 1. clip the frame
    target_frame_path = os.path.join(args.newframepath, video_id)
    if not os.path.exists(target_frame_path):
      os.makedirs(target_frame_path)
    for frame_idx in needed_frame_idxs:
      ori_frame_file = rgb_frames[frame_idx]
      im = cv2.imread(ori_frame_file, cv2.IMREAD_COLOR)
      im = im[int(clip_height):, :]

      cv2.imwrite(os.path.join(
          target_frame_path, "%s_F_%08d.jpg" % (video_id, frame_idx)),
                  im)

    # 2. save the annotations
    traj_path = os.path.join(args.outpath, "traj_2.5fps", "test")
    person_box_path = os.path.join(args.outpath, "anno_person_box", "test")
    other_box_path = os.path.join(args.outpath, "anno_other_box", "test")
    if not os.path.exists(traj_path):
      os.makedirs(traj_path)
    if not os.path.exists(person_box_path):
      os.makedirs(person_box_path)
    if not os.path.exists(other_box_path):
      os.makedirs(other_box_path)

    desfile = os.path.join(traj_path, "%s.txt" % video_id)

    delim = "\t"

    with open(desfile, "w") as f:
      for i, p, x, y in traj_data:
        f.writelines("%d%s%.1f%s%.6f%s%.6f\n" % (i, delim, p, delim, x,
                                                 delim, y))

    with open(os.path.join(person_box_path,
                           "%s.p" % video_id), "wb") as f:
      pickle.dump(person_box_data, f)

    with open(os.path.join(other_box_path,
                           "%s.p" % video_id), "wb") as f:
      pickle.dump(other_box_data, f)