#!/usr/bin/env python3

import os,sys,math
import numpy as np
import cv2

import matplotlib
matplotlib.use('tkagg') #CV2 causes some gear grinding with matplotlib's version of Qt. This fixes it via magic bullshit.
import matplotlib.pyplot as plt

import scipy
import scipy.interpolate as interp

import threading #this is actually useless for multithreaded execution!

import multiprocessing as mp
from multiprocessing import shared_memory

cos=np.cos
sin=np.sin
tan=np.tan
atan=np.arctan
pi = np.pi
sqrt = np.sqrt
exp = np.exp

"""
Script to stitch multiple output images into a full sky hdr image.

Scroll down to the bottom for the main routine.
"""

def load_img(fname):
    """
    Data plumbing to fix all the stupid things that opencv does in terms of color plane order, etc
    output will be an RGBA image (Nx,Ny,4), float32, scaled from 0 to 1, and 
    """

    try:
        img = cv2.imread(fname,cv2.IMREAD_UNCHANGED)
    except:
        print("load_img: error: could not open {} as an image with cv2".format(fname))
        return None
    
    Nx = img.shape[0]
    Ny = img.shape[1]

    imgout = np.zeros((Nx,Ny,4),dtype=np.float32)

    if(len(img.shape)<=2):
        #grayscale?
        imgout[:,:,0] = img[:,:]
        imgout[:,:,1] = img[:,:]
        imgout[:,:,2] = img[:,:]
        imgout[:,:,3] = 255.0

        imgout = imgout/255.0
        imgout = np.clip(imgout,0.0,1.0)
        
        imgout = np.transpose(imgout,(1,0,2))
    else:
        if(img.shape[2]==3):
            #BGR image
            
            imgout[:,:,0] = img[:,:,2]
            imgout[:,:,1] = img[:,:,1]
            imgout[:,:,2] = img[:,:,0]
            imgout[:,:,3] = 255.0

            imgout = imgout/255.0
            imgout = np.clip(imgout,0.0,1.0)
            
            imgout = np.transpose(imgout,(1,0,2))

        elif(img.shape[2]>=4):
            #BGRA image
            
            imgout[:,:,0] = img[:,:,2]
            imgout[:,:,1] = img[:,:,1]
            imgout[:,:,2] = img[:,:,0]
            imgout[:,:,3] = img[:,:,3]

            imgout = imgout/255.0
            imgout = np.clip(imgout,0.0,1.0)
            
            imgout = np.transpose(imgout,(1,0,2))

    return imgout

def save_img(fname,img):
    """
    takes a floating point (0,1) normalized image 
        greyscale: (Nx,Ny) 
        RGB: (Nx,Ny,3) 
        RGBA: (Nx,Ny,4)
    
    saves it to fname via cv2
    """

    Ny = img.shape[0]
    Nx = img.shape[1]
    if(len(img.shape)==2):
        Ncp = 1
        img = img.copy()
        img = img.transpose()
    else:
        Ncp = img.shape[2]
        img = img.copy()
        img = img.transpose(1,0,2)

    im2 = np.zeros((Nx,Ny,4),dtype=np.float32)
    im3 = np.zeros((Nx,Ny,4),dtype=np.uint8)

    if(Ncp==1):
        im2[:,:,0] = img[:,:]
        im2[:,:,1] = img[:,:]
        im2[:,:,2] = img[:,:]
        im2[:,:,3] = 1.0
    elif(Ncp==3):
        im2[:,:,0] = img[:,:,2]
        im2[:,:,1] = img[:,:,1]
        im2[:,:,2] = img[:,:,0]
        im2[:,:,3] = 1.0
    elif(Ncp>=4):
        im2[:,:,0] = img[:,:,2]
        im2[:,:,1] = img[:,:,1]
        im2[:,:,2] = img[:,:,0]
        im2[:,:,3] = img[:,:,3]
    else:
        print("save_img: no way to handle Ncp={}".format(Ncp))

    im2 = im2*255.0
    im2 = np.clip(im2,0.0,255.0)
    im3 = np.uint8(im2)

    try:
        cv2.imwrite(fname,im3)
    except:
        print("save_img: error: could not write to file {} fname".format(fname))


    return

def test_img_loadsave():

    img = load_img("./testimg.png")
    print(img.shape)


    plt.imshow(img.transpose(1,0,2))
    plt.show()

    save_img("./testout.png",img)

    img = np.random.rand(255,128)

    save_img("./testout2.png",img)

    img = np.random.rand(255,128,3)

    save_img("./testout3.png",img)


###
### Some vector math
###

def hodge_dual(v):
    """
    Hodge dual of a vector. The hodge dual is an operation in differentail forms that (in 3d)
    returns a bivector for a vector, the bivector being related to the generator of a rotation matrix.
    """
    v = np.array(v)

    M = np.zeros((3,3))
    M[0,1] = v[2]
    M[1,2] = v[0]
    M[2,0] = v[1]
    M[1,0] = -v[2]
    M[2,1] = -v[0]
    M[0,2] = -v[1]

    M = -M
    return M

def rotmat_axisangle(axis,angle):
    """
    rotmat_axisangle(axis,angle)
    axis - the axis of rotation
    angle - the angle [rad] to rotate about the axis
    returns: 3x3 rotation matrix
    """
    axis = np.array(axis)
    
    E = np.eye(3)
    H = hodge_dual(axis)
    H2 = np.matmul(H,H)

    R = E + H*sin(angle) + H2*(1.0-cos(angle))

    return R

def rotx(angle):
    R = np.array([[1,0,0],[0,cos(angle),-sin(angle)],[0,sin(angle),cos(angle)]])
    return R

def roty(angle):
    R = np.array([[cos(angle),0,sin(angle)],[0,1,0],[-sin(angle),0,cos(angle)]])
    return R

def rotz(angle):
    R = np.array([[cos(angle),-sin(angle),0],[sin(angle),cos(angle),0],[0,0,1]])
    return R

def rot_eulerxyz(x,y,z):

    Rx = rotx(x)
    Ry = roty(y)
    Rz = rotz(z)

    R = np.matmul(Rz,np.matmul(Ry,Rx))

    return R

def rot_cameraxyz(x,y,z):
    R1 = rotx(-90*pi/180)
    R2 = rot_eulerxyz(x,y,z)

    R = np.matmul(R2,R1)

    return R

def rot_eulerzyx(x,y,z):

    Rx = rotx(x)
    Ry = roty(y)
    Rz = rotz(z)
    R = np.matmul(Rx,np.matmul(Ry,Rz))

    return R

def extrinsic(R,t):

    R = np.array(R).copy()
    t = np.array(t).copy()

    t = np.reshape(t,[3,1])

    Rinv = np.linalg.inv(R)
    c = -np.matmul(Rinv,t)

    Ex = np.zeros([4,4])
    Ex[0:3,0:3] = Rinv[:,:]
    Ex[0:3,3] = c.flatten()
    Ex[3,3] = 1.0

    return Ex

def intrinsic_eqasp(fov):
    """
    Equal aspect ratio intrinsic matrix
    """
    In = np.zeros([4,4])
    alpha = fov/2.0
    fi = np.tan(alpha)

    In[0,0] = 1.0
    In[1,1] = -1.0 #invert y
    In[2,2] = fi
    In[3,3] = 1.0 #to allow inversion

    return In

def ray_to_normscreen(Ex,In,ray):
    """
    ray_to_normscreen(Ex,In,ray)
    Ex - 4x4 extrinsic (modelview) matrix
    In - 4x4 projection matrix
    ray - 3x1 ray vector to project to the normalized (-1,1) x (-1,1) screen coordinates
    of a camera with a given orientation, position, and field of view described by Ex and In
    """
    ray = ray.copy()
    rayp = np.zeros([4,1]) #projective ray
    rayp[0:3,0] = np.squeeze(ray)
    rayp[3,0] = 1.0

    r2 = np.matmul(Ex,rayp)
    r3 = np.matmul(In,r2)
    if(r3[2,0]>0.0):
        sx = r3[0,0]/r3[2,0]
        sy = r3[1,0]/r3[2,0]
    else:
        sx = -2.0
        sy = -2.0 #offscreen

    return [sx,sy]

def vnorm(vect):
    vect = np.array(vect)
    vect = vect.copy()

    snrm = np.sqrt(np.sum(vect*vect))
    return snrm

def vnormalize(vect):
    vect = np.array(vect)
    vect = vect.copy()
    snrm = vnorm(vect)
    if(snrm>0.0):
        vect = vect/snrm
    else:
        vect = np.zeros(vect.shape)
    
    return vect

def normscreen_to_ray(Ex,In,sx,sy):

    Ininv = np.linalg.inv(In)
    Exinv = np.linalg.inv(Ex)

    sp = np.zeros([4,1])
    sp[0] = sx
    sp[1] = sy
    sp[2] = 1.0
    sp[3] = 1.0
    rp1 = np.matmul(Ininv,sp)

    rp2 = np.squeeze(rp1[0:3,0])
    rp3 = vnormalize(rp2)

    rp4 = np.zeros([4,1])
    rp4[0:3,0] = rp3[:]

    rp5 = np.matmul(Exinv,rp4)
    ray = np.squeeze(rp5[0:3,0])

    return ray



###
### End vector math
###

def _local_makeint(istr):
    try:
        i = int(istr)
    except:
        i = 0
    return i

def imfilefilter_preload(scenedir,flst):

    fnames = []
    Nimgs = 0
    angles = []

    for f in flst:
        f2 = os.path.join(scenedir,f)
        fb = os.path.splitext(f)[0]
        fext = os.path.splitext(f)[1]

        if(fext.lower()==".png"):
            fsplit = fb.split("_")
            if(len(fsplit)==3):
                i1s = fsplit[1]
                i2s = fsplit[2]
                i1 = _local_makeint(i1s)
                i2 = _local_makeint(i2s)

                fnames.append(f2)
                angles.append([i1,0,i2])
                Nimgs += 1

    angles = np.float32(angles)
    
    return [Nimgs,angles,fnames]

def load_all_images(scenedir, verbose=True):

    flst = os.listdir(scenedir)
    [Nimgs,angles,fnames] = imfilefilter_preload(scenedir,flst)

    imgs = None
    Nx = 0
    Ny = 0
    for I in range(0,Nimgs):
        if(verbose==True):
            print("Loading {} of {}".format(I,Nimgs))

        img = load_img(fnames[I])
        if(I==0):
            Nx = img.shape[0]
            Ny = img.shape[1]
        else:
            if(img.shape[0]!=Nx or img.shape[1]!=Ny):
                print("Error: Image {}, {} has size {},{}, not {},{}".format(I,fnames[I],img.shape[0],img.shape[1],Nx,Ny))
                return None
        
        if(I==0):
            imgs = np.zeros((Nx,Ny,4,Nimgs),dtype=np.float32,order='F')
        
        imgs[:,:,:,I] = img[:,:,:]

    return [angles, imgs]


def ray_to_normscreen(Ex,In,ray)
    """
    ray_to_normscreen(Ex,In,ray)
    Is a given ray within the normalized screen coordinates of a camera described by Ex/In matrices?
    """
    sx = -2.0
    sy = -2.0
    byesno = False
    [sx,sy] = ray_to_normscreen(Ex,In,ray)
    if(sx>=-1.0 and sx<=1.0 and sy>=-1.0 and sy<=1.0):
        byesno = True

    return [byesno, sx, sy]

def check_cameraz(angles):
    #Debug script to check camera orientations
    #Camera z vectors should form a set covering the full spherical field of view if all the rotations are set up right

    Nimgs = angles.shape[0]

    zvs = np.zeros([3,Nimgs])
    for I in range(0,Nimgs):
        R = rot_cameraxyz(angles[I,0]*pi/180,angles[I,1]*pi/180,angles[I,2]*pi/180)
        print("{} {} {}".format(angles[I,0],angles[I,1],angles[I,2]))
        print(R)
        zv = np.squeeze(R[0:3,2])
        print(zv)
        zvs[:,I] = zv[:]
    
    # plt.figure()
    # plt.plot(angles[:,0],angles[:,2],'kd')

    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    plt.plot(zvs[0,:],zvs[1,:],zvs[2,:],'kd')
    ax.set_xlim([-1,1])
    ax.set_ylim([-1,1])
    ax.set_zlim([-1,1])
    plt.show()


    return

def process_pixel(px,py,angles,imgs,fov,Nxout,Nyout):
    """
    Eventually, figure out a way to vectorize this, or rewrite in C

    This takes a given output pixel px,py [0,Nxout) [0,Nyout)
    calculates a ray corresponding to the pixel for the lat-lon sky image
    and searches the available source images for which ones the ray lies within the field of view of.
    It pulls rgba from the source images and returns either an average or first available.
    """

    Nx = imgs.shape[0]
    Ny = imgs.shape[1]
    Nimgs = imgs.shape[3]

    Nxc = np.arange(0,Nx)
    Nyc = np.arange(0,Ny)

    rgba = np.zeros(4)
    rgba_N = 0

    rgba = np.float32([0,0,0,0])

    ax = px/(Nxout)*2.0*pi
    ay = -2.0*(py/(Nyout-1)-0.5)*pi/2

    ray = np.float32([cos(ax)*cos(ay),sin(ax)*cos(ay),sin(ay)])

    # rgba = np.float32([ray[0],ray[1],ray[2],1]) #debug
    # rgba = 0.5*rgba+0.5
    
    for I in range(0,Nimgs):


        R = rot_cameraxyz(angles[I,0]*pi/180,angles[I,1]*pi/180,angles[I,2]*pi/180)
        Ex = extrinsic(R,[0,0,0])
        In = intrinsic_eqasp(fov)

        [byesno,sx,sy] = iswithin_img(ray,Ex,In)

        #debug
        #rgba = np.float32([ray[0]*0.5+0.5,ray[1]*0.5+0.5,0.5*ray[2]+0.5,1])

        if(byesno):
            #interpolate images
            ra = imgs[:,:,0,I]
            ga = imgs[:,:,1,I]
            ba = imgs[:,:,2,I]
            aa = imgs[:,:,3,I]

            ppx = (sx+1.0)/2.0*(Nx-1.0)
            ppy = (-sy+1.0)/2.0*(Ny-1.0)
            
            ri = interp.interpn((Nxc,Nyc),ra,(ppx,ppy),bounds_error=False,fill_value=0.0)
            gi = interp.interpn((Nxc,Nyc),ga,(ppx,ppy),bounds_error=False,fill_value=0.0)
            bi = interp.interpn((Nxc,Nyc),ba,(ppx,ppy),bounds_error=False,fill_value=0.0)
            ai = interp.interpn((Nxc,Nyc),aa,(ppx,ppy),bounds_error=False,fill_value=0.0)
            
            
            rgba1 = np.squeeze(np.float32([ri,gi,bi,ai]))
            rgba = rgba + rgba1
            rgba_N += 1

            #debug
            #rgba = np.float32([0.5*sx+0.5,0.5*sy+0.5,0,1])
            #rgba = np.float32([angles[I,0]/360,0,0,1])

            break

    if(rgba_N>0):
        rgba = rgba/rgba_N


    return rgba

def proc_pixelrange(angles,sm_imgs,imgs_shape,fov,sm_imgout,imgout_shape,threadnum,nthreads):
    """
    Thread function for processing a range of pixels.
    angles - the list of camera angles
    sm_imgs - the shared memory buffer for the source images
    imgs_shape - the shape of the source images
    fov - the field of view angle for the camera source images [rad]
    sm_imgout - the shared memory buffer for the output images
    imgout_shape - the shape of the output image buffer
    threadnum - which thread is this?
    nthreads - total number of threads
    """

    Nxout = imgout_shape[0]
    Nyout = imgout_shape[1]

    #Thread indexing
    #Is - stride index - how many elements should each thread process?
    #I0,I1 - elements processed by this thread
    #N - total number of pixels
    N = Nxout*Nyout
    Is = int(np.floor(N/nthreads))+1
    I0 = threadnum*Is
    I1 = (threadnum+1)*Is
    if(I0>N):
        I0 = N
    if(I1>N):
        I1 = N

    #Unpack shared memory arrays
    s_imgs = np.ndarray(imgs_shape,dtype=np.float32,buffer=sm_imgs.buf)
    s_imgout = np.ndarray(imgout_shape,dtype=np.float32,buffer=sm_imgout.buf)

    for I in range(I0,I1):
        py = int(np.floor(I/Nxout))
        px = I%Nxout

        rgba = process_pixel(px,py,angles,s_imgs,fov,Nxout,Nyout)
        s_imgout[px,py,:] = rgba.flatten()[:] #save rgba to shared-memory output array


    return

def threaded_proc(angles,imgs,fov,imgout):
    """
    Thread manager for parallel processing.
    The actual threading library doesn't work for this purpose because of GIL
    multiprocessing does allow parallel execution, so "threads" are "multiprocessing" Processes

    Responsibility for painting in pixels is split among a number of threads.

    While this is faster than single threaded, it is still pretty slow by compiled code or GPU code standards, and
    I'll probably write a compiled submodule for this in the future. For portability's sake, though, here is the
    native python version of this.
    """
    nthreads = os.cpu_count()-2
    if(nthreads<=0):
        nthreads=1

    threads = []

    #create shared memory arrays
    sm_imgs = shared_memory.SharedMemory(create=True, size=imgs.nbytes)
    s_imgs = np.ndarray(imgs.shape,dtype=np.float32,buffer=sm_imgs.buf)
    s_imgs[:] = imgs[:]

    sm_imgout = shared_memory.SharedMemory(create=True, size=imgout.nbytes)
    s_imgout = np.ndarray(imgout.shape,dtype=np.float32,buffer=sm_imgout.buf)


    for I in range(0,nthreads):
        t = mp.Process(target=proc_pixelrange,args=(angles,sm_imgs,s_imgs.shape,fov,sm_imgout,s_imgout.shape,I,nthreads))
        threads.append(t)
    
    for I in range(0,nthreads):
        threads[I].start()

    for I in range(0,nthreads):
        threads[I].join()
        print("{} %% complete".format((I+1)/nthreads*100))

    imgout[:] = s_imgout[:]

    sm_imgs.close()
    sm_imgs.unlink()
    sm_imgout.close()
    sm_imgout.unlink()

    return imgout

def main():

    [angles,imgs] = load_all_images("./scenerenders")
    fnout = "./scenehdr2.png"

    Nxout = 200 #Dimensions of the output scene image
    Nyout = int(np.floor(Nxout/2.0))

    fov = 50.0*pi/180.0 #field of view angle of the renders in [radians]
    
    imgout = np.zeros((Nxout,Nyout,4),dtype=np.float32)

    # single threaded implementation
    # for I in range(0,Nxout):
    #     print("processing {:1.2f} %% complete.".format(I/Nxout*100.0))
    #     for J in range(0,Nyout):
    #         rgba = process_pixel(I,J,angles,imgs,fov,Nxout,Nyout)
    #         imgout[I,J,:] = rgba[:]
    
    imgout = threaded_proc(angles,imgs,fov,imgout)

    save_img(fnout,imgout)

    return



if(__name__=="__main__"):
    main()
    #test_eulermatrices()

    exit(0)
