import numpy as np
import matplotlib.pyplot as plt

import sys
sys.path.append('../pyAPT/')
from pyAPT import Z806

'''
config file:
pos_Pay:                  current position of motor Pay
pos_Paz:                  current position of motor Paz


measurement file (n.a)
'''

class ParabolicMirror:
    def __init__(self, config, home = False):
        import time
        time.sleep(0.1)
        try:
            self.connect_motors()
            print ('Connnection with motors sucessful!')
        except:
            print ('Connection to motors failed. Close all programs that'+' use the motors or restart the controllers')
            raise
            
        if home:
            self.home()
            
        self.set_config(config)
        
        
    def set_config(self,config):
        self.Pay.goto(config['pos_Pay'])
        self.Paz.goto(config['pos_Paz'])
        
    def get_config(self):
        config = {}                          
        config['pos_Pay'] = float(self.Pay.position())
        config['pos_Paz'] = float(self.Paz.position())
        
        return config
   
    def connect_motors(self):

        self.Paz = Z806(83835804)
        self.Pay = Z806(83835879)
        self.motors = [self.Pay, self.Paz]
        
    def home(self):
        for motor in self.motors:
            motor.home()
        
    def optimize(self, focus, laser, Cz, wanted_focus_um, n_av = 10, sleep_time = 0.2, output = True, maxiter = 100):
        #this variable is necessary to compensate the attenuation of the laser
        global intmult
        intmult = 1.
        
        global reduce_int
        reduce_int = False
        def to_min():
            global intmult
            
            FWHMxs = []
            FWHMys = []
            maxints = []
            for i in range(n_av):
                measurement = focus.measure()
                FWHMxs.append(measurement['FWHMx'])
                FWHMys.append(measurement['FWHMy'])
                maxints.append(measurement['maxit'])
                
            #if output:
            #    focus.plot_measure(measurement)
            #    plt.pause(0.01)
            del measurement
            if output:
                print ('FWHMxs: ',FWHMxs)
                print ('maxints: ',maxints)
                
            res = np.sqrt(np.array(FWHMxs)**2 + np.array(FWHMys)**2)/np.sqrt(2)*4096/np.array(maxints)*intmult
            mean, error = np.mean(res), np.std(res, ddof = -1)
            
            if np.mean(maxints) > 4e3:
                redurce_int = True
                
            if output:
                print ('mean: ',mean,'error: ', error)
                
            if mean*1e3 < wanted_focus_um:
                if output:
                    print ('The optimization converged. Final value: ',mean)
                return
            
            return mean, error


        
        
        def change_parameter(y, z, zc):
            print ('Parameters: ',(y, z, zc))
            config = {'pos_Pay': y, 'pos_Paz': z}
            
            self.set_config(config)
            
            #config_per = per.get_config()
            
            #config_per['deltaz_Cz'] = zc 
            Cz.goto(zc)
            #per.set_config(config_per)
            return True
        
        def opt_f(x):
            y = x[0]
            z = x[1]
            zc = x[2] 
            if output:
                print ('-------------------------------------------')
            change_parameter(y,z, zc)
            
            #wait for oszillation to decay
            import time
            time.sleep(sleep_time)
            
            return to_min()[0]
        
        config0 = self.get_config()
        #config_per0 = per.get_config()
        config0_focus = focus.get_config()
        config0_focus['measure_full']=True
        focus.set_config(config0_focus)
        x0 = [config0['pos_Pay'],
              config0['pos_Paz'], 
              Cz.position()]
              #config_per0['deltaz_Cz']]
        
        def after_it(res):
            print ('###################')
            print ('step done')
            global reduce_int
            if reduce_int:
                config_laser = laser.get_config()
                config_laser['attenuation'] = config_laser['attenuation']/1.2
                laser.set_config(config_laser)
                intmult *= 1.2
                reduce_int = False
                print ('laser attenuated')
            
        
        # directly Nedler-Mead
        import scipy.optimize
        res = scipy.optimize.minimize(opt_f, x0, method='Nelder-Mead',callback = after_it, options={'disp': output,
                                                                          'initial_simplex': None, 
                                                                          'maxiter': maxiter, 
                                                                          'xatol': 0.00005, 
                                                                          'return_all': False, 
                                                                          'fatol': wanted_focus_um})
        
        if output:
            print (res)
        return res