import os, sys, glob
import math
import argparse
import numpy as np
import cv2
from cv2 import aruco
import scipy.spatial
import math
import matplotlib
import matplotlib.pyplot as plt
import GhostScan.Reconstructions.generatephase as gp
from GhostScan.Calibrations.Camera2ScreenCalib import calib
from GhostScan.Reconstructions.normalInWorld import normalInWorld, normalCompare, calAngFeature

#metric_factor = 0.065 / 60 #0.001
# the dimension of glass
#half_length = 0.2875 - 0.005
#half_height = 0.1587 - 0.005
# the dimension of mirror
#half_length = 0.3016 - 0.005
#half_height = 0.1492 - 0.005
# the dimension of round glass samples
#half_length = 0.1524 - 0.005
#half_height = 0.1397 - 0.005

def PhaseMeasuringDeflectometry(rawImgPath, addImgPath, normalPath, metric_factor, cameraMatrix, disCoeffs, ret, rc2s, tc2s, half_length, half_height):
    #load intrinsic matrix and screen2cam transformation
    #sinusoid image path (for calculating phase values)
    #checkerboard image path (for geometric calibration)
    imgFile = addImgPath
    imgUndistortFile = os.path.join(rawImgPath, "0.png")
    #path of output normal map
    normalFile = os.path.join(normalPath, "normal_map.png")
    #cam2screen = np.load("D:/Deflectometry/calibration/Cam2Screen0623_2.npz")
    #rc2s, tc2s = calib(imgFileMirror, intrinsicFile)
    #tc2s = tc2s[0]
    rcamera_dis = rc2s
    tcamera_dis = tc2s
    # additional section: Make sure that the shape of rC2S and tC2S is correct
    tcamera_dis = np.reshape(tcamera_dis, (3,))
    #print("rc2s is ",rcamera_dis)
    #print("tc2s is ",tcamera_dis)   

    #load phasemaps generated by 4-phase shift algorithm
    PhaseMap0, PhaseMap1 = gp.calPhase(imgPath= rawImgPath)
    PhaseMapProjected = np.load('PhaseMapHV.npz')
    PhaseMapHC = PhaseMapProjected.f.arr_0
    PhaseMapVC = PhaseMapProjected.f.arr_1
    #disnplay phase maps
    fig = plt.figure(figsize=(15,10))
    plt.imshow(PhaseMap0)
    plt.title("Horizontal Phase Value")
    plt.colorbar()
    fig.savefig(normalPath + '/Horizontal Phase Value' + '.png')
    plt.show()
    fig = plt.figure(figsize=(15,10))
    plt.plot(PhaseMap0[800, :])
    fig.savefig(normalPath + '/y=800' + '.png')
    plt.show()
    #phaseFile = os.path.join(normalPath, "phase_map_0.png")
    #plt.savefig(phaseFile)
    #cv2.imwrite(phaseFile, PhaseMap0)

    fig = plt.figure(figsize=(15,10))
    plt.imshow(PhaseMap1)
    plt.title("Vertical Phase Value")
    plt.colorbar()
    fig.savefig(normalPath + '/Vertical Phase Value' + '.png')
    plt.show()
    fig = plt.figure(figsize=(15,10))
    plt.plot(PhaseMap1[:, 1000])
    fig.savefig(normalPath + '/x=1000' + '.png')
    plt.show()
    #define board (plane of the object), then estimate the pose of that board using the markers on the object
    allCorners = []
    allIds = []
    decimator = 0
    dictionary = aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
    img = cv2.imread(imgFile)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    cornerRefine = cv2.aruco.CORNER_REFINE_CONTOUR
    [markerCorners,markerIDs, rejectedImgPoints] = cv2.aruco.detectMarkers(gray, dictionary, cameraMatrix=cameraMatrix, distCoeff=disCoeffs)


    board_corner = np.array([np.array([[-half_length, -half_height, 0], [-half_length + 0.04, -half_height, 0], [-half_length + 0.04, -half_height + 0.04, 0], [-half_length, -half_height + 0.04, 0]]),
                    np.array([[half_length - 0.04, -half_height, 0], [half_length, -half_height, 0], [half_length, -half_height + 0.04, 0], [half_length - 0.04, -half_height + 0.04, 0]]),
                    np.array([[half_length - 0.04, half_height - 0.04, 0], [half_length, half_height - 0.04, 0], [half_length, half_height, 0], [half_length - 0.04, half_height, 0]]),
                    np.array([[-half_length, half_height - 0.04, 0], [-half_length + 0.04, half_height - 0.04, 0], [-half_length + 0.04, half_height, 0], [-half_length, half_height, 0]]),
                    np.array([[-0.02, -half_height, 0], [0.02, -half_height, 0], [0.02, -half_height + 0.04, 0], [-0.02, -half_height + 0.04, 0]]),
                    np.array([[-half_length, -0.02, 0], [-half_length + 0.04, -0.02, 0], [-half_length + 0.04, 0.02, 0], [-half_length, 0.02, 0]]),
                    np.array([[half_length - 0.04, -0.02, 0], [half_length, -0.02, 0], [half_length, 0.02, 0], [half_length - 0.04, 0.02, 0]])], dtype=np.float32)

    board_id = np.array([[0], [6], [4], [2], [7], [1], [5]], dtype=np.int32)
    board = cv2.aruco.Board_create(board_corner, dictionary, board_id)

    if len(markerCorners) > 0:
        allCorners.append(markerCorners)
        allIds.append(markerIDs)
        cv2.aruco.drawDetectedMarkers(img, markerCorners, markerIDs, [0,200,0])


    rvecs,tvecs,_objpoints = cv2.aruco.estimatePoseSingleMarkers(markerCorners, 0.04, cameraMatrix, disCoeffs, )
    r_mean = np.mean(rvecs,axis=0)
    t_mean = np.mean(tvecs,axis=0)
    r_mirror = None
    t_mirror = None
    ret, r_mirror, t_mirror, = cv2.aruco.estimatePoseBoard(markerCorners, markerIDs, board, cameraMatrix, disCoeffs, r_mirror, t_mirror)
    r_mirror_mat = np.zeros((3,3))
    r_mirror_mat = cv2.Rodrigues(r_mirror,dst= r_mirror_mat,jacobian=None)
    r_mirror_mat = r_mirror_mat[0]
    #print("r mirror mat is ", r_mirror_mat)
    img_axis = aruco.drawAxis(img, cameraMatrix, disCoeffs, r_mirror, t_mirror, 0.04)
    width = 960
    height = int(img.shape[0]*960/img.shape[1])
    smallimg = cv2.resize(img,(width,height))
    smallimg = cv2.cvtColor(smallimg, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(15,10))
    plt.imshow(smallimg)
    plt.show()
    imsize = img.shape

    #rtcam: camera coordinate system to world coordinate system
    #rtMat: world coordinate system to camera coordinate system
    rtMat = np.hstack((r_mirror_mat, t_mirror))
    rtMat = np.vstack((rtMat, np.array([[0, 0, 0, 1]])))
    rt_cam = np.linalg.inv(rtMat)
    r_cam = rt_cam[0:3, 0:3]
    t_cam = rt_cam[0:3, 3]
    #print("rtMat is",rtMat)
    #print("rtcam in world coor is ", rt_cam)

    #rtdis: display coordinate system to world coordinate system
    #rtdis_inv: world coordinate system to display coordinate system
    t_camdis = np.reshape(tcamera_dis, (tcamera_dis.shape[0],1))
    rtMat_camdis = np.hstack((rcamera_dis, t_camdis))
    rtMat_camdis = np.vstack((rtMat_camdis, np.array([0,0,0,1])))
    rt_dis_inv = np.matmul(rtMat_camdis, rtMat)
    rt_dis = np.linalg.inv(rt_dis_inv)
    r_dis = rt_dis[0:3, 0:3]
    t_dis = rt_dis[0:3,3]
    r_dis_inv = rt_dis_inv[0:3,0:3]
    t_dis_inv = rt_dis_inv[0:3,3]
    #print("rtscreen in world coor is ", rt_dis)
    #print("mirror in screen coor is ",rt_dis_inv)

    #load undistorted image
    img_undistort = cv2.imread(imgUndistortFile)
    gray_undistort = cv2.cvtColor(img_undistort, cv2.COLOR_BGR2GRAY)
    img_undistort_size = gray_undistort.shape

    #reshape t vector
    t_cam = np.reshape(t_cam, (t_cam.shape[0], 1))
    t_dis = np.reshape(t_dis, (t_dis.shape[0], 1))
    t_dis_inv = np.reshape(t_dis_inv, (t_dis_inv.shape[0], 1))

    #cam coor mat:store vectors in camera coordinate system
    #cam world coor mat: store vectors in world coordinate system
    #camera mirror intersect mat: store intersect points between arrival vectors and mirror in world coordinate system
    # display intersect mat: store display intersect points on display coordinate system
    # display intersect mat trans: store display intersect points on world coordinate system
    # reflec mat: store reflectance vectors in world coordinate system 
    # normal mat: store normalized normal vectors in world coordinate system
    # normal mat trans: store normalized normal vectors in camera coordinate system
    # normal mat origin: store mormalized normal vectors in world coordinate system
    img_coor_mat = np.zeros((3, img_undistort_size[0], img_undistort_size[1]))
    cam_coor_mat = np.zeros((3, img_undistort_size[0], img_undistort_size[1]))
    img_coor_mat_rs = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    cam_coor_mat_rs = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    cam_world_coor_mat = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    cam_mirror_intersect_mat = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    reflect_mat = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    display_intersect_mat = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    display_intersect_mat_trans = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    normal_mat = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    normal_mat_trans = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    normal_mat_origin = np.zeros((3, img_undistort_size[0]*img_undistort_size[1]))
    #calculate normal vectors 
    for i in range(img_undistort_size[0]):
        for j in range(img_undistort_size[1]):
            image_coor = np.array([j + 1, i + 1, 1])
            img_coor_mat[:,i,j] = image_coor

    img_coor_mat_rs = np.reshape(img_coor_mat, (3, img_undistort_size[0]*img_undistort_size[1]))
    # calculate arrival vector
    cam_coor_mat_rs = np.matmul(np.linalg.inv(cameraMatrix), img_coor_mat_rs)
    cam_world_coor_mat = np.matmul(r_cam, cam_coor_mat_rs)
    # calculate intersect point
    scale_factor1 = -t_cam[2,0] / cam_world_coor_mat[2,:]
    scale_factor1 = np.reshape(scale_factor1, (1, scale_factor1.shape[0]))
    cam_mirror_intersect_mat = scale_factor1 * cam_world_coor_mat + np.tile(t_cam, (cam_world_coor_mat.shape[1]))
    # reshape phase map (for phase value calculation)
    PhaseMap0_rs = np.reshape(PhaseMap0, (img_undistort_size[0]*img_undistort_size[1],))
    PhaseMap1_rs = np.reshape(PhaseMap1, (img_undistort_size[0]*img_undistort_size[1],))

    # find corresponding position on screen using phase value
    for i in range(img_undistort_size[0]*img_undistort_size[1]):
        phasevalueH = PhaseMap0_rs[i]
        phasevalueV = PhaseMap1_rs[i]
        phasediffH = np.abs(PhaseMapHC - phasevalueH)
        phasediffV = np.abs(PhaseMapVC - phasevalueV)
        h_index = np.argmin(phasediffH)
        v_index = np.argmin(phasediffV)
        h_loc = h_index - img_undistort_size[1]/2
        v_loc = v_index - img_undistort_size[0]/2
        h_loc_metric = h_loc * metric_factor
        v_loc_metric = v_loc * metric_factor
        # x axis adjustment
        h_loc_metric = -h_loc_metric
        display_intersect_point = np.array([h_loc_metric, v_loc_metric, 0])
        display_intersect_mat[:,i] = display_intersect_point

    # calculate arrival and reflectance fector
    display_intersect_mat_trans = np.dot(r_dis, display_intersect_mat) + np.tile(t_dis, (cam_world_coor_mat.shape[1]))
    reflect_mat = display_intersect_mat_trans - cam_mirror_intersect_mat
    arrival_vec_norm = -cam_world_coor_mat
    arrival_vec_norm = arrival_vec_norm / np.linalg.norm(arrival_vec_norm, axis= 0)
    reflect_vec_norm = reflect_mat / np.linalg.norm(reflect_mat, axis= 0)
    # use normalized arrival and reflectance vector to calculate normal fectors, then normalize normal vectors
    normal_mat = arrival_vec_norm + reflect_vec_norm
    normal_mat_origin = normal_mat / np.linalg.norm(normal_mat, axis= 0)
    normal_mat_trans = np.array([normal_mat_origin[1,:], -normal_mat_origin[0,:], -normal_mat_origin[2,:]])

    # restore the shape of matrix and normal vector matrix
    cam_mirror_intersect_mat = np.reshape(cam_mirror_intersect_mat, (3, img_undistort_size[0], img_undistort_size[1]))
    normal_mat_origin = np.reshape(normal_mat_origin, (3, img_undistort_size[0], img_undistort_size[1]))
    normal_mat_trans = np.reshape(normal_mat_trans, (3, img_undistort_size[0], img_undistort_size[1]))
    # the original matrix is (3, 1200, 1900) and we need to transfer it to (1200, 1920, 3) in order to be compatable to normal export codes
    cam_mirror_intersect_rs = np.zeros((img_undistort_size[0], img_undistort_size[1], 3))
    normal_mat_origin_rs = np.zeros((img_undistort_size[0], img_undistort_size[1], 3))
    normal_mat_trans_rs = np.zeros((img_undistort_size[0], img_undistort_size[1], 3))
    for i in range(normal_mat_trans.shape[0]):
        cam_mirror_intersect_rs[:,:,i] = cam_mirror_intersect_mat[i,:,:]
        # Use variance data to filter the invalid normal data (caused by diffused surface)
        if i == 2:
            normal_mat_origin_rs[:,:,i] = normal_mat_origin[i,:,:] 
            normal_mat_trans_rs[:,:,i] = normal_mat_trans[i,:,:] 
        else:
            normal_mat_origin_rs[:,:,i] = normal_mat_origin[i,:,:] 
            # try to magnify x and y elements
            normal_mat_trans_rs[:,:,i] = normal_mat_trans[i,:,:] 

    cam_mirror_intersect_mat = cam_mirror_intersect_rs
    normal_mat_origin = normal_mat_origin_rs
    normal_mat_trans = normal_mat_trans_rs

    #encode the normal vectors into the RGB image
    #np.savez(os.path.join(rootImgPath, "results/intersect_world"),cam_mirror_intersect_mat)
    np.savez(os.path.join(normalPath, "normal_map"), normal_mat_origin)
    #np.savez(os.path.join(rootImgPath, "results/normal_map_trans"), normal_mat_trans)
    normalImage = np.array((normal_mat_trans + 1) / 2.0 * 255, dtype= np.uint8)
    cv2.imwrite(normalFile, cv2.cvtColor(np.array((normal_mat_trans + 1) / 2.0 * 255, dtype= np.uint8), cv2.COLOR_RGB2BGR))
    print("normal map generated.")

    # this part transfer the normal map from camera coordiante into world coordinate
    normalmap_world = normalInWorld(cam_mirror_intersect_mat= cam_mirror_intersect_mat, normal_map= normal_mat_origin, half_length= half_length, half_height = half_height)
    return normalImage, normal_mat_origin, normalmap_world
        

        
