import argparse
import ipyparallel as ipp
import numpy as np
import os
import cv2 as cv2
import time
from retinatools.library import image_to_cone

#Helper function
def RF_FilterFrame(idx, luminance_frame, r_RF):
    """
    Helper function to do the parallelized filtering with RFs.
    """
    from retinatools.library import get_RF

    X = np.linspace(-128, 128, 256)
    Y = np.linspace(-128, 128, 256)
    idx_x = idx%256
    idx_y = np.int(np.floor(idx/256))
    x, y = np.meshgrid(X,Y)    
    xoff = X[idx_x]
    yoff = Y[idx_y]
    RF = get_RF(xoff, yoff, r_RF)
    product = np.sum(luminance_frame * RF)
    return product

def parallel_prediction(idx_x, RFstrip, fitparameters, celltype):    
    ''' 
    Helper function to do the parallelized processing of the filtered signal; 
    This transform center/surround signal to activity with standard parameters.
    '''
    from retinatools.library import get_MC_response

    signal_center   = RFstrip[0, idx_x, :]
    signal_surround = RFstrip[1, idx_x, :]
    R = get_MC_response(signal_center, signal_surround, fitparameters, celltype)
    return R


fitparameters = {'rcenter': 4.0, 'rsurr': 13.0, 'ksurr': 0.45, 'surround_delay': 3, \
                 'lum_factor': 60, 'tau1': 7.3, 'tau2': 38, 'k1': 0.019, 'c1': 0.016, 'tauplus': 110, \
                 'tauA': 20, 'k2': 50, 'dt': 1000/150.}  # for 299  Mon
celltype = "on-cell"

# fitparameters = {'rcenter': 4.5, 'rsurr': 20, 'ksurr': 0.34, 'surround_delay': 8, \
#                     'lum_factor': 1000, 'tau1': 5.0, 'tau2': 16, 'k1': 0.00048, 'c1': 0.0024, 'tauplus': 22, \
#                     'tauA': 33, 'k2': 560, 'dt': 1000/150.}  # For 078Moff.txt
# celltype = "off-cell"
savepath = "/jukebox/tank/schottdorf/tmp/"

N_MAXFRAME = 4500  # 4500 to only process first 2.5 min *60sec *30fps in the mpg -> 30sec ret video

def main(profile):
    rc = ipp.Client(profile=profile)                 # ipyparallel status
    print(rc.ids)

    dview = rc[:]                                    # Set up the engines
    dview.use_cloudpickle()                          # Make sure to fix the closure over luminance_frame
    with dview.sync_imports():
        import numpy as np

    ################################
    #####   Process the movie  #####
    ################################
    vidcap = cv2.VideoCapture('/usr/people/ms81/cells/1x10_256.mpg')
    data = np.ones((2, 256, 256, N_MAXFRAME))            # for center & surround   X 256px  X 256px  X N_frames
    counter = 0
    success = True
    t0 = time.time()
    while success and (counter < N_MAXFRAME):   
        success, image = vidcap.read()
        scone, mcone, lcone = image_to_cone(image)        
        luminance_frame = mcone + lcone
        for tt in ["center", "surr"]:
            if tt == "center":
                r_RF = fitparameters['rcenter']
                luminance_frame = luminance_frame
                idx = 0
            else:
                r_RF = fitparameters['rsurr']
                luminance_frame = luminance_frame
                idx = 1
            rc[:].push(dict( luminance_frame = luminance_frame, RF_FilterFrame = RF_FilterFrame, r_RF = r_RF ))         # Seed namespace for the ipython engines. 
            tmp = rc[:].map_sync( lambda ix: RF_FilterFrame(ix, luminance_frame, r_RF), range(0,256*256) )              # Execute parallel computation on all ipython engines
            data[idx,:,:,counter] = np.reshape(tmp, (256,256))                                                          # Reorganize data back to frame
        counter += 1
        print(counter/N_MAXFRAME)
    t1 = time.time()

    # filename = "rgc-{0}-raw.npy".format(os.environ["SLURM_JOB_ID"])
    # np.save(filename, data)
    print("Part 1 DONE; Time needed [in h]:")
    print( (t1-t0)/3600 )


    ################################
    ####  Apply temp. process  #####
    ################################
    t1 = time.time()
    output_model = np.zeros((256, 256, N_MAXFRAME))
    rc[:].push(dict( parallel_prediction = parallel_prediction))
    for idx_x in range(0,256):
        RFstrip = data[:,idx_x,:,:]
        rc[:].push(dict( RFstrip = RFstrip, celltype = celltype, fitparameters = fitparameters, parallel_prediction = parallel_prediction ))
        tmp = rc[:].map_sync(lambda idx_y: parallel_prediction(idx_y, RFstrip, fitparameters, celltype), range(0,256))
        output_model[idx_x,:,:] = np.reshape(tmp, (256, N_MAXFRAME))
        print(idx_x/256)
    t2 = time.time()

    filename = "-{0}-frames.npy".format(os.environ["SLURM_JOB_ID"])
    pp = savepath + celltype + filename
    np.save(pp, output_model)
    print("Part 2 DONE; Time needed [in h]:")
    print("Saved: " + pp)
    print((t2-t1)/3600)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--profile", required=True, help="Name of IPython profile to use")

    args = parser.parse_args()
    main(args.profile)
    print("Done and done.")