# this part is to compare generated "ground truth" normal map and calculated normal map in world coordinate system
 # scale factor, used to discretize the intersection points in metric system (from world coordinate system to image coordinate system)
from __future__ import print_function
import os, sys, time, math
import subprocess
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib
import GhostScan.Reconstructions.generatephase as gp


def normalInWorld(cam_mirror_intersect_mat, normal_map, half_length, half_height):
    SCALE_FACTOR = 2000

    intersect_point = cam_mirror_intersect_mat
    #varMask = np.load(os.path.join(rootImgPath, "results/mask.npz"))
    #varMask = varMask.f.arr_0
    intersect_x = intersect_point[:,:,0]
    intersect_y = intersect_point[:,:,1]
    normal_map_x = normal_map[:,:,0]
    normal_map_y = normal_map[:,:,1]
    normal_map_z = normal_map[:,:,2]

    
    intersect_x = intersect_x + half_length
    intersect_y = intersect_y + half_height
    intersect_x[intersect_x < 0] = 0
    intersect_x[intersect_x > 2 * half_length] = 0
    intersect_y[intersect_y < 0] = 0
    intersect_y[intersect_y > 2 * half_height] = 0

    print(np.max(intersect_x))
    print(np.max(intersect_y))

    #discretize corld intersect point coordinate.
    intersect_x = intersect_x * SCALE_FACTOR
    intersect_y = intersect_y * SCALE_FACTOR
    intersect_x = np.round(intersect_x)
    intersect_y = np.round(intersect_y)


    normal_map_world = np.zeros((int(np.round(2 * half_height * SCALE_FACTOR) + 1), int(np.round(2 * half_length * SCALE_FACTOR) + 1), 3))
    #varMask_world = np.zeros((int(np.round(2 * half_height * SCALE_FACTOR) + 1), int(np.round(2 * half_length * SCALE_FACTOR) + 1)))
    for i in range(intersect_x.shape[0]):
        for j in range(intersect_x.shape[1]):
            normal_map_world[int(intersect_y[i,j]), int(intersect_x[i,j]), 0] = normal_map_x[i,j] 
            normal_map_world[int(intersect_y[i,j]), int(intersect_x[i,j]), 1] = normal_map_y[i,j]
            normal_map_world[int(intersect_y[i,j]), int(intersect_x[i,j]), 2] = normal_map_z[i,j]
            #varMask_world[int(intersect_y[i,j]), int(intersect_x[i,j])] = varMask[i,j]

    # use mask to filter invalid normals
    #ground_truth_x = ground_truth_x * varMask_world
    #ground_truth_y = ground_truth_y * varMask_world
    #ground_truth_z = ground_truth_z * varMask_world + (varMask_world - 1)
    for i in range(normal_map_world.shape[2]):
        if i == 2:
            normal_map_world[:,:,i] = normal_map_world[:,:,i] #* varMask_world + (varMask_world - 1)
        else:
            normal_map_world[:,:,i] = normal_map_world[:,:,i] #* varMask_world

    return normal_map_world
def normalRotate(normal_map_world, tilt= 0):
    normal_map_rs = np.zeros((3, normal_map_world.shape[0], normal_map_world.shape[1]))
    for i in range(normal_map_world.shape[2]):
        normal_map_rs[i,:,:] = normal_map_world[:,:,i]
    normal_map_rs = np.reshape(normal_map_rs, (3, normal_map_world.shape[0] * normal_map_world.shape[1]))
    tilt_rad = tilt * 2 * math.pi / 360
    rMat = np.array([[1, 0, 0],[0, np.cos(tilt_rad), -np.sin(tilt_rad)],[0, np.sin(tilt_rad), np.cos(tilt_rad)]])
    print(rMat)
    normal_map_rs = np.dot(rMat, normal_map_rs)
    normal_map_temp = np.reshape(normal_map_rs,(3, normal_map_world.shape[0], normal_map_world.shape[1]))
    normal_map_tilted = np.zeros((normal_map_world.shape[0], normal_map_world.shape[1], 3))
    for i in range(normal_map_temp.shape[0]):
        normal_map_tilted[:,:,i] = normal_map_temp[i,:,:]
        
    
    norm = matplotlib.colors.Normalize(vmin= -0.2, vmax = 0.2)
    norm2 = matplotlib.colors.Normalize(vmin= -1, vmax = -0.98)
    # visualize tilted normal map
    """
    plt.figure()
    for i in range(3):
        #plt.figure()
        if i == 2:
            plt.subplot(1,3,i+1)
            plt.imshow(normal_map_tilted[:,:,i], norm= norm2)
            plt.title("Tilt Z")
            plt.tick_params(labelsize=16)
        else:
            plt.subplot(1,3,i+1)
            plt.imshow(normal_map_tilted[:,:,i], norm= norm)
            if i == 1:
                plt.title("Tilt Y")
            else:
                plt.title("Tilt X")
            plt.tick_params(labelsize=16)
        cb = plt.colorbar()
        cb.ax.tick_params(labelsize= 16)
    plt.show()
    """
    return normal_map_tilted

def normalCompare(rootImgPath, normal_map_world, half_length, half_height):
    SCALE_FACTOR = 2500
    # load ground truth normal map generated by Matlab
    intersect_point = np.load(os.path.join(rootImgPath, "results/intersect_world.npz"))
    ground_truth_x = np.loadtxt(os.path.join(rootImgPath, "results/normx_world.txt"))
    ground_truth_y = np.loadtxt(os.path.join(rootImgPath, "results/normy_world.txt"))
    ground_truth_z = np.loadtxt(os.path.join(rootImgPath, "results/normz_world.txt"))
    varMask = np.load(os.path.join(rootImgPath, "results/mask.npz"))
    varMask = varMask.f.arr_0
    intersect_x = intersect_point.f.arr_0[:,:,0]
    intersect_y = intersect_point.f.arr_0[:,:,1]

    intersect_x = intersect_x + half_length
    intersect_y = intersect_y + half_height
    intersect_x[intersect_x < 0] = 0
    intersect_x[intersect_x > 2 * half_length] = 0
    intersect_y[intersect_y < 0] = 0
    intersect_y[intersect_y > 2 * half_height] = 0

    print(np.max(intersect_x))
    print(np.max(intersect_y))

    #discretize corld intersect point coordinate.
    intersect_x = intersect_x * SCALE_FACTOR
    intersect_y = intersect_y * SCALE_FACTOR
    intersect_x = np.round(intersect_x)
    intersect_y = np.round(intersect_y)

    varMask_world = np.zeros((int(np.round(2 * half_height * SCALE_FACTOR) + 1), int(np.round(2 * half_length * SCALE_FACTOR) + 1)))
    for i in range(intersect_x.shape[0]):
        for j in range(intersect_x.shape[1]):
            varMask_world[int(intersect_y[i,j]), int(intersect_x[i,j])] = varMask[i,j]

    norm = matplotlib.colors.Normalize(vmin= -0.2, vmax = 0.2)
    norm2 = matplotlib.colors.Normalize(vmin= -1, vmax = -0.98)
    plt.figure()
    plt.subplot(1,3,1)
    plt.imshow(ground_truth_x, norm= norm)
    plt.title("Ground Truth X")
    cb1 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb1.ax.tick_params(labelsize= 16)
    plt.subplot(1,3,2)
    plt.imshow(ground_truth_y, norm= norm)
    plt.title("Ground Truth Y")
    cb2 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb2.ax.tick_params(labelsize= 16)
    plt.subplot(1,3,3)
    plt.imshow(ground_truth_z, norm= norm2)
    plt.title('Ground Truth Z')
    cb3 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb3.ax.tick_params(labelsize= 16)
    plt.show()
    # tilt normal map to compensate error
    normal_map_world = normalRotate(normal_map_world, tilt=1)
    # calculate errors
    error_x = np.abs(ground_truth_x - normal_map_world[:,:,0]) * varMask_world
    error_y = np.abs(ground_truth_y - normal_map_world[:,:,1]) * varMask_world
    error_z = np.abs(ground_truth_z - normal_map_world[:,:,2]) * varMask_world 
    # try to erode edge errors
    kernel = np.ones((2,2))
    error_x = cv2.erode(error_x, kernel, iterations= 1)
    error_y = cv2.erode(error_y, kernel, iterations= 1)
    error_z = cv2.erode(error_z, kernel, iterations= 1)
    norm3 = matplotlib.colors.Normalize(vmin= 0, vmax = 0.1)
    norm4 = matplotlib.colors.Normalize(vmin= 0, vmax = 0.005)
    plt.figure()
    plt.subplot(1,3,1)
    plt.imshow(error_x, norm= norm3)
    plt.title('Error X')
    cb4 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb4.ax.tick_params(labelsize= 16)
    plt.subplot(1,3,2)
    plt.imshow(error_y, norm= norm3)
    plt.title('Error Y')
    cb5 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb5.ax.tick_params(labelsize= 16)
    plt.subplot(1,3,3)
    plt.imshow(error_z, norm= norm4)
    plt.title('Error Z')
    cb6 = plt.colorbar()
    plt.tick_params(labelsize=16)
    cb6.ax.tick_params(labelsize= 16)
    plt.show()

    mean_err_x = np.sum(error_x) / np.sum(varMask_world)
    mean_err_y = np.sum(error_y) / np.sum(varMask_world)
    mean_err_z = np.sum(np.abs(ground_truth_z - normal_map_world[:,:,2]) * varMask_world) / np.sum(varMask_world)
    print("mean error on x, y, z axis is ", mean_err_x, mean_err_y, mean_err_z)
    # save the calculate non-zero error
    error_x = error_x[error_x != 0]
    error_y = error_y[error_y != 0]
    error_z = error_z[error_z != 0]
    np.savetxt(os.path.join(rootImgPath, "results/error_x.txt"), error_x)
    np.savetxt(os.path.join(rootImgPath, "results/error_y.txt"), error_y)
    np.savetxt(os.path.join(rootImgPath, "results/error_z.txt"), error_z)

def calAngFeature(normal_map_world):
    horizontal = np.sqrt(normal_map_world[:,:,0] ** 2 + normal_map_world[:,:,1] ** 2)
    sine_value = horizontal / np.ones((normal_map_world.shape[0], normal_map_world.shape[1]))
    angle_rad = np.arcsin(sine_value)
    angle_value = np.degrees(angle_rad)
    norm = matplotlib.colors.Normalize(vmin= 0, vmax = 20)
    plt.figure()
    plt.imshow(angle_value, norm= norm)
    plt.colorbar()
    plt.title('Angle Descriptors')
    plt.show()




