'''
Class for
 1) creating metric depth map by fusing LiDAR beams and relative inverse depth map from MiDaS
 2) estimating position of objects defined by bounding boxes
'''

import rospy

import numpy as np
from pathlib import Path
import time

from scipy import stats
from sklearn.linear_model import HuberRegressor
from sklearn.linear_model import RANSACRegressor

# debug
import matplotlib.pyplot as plt

### path
FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]

class DepthFusion:

    def __init__(self):

        ### launch parameters
        try:
            CALIB_PATH = str(rospy.get_param('calib_path'))
        except Exception:
            print("Can't load Fusion parameters from ROS parameter server, using defaults.")
            CALIB_PATH = 'data/calib_params_MS_424_240.npz'
        
        ### parameters from geometry # TODO z yaml
        # Breach / Leela
        # [m] height from lidar to camera (+ when camera above, - when below)
        self.cam_h = -0.080 # -0.080 / 0.083
        # [m] horizontal distance from lidar center to camera center
        self.offset = 0.055 # 0.055 / 0.030
        # [deg] offset of the lidar beam pointing directly in front of the robot. Lidar spins CCW.
        self.lidar_yaw = np.deg2rad(90) # 90 / 270
        # [deg] positive when facing down
        self.cam_pitch = np.deg2rad(-2) # around 0
        # [deg] positive when rotated to the left (CCW)
        self.cam_yaw = np.deg2rad(0) # around 0 +- 3
        # [deg] possible error in camera yaw angle. Fusing chooses error that fits regression the most
        self.cam_yaw_errors = np.deg2rad(np.arange(-2, 2.1, 0.5)) # 0 or around [-3 : 0.5 : 3] to [-2 : 1 : 2]
        self.cam_yaw_errors = np.append(self.cam_yaw_errors,[0]) # add one element

        ### init atributes
        self.DO_DEBUGING = False
        self.regression_mode = 'hr' # ols / (rm) / hr / ran
        self.min_depth = 0.15 # [m] minimal distance in depth map
        self.max_depth = 10. # [m] maximal distance in depth map

        self.scale = 0.0
        self.shift = 0.0

        # inputs
        self.lidar_ranges = None        # lidar scan
        self.lidar_range_min = None	    # minimum lidar distance
        self.lidar_range_max = None	    # maximum lidar distance
        self.depth_map_rel = None       # np array of depth map
        self.boxes = None               # array of bounding boxes

        # calibration files
        calib_path = ROOT / CALIB_PATH
        calib_files = np.load(calib_path)
        self.fx = calib_files['cam_mtx'][0,0]
        self.fy = calib_files['cam_mtx'][1,1]
        self.cx = calib_files['cam_mtx'][0,2]
        self.cy = calib_files['cam_mtx'][1,2]


    '''
    Fuse LiDAR data and relative inverse depth map to obtain metric depth map
        Output:
            self.depth_map_abs -> np.array, metric depth map
    '''
    def fuse_depth(self):

        ### LOAD PARAMETERS 

        # camera image size [px]
        img_w = self.depth_map_rel.shape[1]
        img_h = self.depth_map_rel.shape[0]

        # camera angle of view [rad]
        FOV_x = 2*np.arctan(img_w/2/self.fx)
        FOV_y = 2*np.arctan(img_h/2/self.fy)

        # init for error correction
        regression_fittness = 5 * np.ones((len(self.cam_yaw_errors)),dtype=np.float64)

        ### --- FOR ALL CAMERA YAW ERROR ANGLES --- ###
        ### finds yaw_error that fits the regression the best
        for j, yaw_error in enumerate(self.cam_yaw_errors):

            if j >= self.cam_yaw_errors.shape[0] - 1:
                # take angle with the best fit
                yaw_error = self.cam_yaw_errors[np.argmin(regression_fittness)]
                rospy.loginfo(f'YAW_ERROR: {np.rad2deg(yaw_error)}')
                
            # lidar orientation
            # TODO cam_yaw brat jako promenou z action serveru
            lidar_angle = self.lidar_yaw + self.cam_yaw + yaw_error

            ### CALCULATIONS

            idx_1 = np.int( np.ceil( (lidar_angle - FOV_x/2) * 180.0 / np.pi ) )
            idx_2 = np.int( np.floor( (lidar_angle + FOV_x/2) * 180.0 / np.pi ) )

            ### match lidar distances to corresponding camera pixels

            # angle from left to right (cca -29..29 deg) ... [rad]
            angle_x_list =  np.deg2rad([*range( int(-(idx_2 - idx_1)/2), int((idx_2 - idx_1)/2) + 1 )])
            # distances from camera center - from left to right
            L_list_all = list(np.flip(self.lidar_ranges[idx_1:idx_2]) - self.offset)
            # number of lidar beams in camera FOV
            angle_x_len = len(L_list_all)

            self.y_depth = [] # for regresion plot
            self.x_depth = [] # for regresion plot
            L_list = [] # inverse depth [m^-1]   - for valid lidar beams
            D_list = [] # inverse relative depth - for valid lidar beams

            # for every lidar beam
            for i in range(angle_x_len):

                # skip invalid distances
                if L_list_all[i] >= self.lidar_range_max or L_list_all[i] <= self.lidar_range_min:
                    continue
                
                alpha = np.arctan(L_list_all[i] / self.cam_h)
                beta = np.pi/2 - self.cam_pitch - alpha
                x = int(np.floor( self.fx * np.tan(angle_x_list[i]) + self.cx ))
                y = int(np.round( self.cy + self.fy * np.tan(beta) ))

                # if within image
                if x >= 0 and x < img_w and y >= 0 and y < img_h:
                    self.x_depth.append(x)
                    self.y_depth.append(y)
                    # create inverse depth pairs for visible beams
                    L_list.append( 1 / np.sqrt(L_list_all[i]**2 + self.cam_h**2) )
                    # D_list.append( self.depth_map_rel[y, x] )
                    R = 1
                    try:
                        D_list.append( np.median( self.depth_map_rel[y-R:y+R+1, x-R:x+R+1] ) )
                    except:
                        D_list.append( self.depth_map_rel[y, x] )
                    

            ### fit line: L_list = scale * D_list + shift

            t0 = time.time()

            # Ordinary Least Squares (ols)
            if self.DO_DEBUGING or self.regression_mode == 'ols':
                A = np.vstack([D_list, np.ones(len(D_list))]).T
                B = np.array([L_list]).T
                result_ols = np.linalg.lstsq(A,B,rcond=None)
                scale_ols = result_ols[0][0]
                shift_ols = result_ols[0][1]
                e_ols = result_ols[1][0] # Sum of squared residuals (ideally 0)
                
                t1 = time.time()
                dt1 = np.round(1000*(t1-t0))

            # # robust regression - Repeated Medians (rm)
            # if self.DO_DEBUGING or self.regression_mode == 'rm':
            #     result_rm = stats.siegelslopes(L_list, D_list) # method='separate' doesn't help
            #     scale_rm = result_rm[0]
            #     shift_rm = result_rm[1]

            #     t2 = time.time()
            #     dt2 = np.round(1000*(t2-t1))

            # robust regression - Huber Regressor (hr)
            if self.DO_DEBUGING or self.regression_mode == 'hr':
                D_list_huber = np.array(D_list)[:, np.newaxis] # row to column vector
                huber = HuberRegressor().fit(D_list_huber, L_list)
                scale_hr = huber.coef_[0]
                shift_hr = huber.intercept_
                e_hr = 1. - huber.score(D_list_huber, L_list) # 1-R^2 (ideally 0)

                t3 = time.time()
                dt3 = np.round(1000*(t3-t1))

            # obust regression - RANSAC (ran)
            if self.DO_DEBUGING or self.regression_mode == 'ran':
                D_list_ransac = np.array(D_list)[:, np.newaxis]
                ransac = RANSACRegressor().fit(D_list_huber, L_list)
                scale_ransac = ransac.estimator_.coef_[0]
                shift_ransac = ransac.estimator_.intercept_
                e_ran = 1. - ransac.score(D_list_ransac, L_list) # 1-R^2 (ideally 0)

                t4 = time.time()
                dt4 = np.round(1000*(t4-t3))

            # append regression error
            if self.regression_mode == 'ols':
                regression_fittness[j] = e_ols
            elif self.regression_mode == 'rm':
                regression_fittness[j] = 0
            elif self.regression_mode == 'hr':
                regression_fittness[j] = e_hr
            elif self.regression_mode == 'ran':
                regression_fittness[j] = e_ran

            # log duration
            if self.DO_DEBUGING:
                # rospy.loginfo(f'Regression [ms]: OLS {dt1}; RM {dt2}; HR {dt3}; RAN {dt4}.')
                rospy.loginfo(f'Regression [ms]: OLS {dt1}; HR {dt3}; RAN {dt4}.')

        ### convert depth map to meters

        # choose method
        if self.regression_mode == 'ols':
            scale, shift = scale_ols, shift_ols
        # elif self.regression_mode == 'rm':
        #     scale, shift = scale_rm, shift_rm
        elif self.regression_mode == 'hr':
            scale, shift = scale_hr, shift_hr
        elif self.regression_mode == 'ran':
            scale, shift = scale_ransac, shift_ransac
        
        # find relative values coresponding to min and max metric depth
        rel_max = (1 - self.max_depth * shift) / (self.max_depth * scale)
        rel_min = (1 - self.min_depth * shift) / (self.min_depth * scale)

        # clip relative depth map
        self.depth_map_rel[self.depth_map_rel < rel_max] = rel_max
        self.depth_map_rel[self.depth_map_rel > rel_min] = rel_min

        # compute metric depth map
        self.depth_map_abs = 1.0 / (scale * self.depth_map_rel + shift)

        # log scale and shift
        self.scale = scale
        self.shift = shift

        ### DEBUGGING
        # show regressions
        if self.DO_DEBUGING:
            D_list_line = [np.min(D_list), np.max(D_list)]
            L_list_line_ols = [scale_ols * L + shift_ols for L in D_list_line]
            # L_list_line_rm = [scale_rm * L + shift_rm for L in D_list_line]
            L_list_line_hr = [scale_hr * L + shift_hr for L in D_list_line]
            L_list_line_ran = [scale_ransac * L + shift_ransac for L in D_list_line]
            plt.figure(figsize=[5, 5], dpi=200)
            plt.plot(D_list, L_list, 'b+') # plot relative depths (depth map, lidar)
            plt.plot(D_list_line, L_list_line_ols, 'r-') # draw line from OLS
            # plt.plot(D_list_line, L_list_line_rm, 'k-') # draw line from RM
            plt.plot(D_list_line, L_list_line_hr, 'm-') # draw line from HR
            plt.plot(D_list_line, L_list_line_ran, 'g-') # draw line from ransac
            plt.xlabel('inverzní relativní  vzdálenost z hloubkové mapy [-]')
            plt.ylabel('inverzní vzdálenost z LiDARu [$\mathregular{m^{-1}}$]')
            plt.title('Lineární regrese')
            plt.legend(['lidar beams', f'ols {e_ols:.2f}',
                f'hr {e_hr:.2f}; {np.rad2deg(yaw_error):.1f}°',
                f'ran {e_ran:.2f}; {np.rad2deg(yaw_error):.1f}°'], loc ='best')
            plt.savefig(ROOT / "scripts/test_files/regression.png",bbox_inches='tight')
            plt.close()
        
    '''
    Obtain distance to object from absolute depth map
        Output:
            self.object_distances -> np.array, shape: (pair_count,1), metric distances of pots
    '''
    def get_distance(self):

        ### center of boxes of pots
        self.xs = np.round( ( self.boxes[:,0] + self.boxes[:,2] ) / 2 ).astype('int64')[:, np.newaxis]
        self.ys = np.round( ( self.boxes[:,1] + self.boxes[:,3] ) / 2 ).astype('int64')[:, np.newaxis]

        ### compute distances of pots
        # median distance of nearby area
        R = 2 # "radius of square"

        self.object_distances = np.zeros((self.xs.shape[0],1), dtype=np.float)
        
        for i, (x, y) in enumerate(zip(self.xs,self.ys)):
            try:
                self.object_distances[i] = np.median(
                    self.depth_map_abs[y[0]-R:y[0]+R+1, x[0]-R:x[0]+R+1] )
            except:
                self.object_distances[i] = self.depth_map_abs[y[0], x[0]]

        ### DEBUGGING
        # show depth metric depth map with lidar points and object points
        # center of bounding box
        if self.DO_DEBUGING:
            plt.figure(figsize=[7, 3], dpi=200)
            plt.imshow(self.depth_map_abs, cmap = 'plasma')
            plt.colorbar(label='Depth [m]', orientation='vertical')
            plt.scatter(self.x_depth, self.y_depth, color='green', marker='.') # lidar points
            plt.scatter(self.xs, self.ys, color='red', marker='X') # points of center of objects
            plt.savefig(ROOT / "scripts/test_files/metric_depth.png",bbox_inches='tight')
            plt.close()

            # save relative depth
            plt.figure(figsize=[4, 2], dpi=200)
            plt.imshow(self.depth_map_rel, cmap = 'inferno')
            plt.savefig(ROOT / f"scripts/test_files/relative_depth_{self.scale:0.5f}_{self.shift:0.3f}.png",bbox_inches='tight')
            plt.close()

    '''
    Obtain object position [x,y,z] from 2D point and distance
        Output:
            self.object_positions -> np.array, shape: (pair_count,3), positions of pots
    '''
    def get_position(self):

        # calculate x,y,z (vectorized)
        alfa_u = np.arctan( \
            ( (self.ys-self.cy)**2 - (self.xs-self.cx)**2 ) / \
            (self.fx**2+self.fy**2) )
        d = np.cos(alfa_u) * self.object_distances # z [m]
        sx = d * (self.xs-self.cx) / self.fx # x [m]
        sy = d * (self.ys-self.cy) / self.fy # y [m]

        self.object_positions = np.concatenate((sx,sy,d),1) # x, y, z (in camera coordite system)

