def raytrace (origin,pix,dep,norm,i,j):
    ray_origin = origin
    ray_destination = (i,j,0)
    s1 = (i-ray_origin[0],j-ray_origin[1],0-ray_origin[2])
    n1 = 1
    n2 = 1.33

    N_surf = norm/255

    s1_mg = np.sqrt(sum(i**2 for i in s1))
    N_surf_mg = np.sqrt(sum(i**2 for i in N_surf))

    #print(s1_mg,N_surf_mg)

    RI = n1/n2
    N_ratio = (N_surf/N_surf_mg)
    s1_ratio = (s1/s1_mg)
    part_1 = np.dot(RI,np.cross(N_ratio,(np.cross(-N_ratio,s1_ratio))))
    part_2 = (N_ratio)*np.sqrt(1-np.dot(np.dot((RI)**2,np.cross(N_ratio,s1_ratio)),np.cross(N_ratio,s1_ratio)))
    s2 = part_1-part_2
    s2_mg = np.sqrt(sum(i**2 for i in s2))

    #print(s2, s2_mg)

    theta_1 = math.acos((np.dot(np.negative(s1),N_surf))/(-s1_mg*N_surf_mg))
    theta_2 = math.asin(RI * math.sin(theta_1))
    #print(theta_1,theta_2)
    depth_position = (i,j,dep)
    s2_unit = (s2[0]/s2_mg,s2[1]/s2_mg,s2[2]/s2_mg)

    #actual ray magnitude when the depth is known
    s2_final_mg = dep/math.cos(theta_2)
    s2_final_vertex = (s2_final_mg * s2_unit[0],s2_final_mg * s2_unit[1],s2_final_mg * s2_unit[2])

    #print(s2_final_vertex)
    return s2_final_vertex



def main():
    path_background = "tex0.png"
    path_normal = "normal_001900_1.33.npy"
    path_depth = "depth_001900_1.33.npy"

    in_image = Image.open(path_background)
    out_image = Image.open(path_background)

    pix_in = in_image.load()
    pix_out = out_image.load()
    #print(ref.size)
    #print(pix_in[255,255])


    normal = np.load(path_normal)
    normal = np.array(normal)
    #print(normal[0,0,:])

    depth = np.load(path_depth)
    depth = np.array(depth)
    #print(depth[0,0])
    origin = (123,123,7) #camera right in the middle of the 256x256 image at some height

    for i in range(0,256):
        for j in range(0,256):
            (valx,valy,valz) = raytrace_generic(origin,pix_in[i,j],depth[i,j],normal[i,j,:],i,j)
            print(valx,valy)
            pix_out[i,j] = pix_in[valx*255, valy*255]

    out_image.save('example_raytrace.png')


if __name__ == "__main__":
    main()