# -*- coding: utf-8 -*-
"""
Created on Mon Nov  6 15:32:24 2017

@author: vidyag
"""

import tensorflow as tf
import numpy as np
import math
from scipy import misc
import glob

from autoencode import AutoEncoder
from layer_defs import variable_on_cpu, getConvInitializer

def fftshift(mat2D, dim0, dim1): #fftshift == ifftshift when dimensions are all even
                                 #fftshift only works with even dimensions

    if (dim0==1) and (dim1==1):
        return mat2D

    if (dim0%2) or (dim1%2):
        raise ValueError('Dimensions must be even to use fftshift.')

    dim0=tf.cast(dim0,tf.int32)
    dim1=tf.cast(dim1,tf.int32)

    piece1=tf.slice(mat2D,[0,0],[dim0//2,dim1//2])
    piece2=tf.slice(mat2D,[0,dim1//2],[dim0//2,dim1//2])
    piece3=tf.slice(mat2D,[dim0//2,0],[dim0//2,dim1//2])
    piece4=tf.slice(mat2D,[dim0//2,dim1//2],[dim0//2,dim1//2])

    top=tf.concat([piece4,piece3],axis=1)
    bottom=tf.concat([piece2,piece1],axis=1)

    final=tf.concat([top,bottom],axis=0)
    return final

def propTF_withNA_PSFcustom(u1,H_fresnel,H_NA,PSF,m,incoherent=1):
    #u1 is the source plane field
    #L is the side length of the observation and source fields (assume square fields)

    #NA is the numerical aperture
    #m is u1.shape[0], is a tf.int32
    #dx is L/m

    PSF_phase=2*math.pi*PSF
    PSF_phase=tf.cast(tf.cos(PSF_phase),tf.complex64)-1j*tf.cast(tf.sin(PSF_phase),tf.complex64)

#    PSF_phase=2*tf.constant(math.pi, dtype=tf.complex64)*tf.cast(PSF, tf.complex64) #/tf.cast(tf.reduce_max(PSF), tf.complex64)
#    with tf.device('/cpu:0'):
#        PSF_phase=tf.exp(-1j*PSF_phase)

    PSF_phase=fftshift(PSF_phase,m,m)

    H = H_fresnel*H_NA*PSF_phase

    if incoherent:
        #H=np.fft.ifft2(np.abs(np.fft.fft2(H))**2)
        #H=tf.ifft2d(tf.fft2d(H)*tf.conj(tf.fft2d(H)))
        H=tf.fft2d(tf.ifft2d(H)*tf.conj(tf.ifft2d(H)))
        #H=H/H[0,0]

        U1=fftshift(u1,m,m)
        U1=tf.fft2d(U1)

        U2=H*U1
        u2=fftshift(tf.ifft2d(U2),m,m)

    else:
        U1=fftshift(u1,m,m)
        U1=tf.fft2d(U1)

        U2=H*U1
        u2=fftshift(tf.ifft2d(U2),m,m)
        u2 = u2*tf.conj(u2) # make into intensity object

    return u2

def change_filter_function(u1,H_old,H_new,Nx,reg=1e-10,incoherent=1):

    #u1 is the source plane field
    #Nx is u1.shape[0]
    #dx is L/m
    #H_old and H_new are already fftshifted in Fourier space

    if incoherent:
        H_old = tf.fft2d(tf.ifft2d(H_old)*tf.conj(tf.ifft2d(H_old)))
        H_new = tf.fft2d(tf.ifft2d(H_new)*tf.conj(tf.ifft2d(H_new)))
        #H=H/H[0,0]

        U1=fftshift(u1,Nx,Nx)
        U1=tf.fft2d(U1)

#        U2=H_old*U1 # previous processing
        U2 = (H_new/(H_old+reg))*U1

        u2=fftshift(tf.ifft2d(U2),Nx,Nx)

    else:
        U1=fftshift(u1,Nx,Nx)
        U1=tf.fft2d(U1)

        U2=(H_new/(H_old+reg))*U1

#        U2=H_old*U1 # previous processing
        u2=fftshift(tf.ifft2d(U2),Nx,Nx) # field value
#        u2 = u2*np.conj(u2) # make into intensity object

    return u2

def create_phase_obj_stack(PSF, trainingSet, batch_size, H_fresnel, H_NA, Nx):
    for ii in range(batch_size):

        # low_res_obj is an intensity object
        low_res_obj = propTF_withNA_PSFcustom(trainingSet[ii,:,:],H_fresnel,H_NA,PSF,Nx,incoherent=0)

        low_res_obj = tf.expand_dims(low_res_obj,axis=0)

        if ii == 0:
            low_res_obj_stack = low_res_obj
        else:
            low_res_obj_stack = tf.concat([low_res_obj_stack,low_res_obj],0)

    low_res_obj_stack = tf.cast(low_res_obj_stack, tf.float32)

    return low_res_obj_stack

def create_microscope_img(PSF,trainingSet_sample,Nz,H_fresnel_stack,H_NA,m,num_wavelengths):
    # trainingSet_sample = trainingSet[ii,:,:,:]

    microscopeImg=[]

    for zInd in range(Nz):
        for waveInd in range(num_wavelengths):
            add_layer=propTF_withNA_PSFcustom(trainingSet_sample[:,:,zInd,waveInd],H_fresnel_stack[:,:,zInd,waveInd],H_NA[:,:,waveInd],PSF,m)
            microscopeImg.append(add_layer)

    microscopeImg = tf.add_n(microscopeImg)
    microscopeImg = tf.expand_dims(microscopeImg,axis=0)
    microscopeImg = tf.cast(microscopeImg, tf.float32)
    return microscopeImg

def create_microscopeImgStack(PSF,trainingSet,Nz,batch_size,H_fresnel_stack,H_NA,m,num_wavelengths): #creates microscopeImg for every example in the trainingSet

    for ii in range(batch_size):
        microscopeImg = create_microscope_img(PSF,trainingSet[ii,:,:,:,:],Nz,H_fresnel_stack,H_NA,m,num_wavelengths)
        if ii == 0:
            microscopeImgStack = microscopeImg
        else:
            microscopeImgStack = tf.concat([microscopeImgStack,microscopeImg],0)

    return microscopeImgStack

def add_noise_microscopeImgStack(microscopeImgStack,normalRandomMat1,normalRandomMat2,sqrt_reg,\
                                 poisson_noise_multiplier, gaussian_noise_multiplier, batch_size, library=tf):

    #XXX Fix the Poisson noise for low photon levels

    if library == tf:
        multiplierPoisson = tf.constant(poisson_noise_multiplier,dtype=tf.float32) #6e3 for EPFL
        multiplierGaussian = tf.constant(gaussian_noise_multiplier,dtype=tf.float32)
    else:
        multiplierPoisson = poisson_noise_multiplier
        multiplierGaussian = gaussian_noise_multiplier

    microscopeImgStack2 = microscopeImgStack*multiplierPoisson

    microscopeImgStack3=library.sqrt(library.abs(microscopeImgStack2)+sqrt_reg)*normalRandomMat1+microscopeImgStack2
    microscopeImgStack4=microscopeImgStack3+multiplierGaussian*normalRandomMat2

    zeros = library.zeros([batch_size, microscopeImgStack4.shape[1], microscopeImgStack4.shape[2]], dtype=library.float32)

    microscopeImgStack = library.where(microscopeImgStack4<0,zeros,microscopeImgStack4) #truncate below 0

    microscopeImgStack = microscopeImgStack/multiplierPoisson
    return microscopeImgStack

def F(mat2D,dim0,dim1):
    return fftshift(tf.fft2d(mat2D),dim0,dim1)

def Ft(mat2D,dim0,dim1):
    return tf.ifft2d(fftshift(mat2D,dim0,dim1))

downsamp = lambda x,cen,Np:  x[cen[0]-Np//2:cen[0]-Np//2+Np, \
    cen[1]-Np//2:cen[1]-Np//2+Np]

def sigmoid_stretch(x, stretch, library=tf):
    y = 1. / (1. + library.exp(-x/stretch))
    return y

def HiToLoPatch_singleLED(obj,scale_multiply, Ns, scale, cen0, P, Np, N_obj, LED_i): #stretch

    illumination_weight = scale[LED_i] * scale_multiply[LED_i]

    cen = (cen0-Ns[LED_i,:]).astype(int)
    O=F(obj,N_obj,N_obj)

    Psi0 = downsamp(O,cen,Np)*P

    psi0 = Ft(Psi0,Np,Np) #low resolution field
    intensity_i = psi0*tf.conj(psi0)*tf.cast(illumination_weight, tf.complex64)

    return intensity_i

def HiToLo_singleLED(obj, N_obj, N_patch, scale_multiply, num_patches, Ns_mat, scale_mat, \
           cen0, P, Np, LED_i):

    low_res_patches=[]

    count = 0
    for i,startX in enumerate(np.arange(0,N_obj,N_patch)):
        for j,startY in enumerate(np.arange(0,N_obj,N_patch)):

            # pass the full object to HiToLoPatch
            Ns = Ns_mat[count,:,:]
            scale = scale_mat[count,:]
            low_res_patch_everything = HiToLoPatch_singleLED(obj,scale_multiply, Ns, scale, cen0, P, Np, N_obj, LED_i)
            low_res_patches.append(low_res_patch_everything)
            count += 1

    count = 0

    for i,startX in enumerate(np.arange(0,Np,Np/num_patches)):
        for j,startY in enumerate(np.arange(0,Np,Np/num_patches)):
            # Extract out patch of interest
            low_res_patch=tf.slice(low_res_patches[count],[int(startX),int(startY)],[int(N_patch*Np/N_obj),int(N_patch*Np/N_obj)])
            if j==0:
                low_res_obj_row=low_res_patch
            else:
                low_res_obj_row = tf.concat([low_res_obj_row,low_res_patch],axis=1)

            count += 1

        if i==0:
            low_res_obj = low_res_obj_row
        else:
            low_res_obj = tf.concat([low_res_obj,low_res_obj_row],axis=0)

    low_res_obj = tf.cast(low_res_obj,tf.float32)
    return low_res_obj

def HiToLoPatch(obj,scale_multiply, Ns, scale, cen0, P, H0, Np, N_obj, numLEDs): #stretch

    illumination_weights = scale * scale_multiply #sigmoid_stretch(scale_multiply,stretch) #scale_multiply

    for LED_i in range(numLEDs):
        cen = (cen0-Ns[LED_i,:]).astype(int)
        O=F(obj,N_obj,N_obj)

        Psi0 = downsamp(O*H0,cen,Np)*P

        psi0 = Ft(Psi0,Np,Np) #low resolution field
        intensity_i = psi0*tf.conj(psi0)*tf.cast(illumination_weights[LED_i], tf.complex64)

        if LED_i == 0:
            low_res_patch = intensity_i
        else:
            low_res_patch = low_res_patch + intensity_i

    return low_res_patch

def HiToLo(obj, N_obj, N_patch, scale_multiply, num_patches, Ns_mat, scale_mat, \
           cen0, P, H0, Np, numLEDs): #stretch

    low_res_patches=[]

    count = 0
    for i,startX in enumerate(np.arange(0,N_obj,N_patch)):
        for j,startY in enumerate(np.arange(0,N_obj,N_patch)):

            # pass the full object to HiToLoPatch
            Ns = Ns_mat[count,:,:]
            scale = scale_mat[count,:]
            low_res_patch_everything = HiToLoPatch(obj,scale_multiply, Ns, scale, cen0, P, H0, Np, N_obj, numLEDs) #stretch
            low_res_patches.append(low_res_patch_everything)
            count += 1

    count = 0

    for i,startX in enumerate(np.arange(0,Np,Np/num_patches)):
        for j,startY in enumerate(np.arange(0,Np,Np/num_patches)):
            # Extract out patch of interest
            low_res_patch=tf.slice(low_res_patches[count],[int(startX),int(startY)],[int(N_patch*Np/N_obj),int(N_patch*Np/N_obj)])
            if j==0:
                low_res_obj_row=low_res_patch
            else:
                low_res_obj_row = tf.concat([low_res_obj_row,low_res_patch],axis=1)

            count += 1

        if i==0:
            low_res_obj = low_res_obj_row
        else:
            low_res_obj = tf.concat([low_res_obj,low_res_obj_row],axis=0)

    return low_res_obj

def upsample(low_res_obj, Np, N_obj):

    if Np == 1:
        dense_multiply = tf.cast(tf.squeeze(getConvInitializer(N_obj, N_obj, init_type="trunc_norm")), tf.complex64)
        dense_bias = tf.cast(tf.squeeze(getConvInitializer(N_obj, N_obj, init_type="trunc_norm")), tf.complex64)

#        dense_multiply_i = tf.cast(tf.squeeze(getConvInitializer(N_obj, N_obj, init_type="trunc_norm")), tf.complex64)
#        dense_bias_i = tf.cast(tf.squeeze(getConvInitializer(N_obj, N_obj, init_type="trunc_norm")), tf.complex64)

#        pad0 = (N_obj - Np)//2 -1
#        pad1 = pad0 + 3
#        upsampled_obj = tf.pad(upsampled_obj, [[int(pad0),int(pad1)],[int(pad0),int(pad1)]], 'CONSTANT', \
#                                                constant_values = 100)
#        upsampled_obj_real = low_res_obj*dense_multiply_r + dense_bias_r
#        upsampled_obj_imag = low_res_obj*dense_multiply_i + dense_bias_i
        upsampled_obj = low_res_obj*dense_multiply + dense_bias
#        upsampled_obj = tf.Print(upsampled_obj, [dense_multiply], message='dense_multiply: ')
#        upsampled_obj = tf.Print(upsampled_obj, [dense_bias], message='dense_bias: ')
#        upsampled_obj = tf.Print(upsampled_obj, [low_res_obj], message='low_res_obj: ')

    else:
        upsampled_obj = F(low_res_obj,Np,Np)
        pad = (N_obj - Np)/2
        upsampled_obj = tf.pad(upsampled_obj, [[int(pad),int(pad)],[int(pad),int(pad)]], 'CONSTANT')
        upsampled_obj = Ft(upsampled_obj,N_obj,N_obj)

    return upsampled_obj

def create_upsampled_obj_stack(low_res_obj_stack, batch_size, Np, N_obj):
    low_res_obj_stack = tf.cast(low_res_obj_stack, tf.complex64)
    for ii in range(batch_size):
        upsampled_obj = upsample(low_res_obj_stack[ii,:,:], Np, N_obj)
        upsampled_obj = tf.expand_dims(upsampled_obj,axis=0)
        if ii == 0:
            upsampled_obj_stack = upsampled_obj
        else:
            upsampled_obj_stack = tf.concat([upsampled_obj_stack,upsampled_obj],0)

    upsampled_obj_stack = tf.cast(upsampled_obj_stack, tf.float32)
    return upsampled_obj_stack

def create_FP_img_stack(scale_multiply,trainingSet,batch_size, N_obj, N_patch, num_patches, Ns_mat, \
                        scale_mat, cen0, P, H0, Np, numLEDs):
    for ii in range(batch_size):
        low_res_obj = HiToLo(trainingSet[ii,:,:], N_obj, N_patch, scale_multiply, \
                             num_patches, Ns_mat, scale_mat, \
                             cen0, P, H0, Np, numLEDs)

        low_res_obj = tf.expand_dims(low_res_obj,axis=0)

        if ii == 0:
            low_res_obj_stack = low_res_obj
        else:
            low_res_obj_stack = tf.concat([low_res_obj_stack,low_res_obj],0)

    low_res_obj_stack = tf.cast(low_res_obj_stack, tf.float32)
    return low_res_obj_stack

def calculate_loss_FP(predicted_mat, trainingSet, library = tf):

    if library == tf:
        sum_func = tf.reduce_sum
    else:
        sum_func = np.sum

    loss_l2 = sum_func(library.real((predicted_mat - trainingSet)*library.conj(predicted_mat - trainingSet)))
#    loss_l1 = sum_func(library.abs(predicted_mat - trainingSet))

    return loss_l2

def grad_diff_loss(predicted_mat, trainingSet, library = tf):

    if library == tf:
        sum_func = tf.reduce_sum
    else:
        sum_func = np.sum

    diff_x_actual = trainingSet[:,1:,:]-trainingSet[:,:-1,:]
    diff_y_actual = trainingSet[:,:,1:]-trainingSet[:,:,:-1]

    diff_x_guess = predicted_mat[:,1:,:]-predicted_mat[:,:-1,:]
    diff_y_guess= predicted_mat[:,:,1:]-predicted_mat[:,:,:-1]

    loss_x = sum_func(library.real((diff_x_actual-diff_x_guess)*library.conj(diff_x_actual-diff_x_guess)))
    loss_y = sum_func(library.real((diff_y_actual-diff_y_guess)*library.conj(diff_y_actual-diff_y_guess)))

    return loss_x + loss_y

def convert_net_prediction_list(net_prediction_list, image_modality, batch_size, optical_parameters_dict):
    if (image_modality == 'FP') or (image_modality == 'phase') or (image_modality == 'FP_PP'):

        predicted_mat = tf.cast(net_prediction_list[0],tf.complex64)+1j*tf.cast(net_prediction_list[1],tf.complex64)
#        predicted_mat = tf.cast(net_prediction_list[0],tf.complex64)*tf.exp(1j*tf.cast(net_prediction_list[1],tf.complex64))

        predicted_mat = tf.squeeze(predicted_mat,axis=3)

    elif image_modality == 'STORM':
        Nx = optical_parameters_dict['Nx_highres']
        Nz = optical_parameters_dict['Nz']
        num_wavelengths = optical_parameters_dict['num_wavelengths']

        predicted_mat = tf.stack(net_prediction_list,axis=3)
        predicted_mat = tf.reshape(predicted_mat, [batch_size, Nx, Nx, Nz, num_wavelengths])

    elif image_modality == 'HE':
        predicted_mat = tf.stack(net_prediction_list,axis=3)

        predicted_mat = tf.squeeze(predicted_mat, axis = 4)

    return predicted_mat

def iterative_solver_FP(predicted_mat, optical_element, batch_size, N_obj, N_patch, num_patches, Ns_mat,\
                        scale_mat, cen0, P, H0, Np, numLEDs, low_res_obj_stack, step_size, max_internal_iter, \
                        merit_stopping_point, loss_low_res, i):

#    print 'mii2'
    predicted_mat_real = tf.real(predicted_mat)
    predicted_mat_imag = tf.imag(predicted_mat)

    def single_iter(predicted_mat_real, predicted_mat_imag, loss_low_res, i):

        predicted_mat = tf.cast(predicted_mat_real,tf.complex64)+1j*tf.cast(predicted_mat_imag,tf.complex64)
        low_res_obj_predicted = create_FP_img_stack(optical_element,predicted_mat,batch_size, N_obj, N_patch, num_patches, Ns_mat, \
                        scale_mat, cen0, P, H0, Np, numLEDs) #stretch

        loss_low_res = tf.reduce_sum(tf.square(low_res_obj_stack - low_res_obj_predicted))

#        loss_MSE = calculate_loss_FP(low_res_obj_predicted, low_res_obj_stack)
#        loss_grad = grad_diff_loss(low_res_obj_predicted, low_res_obj_stack)

        #find gradient of loss_low_res with respect to predicted_mat
        predicted_mat_real_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat_real),axis=0)
        predicted_mat_imag_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat_imag),axis=0)

#        predicted_mat_real_gradient = tf.Print(predicted_mat_real_gradient, [loss_low_res], message='loss_low_res: ')

        #update predicted_mat with step_size
        predicted_mat_real = predicted_mat_real - step_size*predicted_mat_real_gradient
        predicted_mat_imag = predicted_mat_imag - step_size*predicted_mat_imag_gradient

        i+=1

        return predicted_mat_real, predicted_mat_imag, loss_low_res, i

    def convergence_cond(predicted_mat_real, predicted_mat_imag, loss_low_res,i):
        result = True

#        result = tf.Print(result, [result], message='result1: ')

        result = tf.cond(i>=max_internal_iter, lambda: False, lambda: result)

#        result = tf.Print(result, [result], message='result2: ')

        result = tf.cond(loss_low_res<merit_stopping_point, lambda: False, lambda: result)

#        result = tf.Print(result, [result], message='result3: ')

        return result

    [predicted_mat_real, predicted_mat_imag, loss_low_res, i] = tf.while_loop(convergence_cond,
                                                                              single_iter,
                                                                              [predicted_mat_real, predicted_mat_imag, loss_low_res,i])

    predicted_mat = tf.cast(predicted_mat_real,tf.complex64)+1j*tf.cast(predicted_mat_imag,tf.complex64)

    return predicted_mat

def iterative_solver_STORM(predicted_mat, PSF, Nz,batch_size, H_fresnel_stack, \
                           H_NA, Nx, low_res_obj_stack, step_size, max_internal_iter, \
                           merit_stopping_point, loss_low_res, i, num_wavelengths):

    def single_iter(predicted_mat, loss_low_res, i):

        predicted_mat_complex=tf.cast(predicted_mat, tf.complex64)
        low_res_obj_predicted = create_microscopeImgStack(PSF,predicted_mat_complex,Nz,batch_size,H_fresnel_stack,H_NA,Nx,num_wavelengths)

        loss_low_res = tf.reduce_sum(tf.square(low_res_obj_stack - low_res_obj_predicted))

        #find gradient of loss_low_res with respect to predicted_mat
        predicted_mat_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat),axis=0)

    #            predicted_mat_gradient = tf.Print(predicted_mat_gradient, [loss_low_res], message='loss_low_res: ')

        #update predicted_mat with step_size
        predicted_mat = predicted_mat - step_size*predicted_mat_gradient

        i+=1

        return predicted_mat, loss_low_res, i

    def convergence_cond(predicted_mat, loss_low_res,i):
        result = True
        result = tf.cond(i>=max_internal_iter, lambda: False, lambda: result)
        result = tf.cond(loss_low_res<merit_stopping_point, lambda: False, lambda: result)
        return result

    [predicted_mat, loss_low_res, i] = tf.while_loop(convergence_cond,
                                                    single_iter,
                                                    [predicted_mat, loss_low_res,i])
    return predicted_mat

def iterative_solver_phase(predicted_mat, PSF, batch_size, H_fresnel, \
                           H_NA, Nx, low_res_obj_stack, step_size, max_internal_iter, \
                           merit_stopping_point, loss_low_res, i):

    predicted_mat_real = tf.real(predicted_mat)
    predicted_mat_imag = tf.imag(predicted_mat)

    def single_iter(predicted_mat_real, predicted_mat_imag, loss_low_res, i):

        predicted_mat = tf.cast(predicted_mat_real,tf.complex64)+1j*tf.cast(predicted_mat_imag,tf.complex64)
        low_res_obj_predicted = create_phase_obj_stack(PSF, predicted_mat, batch_size, H_fresnel, H_NA, Nx)

        loss_low_res = tf.reduce_sum(tf.square(low_res_obj_stack - low_res_obj_predicted))

        #find gradient of loss_low_res with respect to predicted_mat
        predicted_mat_real_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat_real),axis=0)
        predicted_mat_imag_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat_imag),axis=0)

        #update predicted_mat with step_size
        predicted_mat_real = predicted_mat_real - step_size*predicted_mat_real_gradient
        predicted_mat_imag = predicted_mat_imag - step_size*predicted_mat_imag_gradient

        i+=1

        return predicted_mat_real, predicted_mat_imag, loss_low_res, i

    def convergence_cond(predicted_mat_real, predicted_mat_imag, loss_low_res,i):
        result = True
        result = tf.cond(i>=max_internal_iter, lambda: False, lambda: result)
        result = tf.cond(loss_low_res<merit_stopping_point, lambda: False, lambda: result)
        return result

    [predicted_mat_real, predicted_mat_imag, loss_low_res, i] = tf.while_loop(convergence_cond,
                                                                              single_iter,
                                                                              [predicted_mat_real, predicted_mat_imag, loss_low_res,i])

    predicted_mat = tf.cast(predicted_mat_real,tf.complex64)+1j*tf.cast(predicted_mat_imag,tf.complex64)

    return predicted_mat

def iterative_solver_HE(predicted_mat, Nx_highres, h_blur, \
                       high_magn, low_magn, dpix_c, wavelength, \
                       low_NA, low_res_obj_batch, step_size, max_internal_iter, \
                       merit_stopping_point, loss_low_res, i, \
                       batch_size):

    def single_iter(predicted_mat, loss_low_res, i):

        low_res_obj_batch_predicted = change_magn_batch(predicted_mat, Nx_highres, h_blur, \
                                                        high_magn, low_magn, dpix_c, wavelength, low_NA, batch_size)

        low_res_obj_batch_predicted = tf.cast(low_res_obj_batch_predicted, tf.float32)

        loss_low_res = tf.reduce_sum(tf.square(low_res_obj_batch - low_res_obj_batch_predicted))

        #find gradient of loss_low_res with respect to predicted_mat

        predicted_mat_gradient = tf.squeeze(tf.gradients(loss_low_res,predicted_mat),axis=0)

#            predicted_mat_gradient = tf.Print(predicted_mat_gradient, [predicted_mat_gradient], message='predicted_mat_gradient: ')
#            predicted_mat_gradient = tf.Print(predicted_mat_gradient, [loss_low_res], message='loss_low_res: ')

        #update predicted_mat with step_size
        predicted_mat = predicted_mat - step_size*predicted_mat_gradient

        i+=1

        return predicted_mat, loss_low_res, i

    def convergence_cond(predicted_mat, loss_low_res,i):
        result = True
        result = tf.cond(i>=max_internal_iter, lambda: False, lambda: result)
        result = tf.cond(loss_low_res<merit_stopping_point, lambda: False, lambda: result)
        return result

    [predicted_mat, loss_low_res, i] = tf.while_loop(convergence_cond,
                                                     single_iter,
                                                     [predicted_mat, loss_low_res,i])

    return predicted_mat, loss_low_res

def find_predicted_mat(image_modality, num_nets, input_layer_list, Nx, batch_size, layers_dropout, dropout_prob,
                       use_batch_norm, autoencode_init_type, init_type_bias, init_type_resid, kernel_multiplier, variance_reg, training, \
                       num_layers_autoencode, skip_interval, num_blocks, optical_parameters_dict):

    net_prediction_list = []
    for net in range(num_nets):
        with tf.variable_scope("net_" + str(net)):

            curr_input = input_layer_list[net]
            for i in range(num_blocks):
                with tf.variable_scope("ae_" + str(i)):
                    curr_net = AutoEncoder(curr_input,
                                           num_layers_autoencode, Nx,
                                           batch_size, training,
                                           kernel_multiplier,
                                           skip_interval=skip_interval,
                                           conv_activation='maxout',
                                           deconv_activation='maxout',
                                           dropout_count=layers_dropout,
                                           dropout_prob=dropout_prob,
                                           create_graph_viz=False,
                                           use_batch_norm=use_batch_norm,
                                           init_type=autoencode_init_type,
                                           init_type_bias=init_type_bias,
                                           init_type_resid=init_type_resid)
                    curr_input = curr_net.get_prediction()
            net0_prediction = curr_input
            net_prediction_list.append(net0_prediction)

            """
            with tf.variable_scope("ae_0"):
                net_init = AutoEncoder(input_layer_list[net], num_layers_autoencode, Nx, batch_size, training, \
                                           kernel_multiplier, skip_interval=skip_interval, conv_activation='maxout', \
                                           deconv_activation='maxout', dropout_count=layers_dropout, dropout_prob=dropout_prob, \
                                           create_graph_viz=False, use_batch_norm=use_batch_norm, \
                                           init_type=autoencode_init_type, init_type_bias=init_type_bias, init_type_resid=init_type_resid)
                net_init_prediction = net_init.get_prediction()

            with tf.variable_scope("ae_1"):
                net0 = AutoEncoder(net_init_prediction, num_layers_autoencode, Nx, batch_size, training, \
                                       kernel_multiplier, skip_interval=skip_interval, conv_activation='maxout', \
                                       deconv_activation='maxout', dropout_count=layers_dropout, dropout_prob=dropout_prob, \
                                       create_graph_viz=False, use_batch_norm=use_batch_norm, \
                                       init_type=autoencode_init_type, init_type_bias=init_type_bias, init_type_resid=init_type_resid)
                net0_prediction = net0.get_prediction()
                net_prediction_list.append(net0_prediction)
            """

    ### End Neural Network(s)

    ### Convert net_prediction_list to predicted_mat

    predicted_mat = convert_net_prediction_list(net_prediction_list, image_modality, batch_size, optical_parameters_dict)

    return predicted_mat

def tower_loss_all(trainingSet, training, normalRandomMat1, normalRandomMat2, \
               initialize_optical_element_ones, num_elements, use_batch_norm, variance_reg, add_noise, sqrt_reg, \
               dropout_prob, layers_dropout, batch_size, max_internal_iter, merit_stopping_point, optical_parameters_dict, \
               autoencode_init_type, init_type_bias, init_type_resid, \
               kernel_multiplier, \
               poisson_noise_multiplier, gaussian_noise_multiplier, image_modality,
               num_layers_autoencode, skip_interval, num_blocks, \
               training_data_folder, load_optical_element_init,
               lowres_trainingSet=None):

    with tf.variable_scope("optical_transform"):

        Nx = optical_parameters_dict['Nx_highres']
        Nz = optical_parameters_dict['Nz']

        if initialize_optical_element_ones:
            if (image_modality == 'FP') or (image_modality == 'FP_PP'):
                optical_element0 = np.ones([num_elements,], dtype=np.float32)
            elif (image_modality == 'STORM') or (image_modality == 'phase'):
                optical_element0 = np.zeros([num_elements,], dtype=np.float32)
                optical_element0[0] = 1
        elif load_optical_element_init:
            optical_element0 = np.load(training_data_folder + '/optical_element_init.npy')

        else:
            optical_element0 = np.random.rand(num_elements,).astype(np.float32)

        if image_modality == 'FP':

            optical_element = variable_on_cpu('optical_element', optical_element0, tf.float32, \
                                              constraint = lambda x: tf.clip_by_value(x, 0, 1.0))

#            optical_element = variable_on_cpu('optical_element', optical_element0, tf.float32) #stretch

            num_nets = 2 # one neural net for real and one for imaginary

            numLEDs = optical_parameters_dict['numLEDs']
            N_obj = optical_parameters_dict['N_obj']
            N_patch = optical_parameters_dict['N_patch']
            num_patches = optical_parameters_dict['num_patches']
            Ns_mat = optical_parameters_dict['Ns_mat']
            scale_mat = optical_parameters_dict['scale_mat']
            cen0 = optical_parameters_dict['cen0']
            P = optical_parameters_dict['P']
            H0 = optical_parameters_dict['H0']
            Np = optical_parameters_dict['Np']

#            stretch = variable_on_cpu('sigmoid_stretch', 1.0, tf.float32, trainable=False)

            ### Convert trainingSet to low resolution

            low_res_obj_stack = create_FP_img_stack(optical_element,trainingSet,batch_size, N_obj, N_patch, num_patches, Ns_mat, \
                            scale_mat, cen0, P, H0, Np, numLEDs) #stretch

        elif image_modality == 'FP_PP':

            numLEDs = optical_parameters_dict['numLEDs']
            N_obj = optical_parameters_dict['N_obj']
            Np = optical_parameters_dict['Np']
            N_patch = optical_parameters_dict['N_patch']
            num_patches = optical_parameters_dict['num_patches']
            Ns_mat = optical_parameters_dict['Ns_mat']
            scale_mat = optical_parameters_dict['scale_mat']
            cen0 = optical_parameters_dict['cen0']
            P = optical_parameters_dict['P']
            H0 = optical_parameters_dict['H0']

            optical_element = variable_on_cpu('optical_element', optical_element0, tf.float32, \
                                              constraint = lambda x: tf.clip_by_value(x, 0, 1.0))

            num_nets = 2 # one neural net for real and one for imaginary

            low_res_obj_stack = lowres_trainingSet*optical_element
            low_res_obj_stack = tf.reduce_sum(low_res_obj_stack, axis=-1)

        elif image_modality == 'STORM':

            optical_element = variable_on_cpu('optical_element', optical_element0, tf.float32)

            num_wavelengths = optical_parameters_dict['num_wavelengths']
            num_nets = Nz*num_wavelengths

            H_fresnel_stack = optical_parameters_dict['H_fresnel_stack']
            H_NA = optical_parameters_dict['H_NA']
            ZernikePolyMat = optical_parameters_dict['ZernikePolyMat']

            PSF = ZernikePolyMat*optical_element
            PSF = tf.reduce_sum(PSF,axis=2)

            # Process 3D matrices to create the two-dimensional image

            low_res_obj_stack = create_microscopeImgStack(PSF, tf.cast(trainingSet, tf.complex64), Nz, batch_size, H_fresnel_stack, H_NA, Nx, num_wavelengths)

        elif image_modality == 'phase':

            optical_element = variable_on_cpu('optical_element', optical_element0, tf.float32)
            num_nets = 2 # one neural net for real and one for imaginary

            H_fresnel = optical_parameters_dict['H_fresnel']
            H_NA = optical_parameters_dict['H_NA']
            ZernikePolyMat = optical_parameters_dict['ZernikePolyMat']

            PSF = ZernikePolyMat*optical_element
            PSF = tf.reduce_sum(PSF,axis=2)

            low_res_obj_stack = create_phase_obj_stack(PSF, trainingSet, batch_size, H_fresnel, H_NA, Nx)

        low_res_obj_stack_nonoise = low_res_obj_stack

        if add_noise:
            low_res_obj_stack=add_noise_microscopeImgStack(low_res_obj_stack, normalRandomMat1, normalRandomMat2,\
                                                           sqrt_reg, poisson_noise_multiplier, gaussian_noise_multiplier, \
                                                           batch_size)

        if (image_modality == 'FP') or (image_modality == 'FP_PP'):
            # upsample the low_res_obj_stack for FP modality
            input_layer = create_upsampled_obj_stack(low_res_obj_stack, batch_size, Np, N_obj)
            input_layer = tf.expand_dims(input_layer, axis=3)
#            input_layer = tf.sqrt( (input_layer + sqrt_reg) / numLEDs)

        elif (image_modality == 'STORM') or (image_modality == 'phase'):
            input_layer = tf.expand_dims(low_res_obj_stack, axis=3)

    ### Neural Network(s)

    with tf.variable_scope("neural_net"):

        input_layer_list = [input_layer for i in range(num_nets)]
        predicted_mat = find_predicted_mat(image_modality, num_nets, input_layer_list, Nx, batch_size, layers_dropout, dropout_prob, \
                               use_batch_norm, autoencode_init_type, init_type_bias, init_type_resid, kernel_multiplier, variance_reg, training, \
                               num_layers_autoencode, skip_interval, num_blocks, optical_parameters_dict)

        if (image_modality == 'FP_PP') and (max_internal_iter == 0):

            loss = calculate_loss_FP(predicted_mat, trainingSet)

        else:
            ### convert predicted_mat down to low_res_obj_predicted
            if (image_modality == 'FP') or (image_modality == 'FP_PP'):

                low_res_obj_predicted = create_FP_img_stack(optical_element,predicted_mat,batch_size, N_obj, N_patch, num_patches, Ns_mat, \
                                scale_mat, cen0, P, H0, Np, numLEDs) #stretch

            elif image_modality == 'STORM':

                predicted_mat_complex=tf.cast(predicted_mat, tf.complex64)
                low_res_obj_predicted = create_microscopeImgStack(PSF,predicted_mat_complex,Nz,batch_size,H_fresnel_stack,H_NA,Nx,num_wavelengths)

            elif image_modality == 'phase':

                low_res_obj_predicted = create_phase_obj_stack(PSF, predicted_mat, batch_size, H_fresnel, H_NA, Nx)

            loss_low_res = tf.reduce_sum(tf.square(low_res_obj_stack - low_res_obj_predicted))

    ###################
            step_size = variable_on_cpu('step_size', 1e-7, tf.float32, trainable = False) # XXX make step_size trainable

            i = tf.Variable(0, dtype=tf.int32, trainable=False)

            if (image_modality == 'FP') or (image_modality == 'FP_PP'):
#                print 'mii4'
                if max_internal_iter > 0:
#                    print 'mii3'
                    predicted_mat = iterative_solver_FP(predicted_mat, optical_element, batch_size, N_obj, N_patch, num_patches, Ns_mat,\
                                scale_mat, cen0, P, H0, Np, numLEDs, low_res_obj_stack, step_size, max_internal_iter, \
                                merit_stopping_point, loss_low_res, i)

                loss = calculate_loss_FP(predicted_mat, trainingSet)

            elif image_modality == 'STORM':
                if max_internal_iter > 0:
                    predicted_mat = iterative_solver_STORM(predicted_mat, PSF, Nz,batch_size, H_fresnel_stack, \
                                   H_NA, Nx, low_res_obj_stack, step_size, max_internal_iter, \
                                   merit_stopping_point, loss_low_res, i, num_wavelengths)

                loss = tf.reduce_sum(tf.square(predicted_mat - tf.cast(trainingSet, tf.float32))) #l2 loss

            elif image_modality == 'phase':

                if max_internal_iter > 0:
                    predicted_mat = iterative_solver_phase(predicted_mat, PSF, batch_size, H_fresnel, \
                                   H_NA, Nx, low_res_obj_stack, step_size, max_internal_iter, \
                                   merit_stopping_point, loss_low_res, i)

                loss = calculate_loss_FP(predicted_mat, trainingSet)

        loss_grad = grad_diff_loss(predicted_mat, trainingSet)

    return loss, loss_grad, optical_element, low_res_obj_stack_nonoise, low_res_obj_stack, predicted_mat

def average_gradients(tower_grads, take_average=True): # Take
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """
  average_grads = []
  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, v in grad_and_vars:
#      print('*'*100)
#      print(g)
#      print(v)
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(axis=0, values=grads)

    if take_average:
        grad = tf.reduce_mean(grad, 0)
    else:
        grad = tf.reduce_sum(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

def load_optical_parameters(training_data_folder,image_modality):
    optical_parameters_dict = {}
    if image_modality=='FP' or image_modality == 'FP_PP':
        optical_parameters_dict['Ns_mat'] = np.load(training_data_folder + '/Ns_mat.npy')
        optical_parameters_dict['scale_mat'] = np.load(training_data_folder + '/scale_mat.npy')
        optical_parameters_dict['cen0'] = np.load(training_data_folder + '/cen0.npy')
        optical_parameters_dict['P'] = np.load(training_data_folder + '/P.npy')
        optical_parameters_dict['H0'] = np.load(training_data_folder + '/H0.npy')
        optical_parameters_dict['Np'] = np.load(training_data_folder + '/Np.npy')
        optical_parameters_dict['numLEDs'] = np.load(training_data_folder + '/numLEDs.npy')
        optical_parameters_dict['N_obj'] = np.load(training_data_folder + '/N_obj.npy')
        optical_parameters_dict['N_patch'] = np.load(training_data_folder + '/N_patch.npy')
        optical_parameters_dict['num_patches'] = np.load(training_data_folder + '/num_patches.npy')
        optical_parameters_dict['Nx_lowres'] = np.load(training_data_folder + '/Np.npy')
        optical_parameters_dict['Nx_highres'] = np.load(training_data_folder + '/N_obj.npy')
        optical_parameters_dict['NAfilter_synthetic'] = np.load(training_data_folder + '/NAfilter_synthetic.npy')
        optical_parameters_dict['Nz'] = 1

        num_elements = optical_parameters_dict['numLEDs']

    elif image_modality=='STORM':
        optical_parameters_dict['ZernikePolyMat'] = np.load(training_data_folder + '/ZernikePolyMat.npy')
        optical_parameters_dict['H_NA'] = np.load(training_data_folder + '/H_NA.npy')
        optical_parameters_dict['H_fresnel_stack'] = np.load(training_data_folder + '/H_fresnel_stack.npy')
        optical_parameters_dict['Nx_lowres'] = np.load(training_data_folder + '/Nx.npy')
        optical_parameters_dict['Nx_highres'] = np.load(training_data_folder + '/Nx.npy')
        optical_parameters_dict['Nz'] = np.load(training_data_folder + '/Nz.npy')
        optical_parameters_dict['num_wavelengths'] = np.load(training_data_folder + '/num_wavelengths.npy')
        optical_parameters_dict['numCoeff'] = np.load(training_data_folder + '/numCoeff.npy')

        num_elements = optical_parameters_dict['numCoeff']

    elif image_modality=='HE':
        only_noise = np.load(training_data_folder + '/only_noise.npy') # Option to keep magnification/resolution the same, but add noise
        optical_parameters_dict['only_noise'] = only_noise

        optical_parameters_dict['Nx_lowres'] = np.load(training_data_folder + '/Nx_lowres.npy')
        optical_parameters_dict['Nx_highres'] = np.load(training_data_folder + '/Nx_highres.npy')
        optical_parameters_dict['num_wavelengths'] = np.load(training_data_folder + '/num_wavelengths.npy')

        if not(only_noise):
            optical_parameters_dict['wavelength'] = np.load(training_data_folder + '/wavelength.npy')
            optical_parameters_dict['h_blur'] = np.load(training_data_folder + '/h_blur.npy')
            optical_parameters_dict['high_magn'] = np.load(training_data_folder + '/high_magn.npy')
            optical_parameters_dict['low_magn'] = np.load(training_data_folder + '/low_magn.npy')
            optical_parameters_dict['dpix_c'] = np.load(training_data_folder + '/dpix_c.npy')
            optical_parameters_dict['low_NA'] = np.load(training_data_folder + '/low_NA.npy')

        num_elements = None

    elif image_modality=='phase':
        optical_parameters_dict['ZernikePolyMat'] = np.load(training_data_folder + '/ZernikePolyMat.npy')
        optical_parameters_dict['H_NA'] = np.load(training_data_folder + '/H_NA.npy')
        optical_parameters_dict['H_fresnel'] = np.load(training_data_folder + '/H_fresnel.npy')
        optical_parameters_dict['Nx_lowres'] = np.load(training_data_folder + '/Nx.npy')
        optical_parameters_dict['Nx_highres'] = np.load(training_data_folder + '/Nx.npy')
        optical_parameters_dict['numCoeff'] = np.load(training_data_folder + '/numCoeff.npy')
        optical_parameters_dict['Nz'] = 1

        num_elements = optical_parameters_dict['numCoeff']

    return optical_parameters_dict, num_elements

###### Functions for H&E

def NAfilter(L,wavelength,NA,m):

    #L is the side length of the observation and source fields (assume square fields), L is a float
    #wavelength is the free space wavelength

    dx=L/m
    k=1./wavelength #wavenumber #2*pi/wavelength #1./wavelength

    fx=tf.linspace(-1/(2*dx),1/(2*dx)-1/L,m) #freq coords

    FX,FY=tf.meshgrid(fx,fx)

    zeros = tf.cast(tf.zeros([m,m]),dtype=tf.complex64)
    ones = tf.cast(tf.ones([m,m]),dtype=tf.complex64)

    H = tf.where(tf.sqrt(FX**2+FY**2)<=NA*k,ones,zeros)

#    H=fftshift(H,m,m)

    return H

def make_even(Np):

    if np.ceil(Np)%2: # make Np even
        Np = np.floor(Np)
    else:
        Np = np.ceil(Np)

    Np = int(Np)

    if Np % 2:
        Np += 1

    return Np

def change_magn_img(img0, m, h_blur, high_magn, low_magn, dpix_c, wavelength, low_NA):

    img = tf.cast(img0, dtype=tf.complex64)
    img = img[0:m,0:m]

    L = m*dpix_c/high_magn # [m]

    cen = [m//2,m//2]
    Np = m*low_magn/high_magn

    Np = make_even(Np)

    H_NA_R = NAfilter(L,wavelength[0],low_NA,m)
    H_NA_G = NAfilter(L,wavelength[1],low_NA,m)
    H_NA_B = NAfilter(L,wavelength[2],low_NA,m)

    img_r = incoherent_filter_H(img[:,:,0], m, H_NA_R, cen, Np, h_blur)
    img_g = incoherent_filter_H(img[:,:,1], m, H_NA_G, cen, Np, h_blur)
    img_b = incoherent_filter_H(img[:,:,2], m, H_NA_B, cen, Np, h_blur)

    img_r = tf.expand_dims(img_r, axis=2)
    img_g = tf.expand_dims(img_g, axis=2)
    img_b = tf.expand_dims(img_b, axis=2)
    img=tf.concat([img_r,img_g,img_b],axis=2)

#    img = tf.stack([img_r, img_g, img_b], axis=2)

    return img0, img

def incoherent_filter_H(u1, m, H, cen, Np, h_blur):

#    H0 = H

    H = F(Ft(H,m,m)*tf.conj(Ft(H,m,m)),m,m)
    H=H/H[m//2,m//2]

    U1 = F(u1,m,m)

    U2=H*U1

    #convolve with n x n filter with n = high_magn/low_magn
    U2 = U2 * F(h_blur,m,m)

    U2 = downsamp(U2,cen,Np)

    u2=Ft(U2,Np,Np)/(m/Np)**2

    u2 = tf.real(u2)

    # Clip u2 at 0 and 255

    zeros = tf.cast(tf.zeros([Np,Np]),dtype=tf.float32)
    ones = 255*tf.cast(tf.ones([Np,Np]),dtype=tf.float32)

    u2 = tf.where(u2>255,ones,u2)
    u2 = tf.where(u2<0,zeros,u2)

#    u2 = tf.cast(u2, dtype=tf.uint8)
#    u2=Ft(U2,m,m)

    return u2

def filter_function_NA(u1,H_NA,Nx,incoherent=1):
    #u1 is the source plane field
    #Nx is u1.shape[0]
    #dx is L/m

    H = H_NA

    if incoherent:
        H=tf.fft2d(tf.ifft2d(H)*tf.conj(tf.ifft2d(H)))
        #H=H/H[0,0]

        U1=fftshift(u1,Nx,Nx)
        U1=tf.fft2d(U1)

        U2=H*U1
        u2=fftshift(tf.ifft2d(U2),Nx,Nx)

    else:
        U1=fftshift(u1,Nx,Nx)
        U1=tf.fft2d(U1)

        U2=H*U1
        u2=fftshift(tf.ifft2d(U2),Nx,Nx)

#        u2 = u2*np.conj(u2) # make into intensity object

    return u2

def read_img_file(img_path, channels):
    img_file = tf.read_file(img_path)
    img0 = tf.image.decode_image(img_file, channels=channels)
    return img0

def change_magn(img_path, m, h_blur, high_magn, low_magn, dpix_c, wavelength, low_NA, num_wavelengths):

    img0 = read_img_file(img_path, num_wavelengths)

    return change_magn_img(img0, m, h_blur, high_magn, low_magn, dpix_c, wavelength, low_NA)

def change_magn_batch(img_stack, m, h_blur, high_magn, low_magn, dpix_c, wavelength, low_NA, batch_size):

    for ii in range(batch_size):

        ( _ , img) =change_magn_img(img_stack[ii,:,:,:], m, h_blur, high_magn, low_magn, dpix_c, wavelength, low_NA)

        img = tf.expand_dims(img,axis=0)

        if ii == 0:
            img_stack_new = img
        else:
            img_stack_new = tf.concat([img_stack_new,img],0)

#    img_stack_new = tf.cast(img_stack_new, dtype=tf.complex64)

    return img_stack_new

def add_poisson_noise(img0, img, noise_multiplier):

    img = noise_multiplier*tf.cast(img, dtype=tf.float32)
    img = tf.random_poisson(img, [1])
    img = img/noise_multiplier

    # Clip img at 0 and 255

    zeros = tf.cast(tf.zeros_like(img),dtype=tf.float32)
    ones = 255*tf.cast(tf.ones_like(img),dtype=tf.float32)

    img = tf.where(img>255,ones,img)
    img = tf.where(img<0,zeros,img)

    img = tf.cast(img,dtype=tf.uint8)

    img = tf.squeeze(img, axis=0)

    return img0, img

def upsample_rgb(low_res_obj_rgb, Np, N_obj, num_wavelengths):

    for c in range(num_wavelengths):
        upsampled_obj = upsample(low_res_obj_rgb[:,:,c], Np, N_obj)

        upsampled_obj = tf.expand_dims(upsampled_obj, axis=2)

        if c == 0:
            upsampled_obj_stack = upsampled_obj
        else:
            upsampled_obj_stack = tf.concat([upsampled_obj_stack,upsampled_obj],2)

    return upsampled_obj_stack

def make_iterator_numpy(training_dataset, batch_size, num_GPUs, shuffle=False):
    tr_data = tf.data.Dataset.from_tensor_slices(training_dataset)
    tr_indices =tf.data.Dataset.from_tensor_slices(tf.constant(np.arange(0,training_dataset.shape[0],1),dtype = tf.int32))

    tr_data = tf.data.Dataset.zip((tr_data, tr_indices))

    tr_data = tr_data.repeat()

    if shuffle:
        tr_data = tr_data.shuffle(training_dataset.shape[0])

    tr_data = tr_data.batch(batch_size*num_GPUs)

    return tr_data

def make_iterator_FP_PP(training_dataset, lowres_training_dataset, \
                        batch_size, num_GPUs, shuffle=False):

    tr_data = tf.data.Dataset.from_tensor_slices(training_dataset)
    tr_data_lowres = tf.data.Dataset.from_tensor_slices(lowres_training_dataset)
    tr_indices =tf.data.Dataset.from_tensor_slices(tf.constant(np.arange(0,training_dataset.shape[0],1),dtype = tf.int32))

    tr_data = tf.data.Dataset.zip((tr_data, tr_data_lowres, tr_indices))

    tr_data = tr_data.repeat()

    if shuffle:
        tr_data = tr_data.shuffle(training_dataset.shape[0])

    tr_data = tr_data.batch(batch_size*num_GPUs)

    return tr_data

def make_complex_img(img_path, Nx, phase = True, intensity = False, library = tf):
    if library == tf:
        img0 = read_img_file(img_path, None)
        img0 = library.reduce_sum(img0,axis=2)
        img0 = img0[Nx/2:-Nx/2,Nx/2:-Nx/2]
        img0 = library.cast(img0,dtype=library.complex64)
    else:
        files = glob.glob(img_path)
        fileI = files[0]
        img0 = misc.imread(fileI)
        img0 = library.sum(img0, axis=2)
        img0 = img0[0:Nx,0:Nx]
        img0 = img0.astype(library.complex64)

    img0 = library.exp(1j*math.pi*1.0*(255.0*3-img0)/(255.0*3))
    return img0

def make_iterator_filename_FP(training_dataset, batch_size, num_GPUs, optical_parameters_dict, \
                              shuffle=False):

    Nx = optical_parameters_dict['Nx_highres']
    H_NA = optical_parameters_dict['NAfilter_synthetic']

    training_dataset = tf.constant(training_dataset.tolist())
    tr_data = tf.data.Dataset.from_tensor_slices(training_dataset)
    tr_indices =tf.data.Dataset.from_tensor_slices(tf.constant(np.arange(0,int(training_dataset.shape[0]),1),dtype = tf.int32))

    filter_function_NA_lambda = lambda u1: filter_function_NA(u1,H_NA,Nx,incoherent=0)
    make_complex_img_lambda = lambda img_path: make_complex_img(img_path,Nx)

    tr_data = tr_data.map(make_complex_img_lambda)
    tr_data = tr_data.map(filter_function_NA_lambda)
    tr_data = tf.data.Dataset.zip((tr_data, tr_indices))

    tr_data = tr_data.repeat()

    if shuffle:
        tr_data = tr_data.shuffle(training_dataset.shape[0])

    tr_data = tr_data.batch(batch_size*num_GPUs)

    return tr_data

def create_STORM_stack(img_path, Nx):
    img0 = read_img_file(img_path, None)
    img0 = img0[0:Nx,0:Nx,:]
    img0 = tf.expand_dims(img0, axis = 2)
    img0 = tf.cast(img0, tf.float32)
    return img0

def make_iterator_filename_STORM(training_dataset, batch_size, num_GPUs, optical_parameters_dict, \
                                 shuffle=False):

    Nx = optical_parameters_dict['Nx_highres']

    training_dataset = tf.constant(training_dataset.tolist())
    tr_data = tf.data.Dataset.from_tensor_slices(training_dataset)
    tr_indices =tf.data.Dataset.from_tensor_slices(tf.constant(np.arange(0,int(training_dataset.shape[0]),1),dtype = tf.int32))

    create_STORM_stack_lambda = lambda img_path: create_STORM_stack(img_path, Nx)

    tr_data = tr_data.map(create_STORM_stack_lambda)
    tr_data = tf.data.Dataset.zip((tr_data, tr_indices))

    tr_data = tr_data.repeat()

    if shuffle:
        tr_data = tr_data.shuffle(training_dataset.shape[0])

    tr_data = tr_data.batch(batch_size*num_GPUs)

    return tr_data

def make_iterator_filename(training_dataset, add_noise, poisson_noise_multiplier, batch_size, \
                           num_GPUs, optical_parameters_dict, shuffle = False):

    only_noise = optical_parameters_dict['only_noise']
    num_wavelengths = optical_parameters_dict['num_wavelengths']

    training_dataset = tf.constant(training_dataset.tolist())
    tr_data = tf.data.Dataset.from_tensor_slices(training_dataset)
    tr_indices =tf.data.Dataset.from_tensor_slices(tf.constant(np.arange(0,int(training_dataset.shape[0]),1),dtype = tf.int32))

    if only_noise:
        read_img_file_lambda = lambda img_path: read_img_file(img_path, num_wavelengths)
        tr_data = tr_data.map(lambda img_path: (read_img_file_lambda(img_path),read_img_file_lambda(img_path)))
    else:
        change_magn_lambda = lambda img_path: change_magn(img_path,
                                                  optical_parameters_dict['Nx_highres'],
                                                  optical_parameters_dict['h_blur'],
                                                  optical_parameters_dict['high_magn'],
                                                  optical_parameters_dict['low_magn'],
                                                  optical_parameters_dict['dpix_c'],
                                                  optical_parameters_dict['wavelength'],
                                                  optical_parameters_dict['low_NA'],
                                                  num_wavelengths)
        tr_data = tr_data.map(change_magn_lambda)

    add_poisson_noise_lambda = lambda img0, img: add_poisson_noise(img0, img, poisson_noise_multiplier)

    if only_noise and not(add_noise):
        print('Warning: only_noise = True and add_noise = False!!')

    if add_noise:
        tr_data = tr_data.map(add_poisson_noise_lambda)

    tr_data = tf.data.Dataset.zip((tr_data, tr_indices))

    tr_data = tr_data.repeat()

    if shuffle:
        tr_data = tr_data.shuffle(int(training_dataset.shape[0]))

    tr_data = tr_data.batch(batch_size*num_GPUs)

    return tr_data

def tower_loss_HE(low_res_obj_batch, \
                  high_res_obj_batch, training, \
                  use_batch_norm, variance_reg, sqrt_reg, \
                  dropout_prob, layers_dropout, batch_size, max_internal_iter, merit_stopping_point, \
                  optical_parameters_dict, autoencode_init_type, init_type_bias, init_type_resid, \
                  kernel_multiplier, num_layers_autoencode, skip_interval, num_blocks):

    only_noise = optical_parameters_dict['only_noise']
    Nx_lowres = optical_parameters_dict['Nx_lowres']
    Nx_highres = optical_parameters_dict['Nx_highres']
    num_wavelengths = optical_parameters_dict['num_wavelengths']

    if not(only_noise):
        wavelength = optical_parameters_dict['wavelength']
        h_blur = optical_parameters_dict['h_blur']
        high_magn = optical_parameters_dict['high_magn']
        low_magn = optical_parameters_dict['low_magn']
        dpix_c = optical_parameters_dict['dpix_c']
        low_NA = optical_parameters_dict['low_NA']

    with tf.variable_scope("optical_transform"):

        if only_noise:
            upsampled_obj_stack = low_res_obj_batch
        else:
            # upsample the low_res_obj_stack
            simple_upsample = 0
            if simple_upsample:
                low_res_obj_batch = tf.cast(low_res_obj_batch, tf.complex64)
                for ii in range(batch_size):
                    upsampled_obj = upsample_rgb(low_res_obj_batch[ii,:,:,:], Nx_lowres, Nx_highres, num_wavelengths)

                    upsampled_obj = tf.expand_dims(upsampled_obj,axis=0)
                    if ii == 0:
                        upsampled_obj_stack = upsampled_obj
                    else:
                        upsampled_obj_stack = tf.concat([upsampled_obj_stack,upsampled_obj],0)
            else:
                upsampled_obj_stack = tf.image.resize_images(low_res_obj_batch, [int(Nx_highres), int(Nx_highres)], \
                                                       method=tf.image.ResizeMethod.BICUBIC, align_corners=False)

        upsampled_obj_stack = tf.cast(upsampled_obj_stack, tf.float32)

    with tf.variable_scope("neural_net"):

        num_nets = num_wavelengths
        input_layer_list = []

        for w in range(num_wavelengths):

            input_layer = tf.expand_dims(upsampled_obj_stack[:,:,:,w], axis=3)
            input_layer_list.append(input_layer)

#        len(input_layer) should equal num_nets

#        input_layer_r = tf.expand_dims(upsampled_obj_stack[:,:,:,0], axis=3)
#        input_layer_g = tf.expand_dims(upsampled_obj_stack[:,:,:,1], axis=3)
#        input_layer_b = tf.expand_dims(upsampled_obj_stack[:,:,:,2], axis=3)
#
#
#        input_layer_list = [input_layer_r, input_layer_g, input_layer_b]

        predicted_mat = find_predicted_mat('HE', num_nets, input_layer_list, Nx_highres, batch_size, layers_dropout, dropout_prob,
                       use_batch_norm, autoencode_init_type, init_type_bias, init_type_resid, kernel_multiplier, variance_reg, training, \
                       num_layers_autoencode, skip_interval, num_blocks, optical_parameters_dict)

    with tf.variable_scope("neural_net"):

#        ### convert predicted_mat down to low_res_obj_predicted
        if only_noise:
            low_res_obj_batch_predicted = low_res_obj_batch

        else:
            low_res_obj_batch_predicted = change_magn_batch(predicted_mat, Nx_highres, h_blur, \
                                                            high_magn, low_magn, dpix_c, wavelength, low_NA, batch_size)

            low_res_obj_batch_predicted = tf.cast(low_res_obj_batch_predicted, tf.float32)
            low_res_obj_batch = tf.cast(low_res_obj_batch, tf.float32)
            loss_low_res = tf.reduce_sum(tf.square(low_res_obj_batch - low_res_obj_batch_predicted))

###################
        step_size = variable_on_cpu('step_size', 1e-7, tf.float32, trainable = False) # XXX make step_size trainable

    with tf.variable_scope("iteration_variable"):
        i = tf.Variable(0, dtype=tf.int32, trainable=False)

    with tf.variable_scope("neural_net"):

        if not(only_noise):
            if max_internal_iter > 0:
                predicted_mat, loss_low_res = iterative_solver_HE(predicted_mat, Nx_highres, h_blur, \
                               high_magn, low_magn, dpix_c, wavelength, \
                               low_NA, low_res_obj_batch, step_size, max_internal_iter, \
                               merit_stopping_point, loss_low_res, i, \
                               batch_size)

        high_res_obj_batch = tf.cast(high_res_obj_batch, tf.float32)

        loss = tf.reduce_sum(tf.square(high_res_obj_batch - predicted_mat))

        if only_noise:
            loss_low_res = loss

        loss_grad = grad_diff_loss(predicted_mat, high_res_obj_batch)

    return loss, loss_grad, loss_low_res, low_res_obj_batch_predicted, low_res_obj_batch, predicted_mat
    #output names are: loss, loss_grad, optical_element, low_res_obj_stack_nonoise, low_res_obj_stack, predicted_mat
    #for only_noise = True, loss_low_res is set to loss, and low_res_obj_batch_predicted is set to low_res_obj_batch

ENGR 090. Fourier ptychographic microscopy for biological samples/TensorFlowFunctions.py


EnsureSConsVersion(1, 1, 0)

import os
import re
from os.path import join as pjoin

opts = Variables('build.py')

env = Environment(options=opts,
                  ENV = os.environ.copy(),
                  tools=['default'])

conf = Configure(env, custom_tests = {})
conf.env['NODE'] = conf.env.WhereIs('node')
conf.env.AppendUnique(RPATH = conf.env.get('LIBPATH'))
env = conf.Finish()

Export("env")

source = SConscript("lib/SConscript")

env["JSLINT"] = "$NODE lib/extern/node-jslint/bin/jslint.js"
jslint = [env.Command(pjoin('.jslint/', str(x))+".jslint", x, ["$JSLINT $SOURCE || exit 0"]) for x in source]
env.AlwaysBuild(jslint)

env.Alias('jslint', jslint)

targets = []

env.Default(targets)


import requests
import json
import sys
import urllib.parse
import os.path
import hashlib
import argparse

# import sqlite3
# from tplib3 import *
# import datetime
# import time
import platform

from sqlite3 import Error

# updategroupconfig
#
# Fetches siteconfig from config file
# fetches groupconfig from cloud
# fetches from disk
# if disk file is missing or there are changes, write new config to disk.
#
# Python version: 3
#
# Modifictions
# 	5/8-20 initial creation
#
# Usage:
# python updategroupconfig [--configfile filename]

progVersion = "1.0 05082020"

if platform.system().lower() == "linux":
    siteConfigFile = "/home/pi/conf/siteconfig.json"
else:
    siteConfigFile = "C:\\usr\\demo_projects\\pi\\hytta\\config.json"

def getGroups(siteConfig):
    #
    #  Fetch all groups
    #
    headers = {"content-type": "application/json"}
    try:
        g = requests.get(siteConfig["groupConfigURL"], headers=headers)
    except IOError as err:
        print("REST API Call for: " + siteConfig["groupConfigURL"] + " failed: ", err)
        return False
    localConfig = []
    if g.ok:
        #
        # Fetch memerships
        #
        try:
            m = requests.get(siteConfig["membersURL"], headers=headers)
            if m.ok:
                allGroups = g.json()["items"]
                allmembers = m.json()["items"]
                #
                # Iterate through and add member array
                #
                for group in list(allGroups):
                    #
                    # iterate through all memers of a group
                    #
                    members = []
                    for member in list(allmembers):
                        if member["groupname"] == group["groupname"]:
                            members.append(member["membername"])
                    group["members"] = members
                    localConfig.append(group)
                f = open(siteConfig["groupConfigFile"], "w")
                f.write(json.dumps(localConfig))
                f.close()
                return localConfig
            else:
                print(
                    "API call for members failed, JSON Parse Error: "
                    + str(m.status_code)
                )
                return False
        except exception as e:
            print("API Call generated invalid result")
            print(siteConfig["membersURL"])
            print(headers)
            print(str(e))
            return False
    else:
        print("API call for groups failed: " + str(g.status_code))
        return False

#
def loadSiteConfig(siteConfigFile):
    #
    # Load master config file to get groupConfigFile path,
    # confURL, for static group membership
    # tempConfigURL for dynamic temp setting
    # groupconfig, contains statig gruop names and group sensor setup
    #
    try:
        f = open(siteConfigFile, "r")
        strsiteconfig = f.read()
        siteConfig = json.loads(strsiteconfig)
        f.close()
        return siteConfig
    except IOError:
        print("Site Config File" + siteConfigFile + " is not accessible")
        return False
    return null

def main():
    global siteConfigFile
    #
    # Identify program
    #
    print()
    print("update group config Version: " + progVersion)
    print()
    #
    # process comand line onl one optin available, -f filename for sysconfig fiel
    #
    argsParser = argparse.ArgumentParser(description="updategroupconfig program")
    argsParser.add_argument(
        "--configfile", default=siteConfigFile, type=str, help="Site Config File"
    )
    args = argsParser.parse_args()
    siteConfigFile = args.configfile
    #
    # Load static site config information
    #
    siteConfig = loadSiteConfig(siteConfigFile)
    if not siteConfig:
        print("Error Loading confguration")
        return 1
    #
    #  get current group config file
    #
    if os.path.isfile(siteConfig["groupConfigFile"]):
        try:
            f = open(siteConfig["groupConfigFile"], "r")
            strconfig = f.read()
            if len(strconfig) > 1:
                currentGroups = json.loads(strconfig)
            else:
                currentGroups = "{}"
            f.close()
        except IOError:
            print(
                "Member Config File"
                + siteConfig["groupConfigFile"]
                + " exists but is not accessible"
            )
            return 2
        groups = getGroups(siteConfig)
        if groups != False:
            if (
                hashlib.md5(json.dumps(groups).encode("utf-8")).digest()
                != hashlib.md5(json.dumps(currentGroups).encode("utf-8")).digest()
            ):
                print("New version found, saving.....")
                f = open(siteConfig["groupConfigFile"], "w")
                f.write(json.dumps(groups))
                f.close()
            else:
                print("Files are in sync")
            sys.exit(0)
        else:
            print("getGroups failed")
            sys.exit(1)
        sys.exit(2)
    else:
        #
        # Group Config did not exists need to fetch it from iot database
        #
        groups = getGroups(siteConfig)

    print("groupconfig.json synchronized")

if __name__ == "__main__":
    main()


import sys
from PyQt4.QtGui import *
from PyQt4.QtCore import *
from remote_class import *
from timelapse_class import *

class DollyProject(QWidget):

    def __init__(self, stack):
        #super(DollyProject, self).__init__()
        QWidget.__init__(self, stack)
        self.stack = stack
        self.initUI()
        self.mode_current = 0

    def initUI(self):

        #set background
        palette = QPalette()
        palette.setColor(QPalette.Background,QColor(36, 36, 36))
        self.setPalette(palette)

        self.stack1 = QWidget()
        self.mainwindow()
        #self.howto1()

        self.setGeometry(0,0,480,320)
        self.setWindowTitle('DollyProject')
        #self.showFullScreen()
        self.show()

    def mainwindow(self):
        print ('call mainwindow')
        self.lbl = QLabel(self)
        pixmap = QPixmap("img/logo_menu.png")
        self.lbl.setPixmap(pixmap)
        self.lbl.move(140, 30)

        self.btn1 = QPushButton(self)
        self.btn1.setIcon(QIcon('img/mode_1.png'))
        self.btn1.setIconSize(QSize(138,145))
        self.btn1.resize(138,145)
        self.btn1.move(20, 90)
        self.btn1.clicked.connect(self.howto1)

        self.btn2 = QPushButton(self)
        self.btn2.setIcon(QIcon('img/mode_2.png'))
        self.btn2.setIconSize(QSize(138,145))
        self.btn2.resize(138,145)
        self.btn2.move(170, 90)

        self.btn3 = QPushButton(self)
        self.btn3.setIcon(QIcon('img/mode_3.png'))
        self.btn3.setIconSize(QSize(138,145))
        self.btn3.resize(138,145)
        self.btn3.move(320, 90)
        self.btn3.clicked.connect(self.howto3)

        self.aboutBtn = QPushButton(self)
        self.aboutBtn.setIcon(QIcon('img/info.png'))
        self.aboutBtn.setIconSize(QSize(53,53))
        self.aboutBtn.resize(53,53)
        self.aboutBtn.move(10, 260)
        self.aboutBtn.clicked.connect(self.aboutme)

        self.close = QPushButton(self)
        self.close.setIcon(QIcon('img/power.png'))
        self.close.setIconSize(QSize(53,53))
        self.close.resize(53,53)
        self.close.move(420, 260)
        self.close.clicked.connect(QCoreApplication.instance().quit)

        #HIDE ITEM ALL-------------------------------------------
        self.lbl2 = QLabel(self)
        pixmap2 = QPixmap("img/howto_1.png")
        self.lbl2.setPixmap(pixmap2)
        self.lbl2.move(-1000, -1000)

        self.lbm3 = QLabel(self)
        pixmap = QPixmap("img/howto_3.png")
        self.lbm3.setPixmap(pixmap)
        self.lbm3.move(-1000, -1000)

        self.nextBtn = QPushButton('', self)
        self.nextBtn.setIcon(QIcon('img/nextBtn.png'))
        self.nextBtn.setIconSize(QSize(127,66))
        self.nextBtn.resize(127,66)
        self.nextBtn.move(-1000,-1000)
        self.nextBtn.clicked.connect(self.nextstep)

        self.homeBtn = QPushButton('', self)
        self.homeBtn.setIcon(QIcon('img/home_btn.png'))
        self.homeBtn.setIconSize(QSize(50,46))
        self.homeBtn.resize(50,46)
        self.homeBtn.move(-1000,-1000)
        self.homeBtn.clicked.connect(self.backhome)

        self.bgwin = QLabel(self)

    def howto1(self):
        print ('call howto1')
        self.mode_current = 1
        self.deleteObject()
        self.lbl2.move(0, 0)
        self.nextBtn.move(177, 237)
        self.homeBtn.move(5, 271)
        #self.aboutBtn.clicked.connect(self.backhome)
    def howto3(self):
        print ('call howto3')
        self.mode_current = 3
        self.deleteObject()
        self.lbm3.move(0, 0)
        self.nextBtn.move(177, 237)
        self.homeBtn.move(5, 271)

        #self.aboutBtn.clicked.connect(self.backhome)
    def nextstep(self):
        if self.mode_current == 1:
            #hide howto
            self.lbl2.move(-1000, -1000)
            self.nextBtn.move(-1000, -1000)
            self.remote()
        elif self.mode_current == 3:
            self.lbm3.move(-1000, -1000)
            self.nextBtn.move(-1000, -1000)
            self.timelapse()

    def backhome(self):
        print ('mode = ',self.mode_current)
        if self.mode_current == 1:
            self.lbl2.move(-1000, -1000)
        elif self.mode_current == 3:
            self.lbm3.move(-1000, -1000)
        self.nextBtn.move(-1000,-1000)
        self.homeBtn.move(-1000,-1000)
        self.refreshMain()

    def remote(self):
        print ('Mode remote is work')
        self.stack.setPage2()

    def timelapse(self):
        print ('Time Lapse is work')
        #self.lbm3.move(-1000, -1000)
        #self.nextBtn.move(-1000, -1000)

        self.stack.setPage3()

        #init_time(self, 1)
        #callwin_timelapse(self)
        #minusBtn_start.clicked.connect(self.startDown)

    def aboutme(self):
        msgBox = QMessageBox()
        msgBox.setText("Created by\nNipitpon Chantada\n5533684623")
        msgBox.setIcon(QMessageBox.Information)
        retval = msgBox.exec_()

    def deleteObject(self):
        print ("delete")
        #hide all
        self.lbl.move(-1000,-1000)
        self.btn1.move(-1000,-1000)
        self.btn2.move(-1000,-1000)
        self.btn3.move(-1000,-1000)
        self.aboutBtn.move(-1000,-1000)
        self.close.move(-1000,-1000)

    def refreshMain(self):
        self.lbl.move(140, 30)
        self.btn1.move(20, 90)
        self.btn2.move(170, 90)
        self.btn3.move(320, 90)
        self.aboutBtn.move(10, 260)
        self.close.move(420, 260)

    def goIndex(self):
        self.stack.setPage1()

class StackedWidget(QStackedWidget):

    def __init__(self, parent = None):
        QStackedWidget.__init__(self, parent)

    def setCurrentIndex(self, index):
        self.fader_widget = DollyProject(self.currentWidget())
        QStackedWidget.setCurrentIndex(self, index)

    def setPage1(self): #page index
        self.setCurrentIndex(0)

    def setPage2(self): #page remote

        self.setCurrentIndex(1)

    def setPage3(self): #page face-detection
        self.setCurrentIndex(2)

    def setPage4(self): #page timelapse
        self.setCurrentIndex(3)

def main():
    app = QApplication(sys.argv)
    #set foreground color
    palette = QPalette()
    palette.setColor(QPalette.Background,QColor(36, 36, 36))

    window = QWidget()
    window.setPalette(palette)
    #ex = DollyProject()
    stack = StackedWidget()
    stack.addWidget(DollyProject(stack))
    stack.addWidget(Remote(stack))
    stack.addWidget(Timelapse(stack))

    layout = QHBoxLayout(window)
    layout.addWidget(stack)
    window.move(0,0)
    window.resize(480,320)
    window.layout().setContentsMargins(0, 0, 0, 0)
    #window.show()
    window.showFullScreen()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()


import sys
sys.path.append("../")
import pymysql
import math
import pandas as pd
import numpy as np
from scipy.spatial import distance
import csv
from InsertDataBase.CreateTables import *
from DBtools.init_db import init_DB

def GetMinDistanceLane(local_x, local_y):
    radius = 5
    node_list = list()
    way_list = list()
    while len(node_list) < 3:
        node_list = list()
        way_list = list()
        min_x = local_x - radius
        max_x = local_x + radius
        min_y = local_y - radius
        max_y = local_y + radius
        sql = "select Node_To_Way.node_id, way_id " \
              "from Node_To_Way " \
              "join Node_Info on Node_To_Way.node_id = Node_Info.node_id " \
              "where local_x BETWEEN %s and %s and local_y BETWEEN %s and %s group by Node_To_Way.node_id,way_id" \
              % (min_x, max_x, min_y, max_y)
        cursor.execute(sql)
        node_list = list()
        for i in cursor.fetchall():
            node_list.append(i[0])
            way_list.append(i[1])
        radius = 2 * radius

    way_list = list(set(way_list))
    vehicle_position_list = np.array([[local_x, local_y]])
    dist = np.inf
    lane = None
    for way_id in way_list:
        temp_position_list = list()
        sql = "select Node_Info.node_id, local_x, local_y " \
              "from Node_Info " \
              "join Node_To_Way on Node_Info.node_id = Node_To_Way.node_id " \
              "where way_id = %s" % (way_id)
        cursor.execute(sql)
        for i in cursor.fetchall():
            temp_position_list.append([float(i[1]), float(i[2])])
        temp_position_list = np.array(temp_position_list)
        new_dist = distance.cdist(vehicle_position_list, temp_position_list).min(axis=1).min(axis=0)
        if new_dist < dist:
            dist = new_dist
            lane = str(way_id)
    return lane

def InsertTable(cursor, VehicleInfo, table):
    insertTimingSql = "insert into Traffic_timing_state" + table + "(time_stamp,vehicle_id,local_x,local_y,orientation,lane_id) " \
                "values(%s,%s,%s,%s,%s,%s)"
    for i in range(len(VehicleInfo)):
        time_stamp = int(float(VehicleInfo[i][0]) * 1e+3)
        id_str = VehicleInfo[i][1].split('-')
        vehicle_id = int(id_str[-1])
        local_x = float(VehicleInfo[i][3])
        local_y = float(VehicleInfo[i][4])
        lane = GetMinDistanceLane(local_x, local_y)
        orien = 0.000
        cursor.execute(insertTimingSql,(time_stamp, vehicle_id, local_x, local_y, orien, lane))
        if (i % 100 == 0):
            print(i)

    insertPropertySql = "insert into Traffic_participant_property" + table + "(vehicle_id,vehicle_class) " \
                "values(%s,%s) ON DUPLICATE KEY UPDATE vehicle_class = vehicle_class"
    for i in range(len(VehicleInfo)):
        vehicle_type = VehicleInfo[i][2]
        id_str = VehicleInfo[i][1].split('-')
        vehicle_id = int(id_str[-1])
        vehicle_class = int(1)
        cursor.execute(insertPropertySql,(vehicle_id, vehicle_class))
        if (i % 100 == 0):
            print(i)

    UpdateParticipantSql = "Update Traffic_timing_state" + table + " set orientation = %s where time_stamp = %s and vehicle_id = %s"
    stampsql = "select time_stamp from Traffic_timing_state" + table
    vehiclesql = "select vehicle_id from Traffic_timing_state" + table + " where time_stamp = %s"
    positionsql = "select local_x, local_y from Traffic_timing_state" + table + " where time_stamp = %s and vehicle_id = %s"
    oriensql = "select orientation from Traffic_timing_state" + table + " where time_stamp = %s and vehicle_id = %s"
    cursor.execute(stampsql)
    timestampresult = cursor.fetchall()
    timestamp_set = set()
    for i in range(len(timestampresult)):
        timestamp_set.add(timestampresult[i][0])
    timestamp_list = list(timestamp_set)
    timestamp_list.sort()
    vehicle_zero_orien_list = list()
    for t in range(len(timestamp_list) - 1):
        vehicle_list = list()
        cursor.execute(vehiclesql, timestamp_list[t + 1])
        for vehicle in cursor.fetchall():
            vehicle_list.append(vehicle[0])
        if t == 0:
            cursor.execute(vehiclesql, timestamp_list[t])
            for vehicle_zero in cursor.fetchall():
                vehicle_zero_orien_list.append(vehicle_zero[0])
        for vehicle in vehicle_list:
            local_x = 0
            local_y = 0
            last_x = 0
            last_y = 0
            cursor.execute(positionsql, (timestamp_list[t + 1], vehicle))
            for i in cursor.fetchall():
                local_x = i[0]
                local_y = i[1]
            cursor.execute(positionsql, (timestamp_list[t], vehicle))
            for j in cursor.fetchall():
                last_x = j[0]
                last_y = j[1]
            if local_x != None and local_y != None and last_x != None and last_y != None:
                if last_x == 0 and last_y == 0:
                    vehicle_zero_orien_list.append(vehicle)
                    continue
                direction = [local_x - last_x, local_y - last_y]
                orien = math.atan2(direction[1], direction[0])
                if local_x == last_x and local_y == last_y:
                    cursor.execute(oriensql, (timestamp_list[t], vehicle))
                    last_orien = cursor.fetchall()[0][0]
                    cursor.execute(UpdateParticipantSql, (last_orien, timestamp_list[t + 1], vehicle))
                else:
                    cursor.execute(UpdateParticipantSql, (orien, timestamp_list[t + 1], vehicle))
                if vehicle in vehicle_zero_orien_list:
                    cursor.execute(UpdateParticipantSql, (orien, timestamp_list[t], vehicle))
                    vehicle_zero_orien_list.remove(vehicle)

if __name__ == '__main__':
    conn, cursor = init_DB("Argoverse_MIA_Scenario_DB")
    csv_reader = csv.reader(open("../Annotator/sample_record.csv", encoding='utf-8'))
    for i, rows in enumerate(csv_reader):
        table = str(rows[0])
        print("table: ", table)
        VehicleInfo = list()
        file = "../../argoverse_trajectory/forecasting_train_v1.1/train/data/" + table + ".csv"
        VehicleInfo = pd.read_csv(file, decimal=",", low_memory=False)
        VehicleInfo = np.array(VehicleInfo)
        table = "_" + table
        CreateTrafficParticipantPropertyTable(cursor, table)
        CreateTrafficTimingStateTable(cursor, table)
        InsertTable(cursor, VehicleInfo, table)

    cursor.close()
    conn.commit()
    conn.close()


"""Open merge example.

Trains a a small percentage of rl vehicles to dissipate shockwaves caused by
on-ramp merge to a single lane open highway network.
"""
from flow.core.params import SumoParams, EnvParams, InitialConfig
from flow.core.params import NetParams, InFlows, SumoCarFollowingParams
from SingalNetwork import ADDITIONAL_NET_PARAMS
from flow.core.params import VehicleParams
from flow.controllers import RLController
from SingalEnv import SingalEnv
from SingalNetwork import SingalNetwork

HORIZON = 600
N_ROLLOUTS = 20
N_CPUS = 6
FLOW_RATE = 10000
NUM_RL = 250

additional_net_params = ADDITIONAL_NET_PARAMS.copy()

vehicles = VehicleParams()
vehicles.add(
    veh_id="human",
    car_following_params=SumoCarFollowingParams(
        speed_mode="aggressive",
    ),
    num_vehicles=0)
vehicles.add(
    veh_id="rl",
    acceleration_controller=(RLController, {}),
    car_following_params=SumoCarFollowingParams(
        speed_mode="aggressive",
    ),
    num_vehicles=0)
# vehicles.add(
#     veh_id="rl_truck",
#     acceleration_controller=(RLController, {}),
#     car_following_params=SumoCarFollowingParams(
#         speed_mode="aggressive",
#     ),
#     num_vehicles=0)

inflow = InFlows()
inflow.add(
    veh_type="human",
    edge="inflow_highway",
    vehs_per_hour= FLOW_RATE * 0.1,
    departLane="free",
    departSpeed=10)

inflow.add(
    veh_type="rl",
    edge="inflow_highway",
    vehs_per_hour=FLOW_RATE * 0.9,
    departLane="free",
    departSpeed=10
    )

inflow.add(
    veh_type="rl",
    edge="inflow_merge_west",
    vehs_per_hour=3000,
    departLane="free",
    departSpeed=7.5
    )
inflow.add(
    veh_type="rl",
    edge="inflow_merge_east",
    vehs_per_hour=3000,
    departLane="free",
    departSpeed=7.5
    )

flow_params = dict(
    exp_tag="SingalEnv",

    env_name=SingalEnv,

    network=SingalNetwork,

    simulator='traci',

    sim=SumoParams(
        sim_step=0.2,
        render=False,
        restart_instance=True,
    ),

    env=EnvParams(
        horizon=HORIZON,
        # evaluate = True,
        sims_per_step=5,
        warmup_steps=0,
        additional_params={
            "max_accel": 1,
            "max_decel": 1.5,
            "target_velocity": 20,
            "num_rl": NUM_RL,
        },
    ),

    net=NetParams(
        inflows=inflow,
        additional_params=additional_net_params,
    ),
    veh=vehicles,
    initial=InitialConfig(),
)


# pylint: disable=W0621,C0114,C0116,W0212,W0613
import pathlib
import textwrap
import pytest
import numpy as np

from dae.genomic_resources.testing import build_inmemory_test_repository
from dae.genomic_resources.repository import GR_CONF_FILE_NAME, \
    GenomicResourceRepo
from dae.gene.gene_scores import build_gene_score_from_resource, \
    GeneScoreImplementation

@pytest.fixture
def scores_repo(tmp_path: pathlib.Path) -> GenomicResourceRepo:
    scores_repo = build_inmemory_test_repository({
        "LinearHist": {
            GR_CONF_FILE_NAME: """
                type: gene_score
                filename: linear.csv
                scores:
                - id: linear
                  desc: linear gene score
                  histogram:
                    type: number
                    number_of_bins: 3
                    x_log_scale: false
                    y_log_scale: false
                """,
            "linear.csv": textwrap.dedent("""
                gene,linear
                G1,1
                G2,2
                G3,3
                G4,1
                G5,2
                G6,3
            """),
            "statistics": {
                "histogram_linear.yaml": textwrap.dedent("""
                    bars:
                    - 2
                    - 2
                    - 2
                    bins:
                    - 1.0
                    - 1.665
                    - 2.333
                    - 3.0
                    config:
                      type: number
                      number_of_bins: 3
                      score: linear
                      view_range:
                        max: 3.0
                        min: 1.0
                      x_log_scale: false
                      y_log_scale: false
                """)
            }
        },
        "LogHist": {
            GR_CONF_FILE_NAME: """
                type: gene_score
                filename: log.csv
                scores:
                - id: log
                  desc: log gene score
                  histogram:
                    type: number
                    number_of_bins: 5
                    view_range:
                      min: 0.0
                      max: 1.0
                    x_min_log: 0.001
                    x_log_scale: true
                    y_log_scale: false
                """,
            "log.csv": textwrap.dedent("""
                gene,log
                G1,0
                G2,0.0001
                G3,0.001
                G4,0.01
                G5,0.1
                G6,1.0
            """),
            "statistics": {
                "histogram_log.yaml": textwrap.dedent("""
                    bars:
                    - 2
                    - 1
                    - 1
                    - 1
                    - 1
                    bins:
                    - 0.0,
                    - 0.001,
                    - 0.005623413251903491,
                    - 0.03162277660168379,
                    - 0.1778279410038923,
                    - 1.0
                    config:
                      type: number
                      number_of_bins: 5
                      view_range:
                        min: 0.0
                        max: 1.0
                      x_min_log: 0.001
                      x_log_scale: true
                      y_log_scale: false
                """)
            }
        },
        "NaNTest": {
            GR_CONF_FILE_NAME: """
                type: gene_score
                filename: scores.csv
                scores:
                - id: score1
                  desc: linear gene score
                  histogram:
                    type: number
                    number_of_bins: 3
                    x_log_scale: false
                    y_log_scale: false
                """,
            "scores.csv": textwrap.dedent("""
                gene,score1
                G1,1
                G2,2
                G3,nan
                G4,1
                G5,nan
                G6,3
            """),
            "statistics": {
                "histogram_linear.yaml": textwrap.dedent("""
                    bars:
                    - 2
                    - 2
                    - 2
                    bins:
                    - 1.0
                    - 1.665
                    - 2.333
                    - 3.0
                    config:
                      type: number
                      number_of_bins: 3
                      score: linear
                      view_range:
                        max: 3.0
                        min: 1.0
                      x_log_scale: false
                      y_log_scale: false
                """)
            }
        },
        "Oops": {
            GR_CONF_FILE_NAME: "",
        },
        "OopsHist": {
            GR_CONF_FILE_NAME: """
                type: gene_score
                filename: oops.csv
                scores:
                - id: linear
                  desc: linear gene score
                """,
            "oops.csv": textwrap.dedent("""
                gene,linear
                G1,1
                G2,2
                G3,3
            """)
        },
        "OopsScores": {
            GR_CONF_FILE_NAME: """
                type: gene_score
                filename: oops.csv
                """,
            "oops.csv": textwrap.dedent("""
                gene,linear
                G1,1
                G2,2
                G3,3
            """)
        },
    })
    return scores_repo

def test_load_linear_gene_scores_from_resource(
        scores_repo: GenomicResourceRepo) -> None:

    res = scores_repo.get_resource("LinearHist")
    assert res.get_type() == "gene_score"

    result = build_gene_score_from_resource(res)
    scores = result.get_scores()
    assert len(scores) == 1
    score_id = scores[0]

    assert result.get_x_scale(score_id) == "linear"
    assert result.get_y_scale(score_id) == "linear"

    hist = result.get_histogram(score_id)
    assert hist is not None
    assert len(hist.bins) == 4

    assert np.all(hist.bars == np.array([2, 2, 2]))

def test_load_log_gene_scores_from_resource(
        scores_repo: GenomicResourceRepo) -> None:

    res = scores_repo.get_resource("LogHist")
    assert res.get_type() == "gene_score"

    result = build_gene_score_from_resource(res)
    scores = result.get_scores()
    assert len(scores) == 1
    score_id = scores[0]

    assert result.get_x_scale(score_id) == "log"
    assert result.get_y_scale(score_id) == "linear"

    hist = result.get_histogram(score_id)
    assert hist is not None
    assert len(hist.bins) == 6

    assert np.all(hist.bars == np.array([2, 1, 1, 1, 1]))

def test_load_wrong_resource_type(
        scores_repo: GenomicResourceRepo) -> None:
    res = scores_repo.get_resource("Oops")
    with pytest.raises(ValueError, match="invalid resource type Oops"):
        build_gene_score_from_resource(res)

def test_load_gene_score_without_histogram(
        scores_repo: GenomicResourceRepo) -> None:
    res = scores_repo.get_resource("OopsHist")
    with pytest.raises(
        ValueError,
        match="Missing histogram config for linear in OopsHist"
    ):
        build_gene_score_from_resource(res)

def test_load_gene_score_without_gene_scores(
        scores_repo: GenomicResourceRepo) -> None:
    res = scores_repo.get_resource("OopsScores")
    with pytest.raises(ValueError,
                       match="missing scores config in OopsScores"):
        build_gene_score_from_resource(res)

def test_gene_score(scores_repo: GenomicResourceRepo) -> None:

    res = scores_repo.get_resource("LinearHist")
    gene_score = build_gene_score_from_resource(res)

    assert gene_score is not None

    assert gene_score.get_gene_value("linear", "G2") == 2
    assert gene_score.get_gene_value("linear", "G3") == 3

def test_gene_score_nan(scores_repo: GenomicResourceRepo) -> None:
    res = scores_repo.get_resource("NaNTest")
    gene_score = build_gene_score_from_resource(res)

    assert len(gene_score.df) == 6
    df = gene_score.get_score_df("score1")
    assert len(df) == 4

def test_calculate_histogram(scores_repo: GenomicResourceRepo) -> None:
    res = scores_repo.get_resource("LinearHist")
    result = build_gene_score_from_resource(res)
    assert result is not None

    histogram = GeneScoreImplementation._calc_histogram(res, "linear")
    assert histogram is not None
    print(histogram.config.view_range[0])
    print(type(histogram.config.view_range[0]))
    print(histogram.serialize())


import sys
from PyQt4.QtGui import (QApplication, QWidget, QFont, QListWidget,
    QHBoxLayout, QVBoxLayout, QShortcut, QKeySequence)
import numpy as np
from spyderlib.widgets.sourcecode.codeeditor import CodeEditor
from spyderlib.widgets.internalshell import InternalShell
from spyderlib.widgets.dicteditor import DictEditorWidget

class Demo(QWidget):
    def __init__(self):
        super(Demo, self).__init__()
        self.code_editor = CodeEditor(self)
        self.code_editor.setup_editor(
            language = "python",
            font = QFont("Courier New")
        )
        run_sc = QShortcut(QKeySequence("F5"), self, self.run)

        self.shell = InternalShell(self, {"demo":self},
            multithreaded = False,
            max_line_count = 3000,
            font = QFont("Courier new", 10),
            message='caonima'
        )

        self.dict_editor = DictEditorWidget(self, {})
        self.dict_editor.editor.set_filter(self.filter_namespace)
        self.dict_editor.set_data(self.shell.interpreter.namespace)
        vbox = QVBoxLayout()
        vbox.addWidget(self.code_editor)
        vbox.addWidget(self.shell)

        hbox = QHBoxLayout()
        hbox.addWidget(self.dict_editor)
        hbox.addLayout(vbox)

        self.setLayout(hbox)
        self.resize(800, 600)

    def filter_namespace(self, data):
        result = {}
        support_types = [np.ndarray, int, float, str, tuple, dict, list]
        for key, value in data.items():
            if not key.startswith("__") and type(value) in support_types:
                result[key] = value
        return result

    def run(self):
        code = str(self.code_editor.toPlainText())
        namespace = self.shell.interpreter.namespace
        exec (code,namespace )
        self.dict_editor.set_data(namespace)

if __name__ == '__main__':
    import sys
    app = QApplication(sys.argv)
    try:
        demo = Demo()
        demo.show()
    except Exception as ex:
        import traceback
        sys.__stdout__.write(traceback.format_exc())
    sys.exit(app.exec_())


from odoo import fields, models

class ResPartnerAcademy(models.Model):
    """ National Education Academy """

    _name = 'res.partner.academy'
    _description = 'Academy'
    _rec_name = 'name'
    _order = 'name'

    partner_id = fields.Many2one('res.partner', 'Rectorate')
    name = fields.Text('Name', required=True)
    email_domain = fields.Char(
        help="Filling this domain (part after the @) will help to "
        "quickly identify a partner academy using its e-mail suffix"
    )
    logo = fields.Binary(
        related='partner_id.image',
        string="Academy Logo",
        readonly=False,
    )
    state_id = fields.Many2one(
        "res.country.state",
        related="partner_id.state_id",
        string='State',
        readonly=False,
    )
    country_id = fields.Many2one(
        "res.country",
        related="partner_id.country_id",
        string='Country',
    )
    department_ids = fields.Many2many(
        'res.country.department',
        string='Departments',
        help='Departments of this academy',
        domain="[('country_id', '=', country_id), ('state_id', '=', state_id)]",
    )


import numpy as np
import _pickle as pickle
import copy

from collections import OrderedDict
from ConstantVariables import *
from pathlib import Path

class Network:
    def __init__(self, name='', inputs=0, outputs=0):
        self.name = name
        self.inputs = inputs
        self.outputs = outputs
        self.layers_info = []
        self.layers = {}
        self.connections = OrderedDict()
        self.input_data_set = []
        self.output_data_set = []

    def add_layer(self, amount):
        self.layers_info.append(amount)

    def add_data_set(self, input_data_set, output_data_set):
        if len(input_data_set) == len(output_data_set):
            self.input_data_set = input_data_set
            self.output_data_set = output_data_set

    def connect(self):
        self.layers['input'] = 0

        current_layer = self.inputs
        current_layer_name = 'input'
        for index, value in enumerate(self.layers_info):
            self.layers['hidden_' + str(index)] = 0
            self.connections[current_layer_name +'_to_' + 'hidden_' + str(index)] = 2 * np.random.random((current_layer, value)) - 1
            current_layer = value
            current_layer_name = 'hidden_' + str(index)

        self.layers['output'] = 0
        self.connections[current_layer_name +'_to_output'] = 2 * np.random.random((current_layer, self.outputs)) - 1

    def activate(self, input_data):
        self.layers['input'] = input_data
        current_layer_name = 'input'
        for index, value in enumerate(self.layers_info):
            connections = self.connections[current_layer_name +'_to_' + 'hidden_' + str(index)]

            self.layers['hidden_' + str(index)] = self.sigmoid(np.dot(self.layers[current_layer_name], connections))

            current_layer_name = 'hidden_' + str(index)

        connections = self.connections[current_layer_name + '_to_output']
        self.layers['output'] = self.sigmoid(np.dot(self.layers[current_layer_name], connections))

    def learn(self, learning_rate=0.1):
        self.activate(self.input_data_set)

        l2_error = self.output_data_set - self.layers['output']
        print("Error:" + str(np.mean(np.abs(l2_error))))

        error = l2_error
        delta = None
        for name in reversed(self.connections):
            start, end = name.split('_to_')

            self.connections[name], delta = self.get_weight(
                self.connections[name],
                self.layers[start],
                self.layers[end],
                error,
                learning_rate
            )

            error = delta.dot(self.connections[name].T)

    def get_weight(self, connections_layer, start_layer, end_layer, error, learning_rate):
        delta = error * self.sigmoid(end_layer, deriv=True)

        connections_layer += start_layer.T.dot(delta) * learning_rate

        return connections_layer, delta

    def print_output(self, converted=False):
        output = self.get_output(convert=converted)
        if converted:
            if output.shape == (4,):
                if output[0]:
                    print('up')
                if output[1]:
                    print('right')
                if output[2]:
                    print('down')
                if output[3]:
                    print('left')
            else:
                for result in output:
                    print(result)
                    if result[0]:
                        print('up')
                    if result[1]:
                        print('right')
                    if result[2]:
                        print('down')
                    if result[3]:
                        print('left')
                    print('---------')
        else:
            print(self.layers['output'])

    def convert_output(self):
        if len(self.layers['output'].shape) > 1:
            for (x, y), result in np.ndenumerate(self.layers['output']):
                if result > 0.6:
                    self.layers['output'][x][y] = 1
                else:
                    self.layers['output'][x][y] = 0
        else:
            for x, result in np.ndenumerate(self.layers['output']):
                if result > 0.6:
                    self.layers['output'][x] = 1
                else:
                    self.layers['output'][x] = 0

    def get_output(self, convert=False):
        if convert:
            self.convert_output()
        return self.layers['output']

    def sigmoid(self, x, deriv=False):
        if deriv:
            return x * (1 - x)
        return 1 / (1 + np.exp(-x))

    def save(self):
        with open(networks_directory_save + 'Network_' + self.name + '.dump', 'wb') as output:
            pickle.dump(copy.copy(self), output, -1)

    def load(self):
        file = Path(networks_directory_save + 'Network_' + self.name + '.dump')
        if file.exists():
            with open(networks_directory_save + 'Network_' + self.name + '.dump', 'rb') as input:
                network = pickle.load(input)
                return network
        else:
            with open(networks_directory_save + 'Network_1.dump', 'rb') as input:
                network = pickle.load(input)
                network.name = self.name
                # network.save()
                return network


import solution

class Solution(solution.Solution):
    def solve(self, test_input=None):
        logs, k = test_input
        return self.findingUsersActiveMinutes([x[:] for x in logs], k)

    def findingUsersActiveMinutes(self, logs, k):
        """
        :type logs: List[List[int]]
        :type k: int
        :rtype: List[int]
        """
        from collections import defaultdict
        d = defaultdict(set)
        for i,t in logs:
            d[i].add(t)
        ans = [0] * k
        for key in d:
            ans[len(d[key]) - 1] += 1
        return ans