# Copyright 2016 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Utilities for processing depth images.
"""
import numpy as np
import src.rotation_utils as ru
import src.utils as utils

def get_camera_matrix(width, height, fov):
  """Returns a camera matrix from image size and fov."""
  xc = (width-1.) / 2.
  zc = (height-1.) / 2.
  f = (width / 2.) / np.tan(np.deg2rad(fov / 2.))
  camera_matrix = utils.Foo(xc=xc, zc=zc, f=f)
  return camera_matrix

def get_point_cloud_from_z(Y, camera_matrix):
  """Projects the depth image Y into a 3D point cloud.
  Inputs:
    Y is ...xHxW
    camera_matrix
  Outputs:
    X is positive going right
    Y is positive into the image
    Z is positive up in the image
    XYZ is ...xHxWx3
  """
  x, z = np.meshgrid(np.arange(Y.shape[-1]),
                     np.arange(Y.shape[-2]-1, -1, -1))
  for i in range(Y.ndim-2):
    x = np.expand_dims(x, axis=0)
    z = np.expand_dims(z, axis=0)
  X = (x-camera_matrix.xc) * Y / camera_matrix.f
  Z = (z-camera_matrix.zc) * Y / camera_matrix.f
  XYZ = np.concatenate((X[...,np.newaxis], Y[...,np.newaxis],
                        Z[...,np.newaxis]), axis=X.ndim)
  return XYZ

def make_geocentric(XYZ, sensor_height, camera_elevation_degree):
  """Transforms the point cloud into geocentric coordinate frame.
  Input:
    XYZ                     : ...x3
    sensor_height           : height of the sensor
    camera_elevation_degree : camera elevation to rectify.
  Output:
    XYZ : ...x3
  """
  R = ru.get_r_matrix([1.,0.,0.], angle=np.deg2rad(camera_elevation_degree))
  XYZ = np.matmul(XYZ.reshape(-1,3), R.T).reshape(XYZ.shape)
  XYZ[...,2] = XYZ[...,2] + sensor_height
  return XYZ

def bin_points(XYZ_cms, map_size, z_bins, xy_resolution):
  """Bins points into xy-z bins
  XYZ_cms is ... x H x W x3
  Outputs is ... x map_size x map_size x (len(z_bins)+1)
  """
  sh = XYZ_cms.shape
  XYZ_cms = XYZ_cms.reshape([-1, sh[-3], sh[-2], sh[-1]])
  n_z_bins = len(z_bins)+1
  map_center = (map_size-1.)/2.
  counts = []
  isvalids = []
  for XYZ_cm in XYZ_cms:
    isnotnan = np.logical_not(np.isnan(XYZ_cm[:,:,0]))
    X_bin = np.round(XYZ_cm[:,:,0] / xy_resolution + map_center).astype(np.int32)
    Y_bin = np.round(XYZ_cm[:,:,1] / xy_resolution + map_center).astype(np.int32)
    Z_bin = np.digitize(XYZ_cm[:,:,2], bins=z_bins).astype(np.int32)

    isvalid = np.array([X_bin >= 0, X_bin < map_size, Y_bin >= 0, Y_bin < map_size,
                        Z_bin >= 0, Z_bin < n_z_bins, isnotnan])
    isvalid = np.all(isvalid, axis=0)

    ind = (Y_bin * map_size + X_bin) * n_z_bins + Z_bin
    ind[np.logical_not(isvalid)] = 0
    count = np.bincount(ind.ravel(), isvalid.ravel().astype(np.int32),
                         minlength=map_size*map_size*n_z_bins)
    count = np.reshape(count, [map_size, map_size, n_z_bins])
    counts.append(count)
    isvalids.append(isvalid)
  counts = np.array(counts).reshape(list(sh[:-3]) + [map_size, map_size, n_z_bins])
  isvalids = np.array(isvalids).reshape(list(sh[:-3]) + [sh[-3], sh[-2], 1])
  return counts, isvalids