#####
# Need input .cvs file with energy, gap.  Will not work with column headers
# Use doc string of set_table() to check/test table before uploading
# Epics table generated by interpolation, but one could add an equation to autogenerate or develop a fit function.
# Upper and lower bounds currently set by input.  Can add another parameter and enable extrapolation for the interpolation function
#
# 09/13/17
#
#update on 08/14/19
#using lmfit to create forward and reverse tables. improved ease of use.  some libaries are no longer needed.  New implmentatons is 
#using the function set_EPUioc_table().  See doc-string via set_EPUioc_table??

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import csv
from lmfit.models import PolynomialModel 
from matplotlib.ticker import FormatStrFormatter 

def test_table(Emeas,Gmeas,Earraysize, Garraysize,minE, maxE, minG, maxG):
    '''
    This function requires that you hard code your fit to the gap vs energy data.  energy is the x-variable for the forward calculation.
    This function also requires a hard code fit for the reverse interpolation.  gap is the x-varyable.
    '''
    #print(file)
    print('Input energy min and max are:\t', np.min(Emeas),np.max(Emeas))
    print('Input gap min and max are:\t', np.min(Gmeas),np.max(Gmeas))
    
    
    Esp = (np.max(Emeas)-np.min(Emeas))/Earraysize
    Gsp = (np.max(Gmeas)-np.min(Gmeas))/Earraysize
    print('\nForward interpolation has energy and gap spacing of {:.3f} and {:.4f}'.format(Esp,Gsp))
    
    x = np.linspace(minE, maxE, Earraysize)
    y = np.interp(x, Emeas, Gmeas)
    
    
    
    plt.figure(figsize=(7,7))
    plt.plot(Emeas, Gmeas, 'o', label='input data ')#{}'.format(file))
    plt.plot(x,y,'--r', label='np.interp forward')
    plt.grid(True)
    plt.legend(loc=2)
    plt.show()
    
    
    x2 = np.linspace(minG, maxG, Garraysize)
    y2 = np.interp(x2, Gmeas, Emeas)

    Esp = (np.max(Emeas)-np.min(Emeas))/Garraysize
    Gsp = (np.max(Gmeas)-np.min(Gmeas))/Garraysize
    print('\nReverse interpolation has energy and gap spacing of {:.3f} and {:.4f}'.format(Esp,Gsp))
    
    plt.figure(figsize=(7,7))
    
    plt.plot(Gmeas,Emeas, 'o', label='input data ')#{}'.format(file))
    plt.plot(x2,y2,'--r', label='np.interp reverse')
    plt.grid(True)
    plt.legend(loc=2)
    plt.show()

    return x, y, x2, y2

def make_table(output,IOCtable, EPUnum):
    x=output[0]
    y=output[1]
    x2=y#output[2]  ##2 and 3 are from previous version that needed to reverse interpolation
    y2=x#output[3]
    
    #print('hello')
    
    #pv_stringi = ('XF:02ID-ID{EPU:1-RLT}Val:Table%sa-Wfrm'%(IOCtable))
    #print(pv_string)        
     
    fp = np.concatenate((np.array([x[0], x[1] - x[0], y.size, -1, 1, 3]), y, y, y))
    ######pv = epics.PV('XF:02ID-ID{EPU:1-FLT}Val:Table1a-Wfrm')  ## example to test
    pv = epics.PV('XF:21ID-ID{EPU:%d-FLT}Val:Table%sa-Wfrm'%(EPUnum,IOCtable))
    print(f'\nWriting to PV:\t{pv}\n{fp}\n\n')
    pv.put(fp)
        
    fp = np.concatenate((np.array([x2[0], x2[1] - x2[0], y2.size, -1, 1, 3]), y2, y2, y2))
    ######pv = epics.PV('XF:02ID-ID{EPU:1-RLT}Val:Table1a-Wfrm')  ## example to test
    pv = epics.PV('XF:21ID-ID{EPU:%d-RLT}Val:Table%sa-Wfrm'%(EPUnum,IOCtable))
    print(f'\nWriting to PV:\t{pv}\n{fp}\n\n')
    pv.put(fp)
    

def load_input(file):
    global readCSV
    print('Input file is:\t', file,'\n\n')
    with open(file) as csvfile:
        readCSV = csv.reader(csvfile, delimiter=',')
        # Assume one header line and skip it
        next(readCSV)
        xe = []
        yg = []

        try:
            for row in readCSV:
                print(row)  #use just this line to figure out the delimiter  delete csv text in first line if not a number
                en = float(row[0])
                gp = float(row[1])

                xe.append(en)
                yg.append(gp)
        except ValueError as ex:
            print('\n\n**********\nError in extracting data from CSV.  Possibly due to header.  Remove non-numeric data.\n**********\n\n')
            raise
        return xe, yg


def set_EPU_table(file,pts = [1000,1000], params = None, IOCtable=1, final=False):
    '''
    pts is the number points for energy, gap:
        [0]\t output energy array length, 1000 by default
        [1]\t output gap array length, 1000 by default
        
        
    params is a list of 6 elements being:
        [0]\t output energy array length
        [1]\t output gap array length
        [2]\t output minimum energy
        [3]\t output maximum energy
        [4]\t output minimum gap
        [5]\t output maximum gap
	
	Use default value for intial test
	
    IOCtable\t EPICS PV for the waveform table number
	
    final\t change to True to write to the waveform table PV

     '''
    data = load_input(file)
    if params is None:
        E0 = np.min(data[0])
        E1 = np.max(data[0])
        G0 = np.min(data[1])
        G1 = np.max(data[1])
        
        output = test_table(data[0],data[1],pts[0],pts[1],E0,E1,G0,G1)
    else:
        output = test_table(data[0],data[1],params[0], params[1],params[2], params[3], params[4], params[5])
    if final is True:
        print('Laoding tables to IOC')
        if __name__ == "__main__":
            print('about to make the table') 
            make_table(output, EPUnum)
        else:
            print('the name space is', __name__)

def fit_input_make_output(xe, yg, Epts = 1000, Emin = None, Emax = None, polydeg = 7, file=None, one2three=False):
    mod = PolynomialModel(polydeg)
    data = np.array([xe,yg])
    pars = mod.guess(yg, x=xe)
    out = mod.fit(yg, pars, x=xe)
    
    if Emin == None:
        Emin = np.min(xe)
    if Emax == None:
        Emax = np.max(xe)
    if file == None:
        file = 'No File'

    xnew = np.linspace(Emin, Emax, Epts)
    ynew = out.eval(x=xnew)
    output = np.array([xnew, ynew])

    if one2three is True:
        harmonic = 'EPU gap from 3xEnergy'
    else:
        harmonic = 'EPU gap with measured Energy'
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_title(f'{file}\n {harmonic} --- fit: deg {mod.poly_degree} polynomial')
    ax.set_xlabel('Energy [eV]')
    ax.set_ylabel('Gap [mm]')
    ax.plot(xnew, ynew, 'r-', label = 'fit data for table')
    ax.plot(xe, yg,'bo' , label='measured data')
    ax.plot(xe,out.residual,'g:',label = 'residual')
    ax.grid(True)
    ax.legend()
    print(f'Report from Fit:\[n {out.fit_report(show_correl = False)}\n')
    print(f'Created table data with {Epts} points and {(xnew[1]-xnew[0])*1000 :.2f} meV resolution.')

    return output

def set_EPUioc_table(file,  Epts=1000, ioc_table_num=None, write_to_ioc=False, EPUnum=None, Emin = None, Emax = None, third_fm_first = False, polydeg = 7 ):
    '''Function fits epu gap vs beamline energy to create a lookup table for automatically setting the epu gap for a given energy and polarization.
    You can find the branch for this function at https://github.com/NSLS-II-CSX/undcontrol/tree/six
    variable\t      :   \tdetails
    file\t          :   \t.csv file.  format used from Valentina.  See example format and crate file with same (entry number, energy, gap)
    Epts\t          :   \tnumber of points.  (Emax-Emin)/Epts = energy resolution of the table
    ioc_table_no\t  :   \tcreates PV assocaite with specific table in the beamline undcontrol IOC. See the undcontrolApp associated with the gitrepo.
    write_to_ioc\t  :   \tdefault False to test function and visualize results before writing to ioc. set to True to write to ioc PV.
    Emin\t          :   \tMinimum energy of output to ioc.  Default is the minimum energy found in the .csv input file
    Emax\t          :   \tMaximum energy of output to ioc.  Default is the maximum energy found in the .csv input filei
    third_fm_first\t:   \tSet to true to create the EPU 3rd harmonic lookup table given that the input is for the 1st harmonic    
    polydeg\t       :   \tparameter for the polynomial fit (degree).  7 is default and maximum number of degrees.
    '''
    xe, yg = load_input(file)
    if third_fm_first is True:
        xe = np.array(xe)*3
    
    if Emin == None:
        Emin = np.min(xe)
    if Emax == None:
        Emax = np.max(xe)

    output = fit_input_make_output(xe, yg, Epts, Emin, Emax, polydeg, file, third_fm_first)
    resol = output[0][1]-output[0][0]
    print('\tminimum Energy and Gap:\t', [np.min(x) for x in output])
    print('\tmaximum Energy and Gap:\t', [np.max(x) for x in output])
    
    dE_x, dG_x = [np.round(np.convolve(x, np.ones(2), 'valid') / 2,3) for x in output]  
    dE_y, dG_y = [np.diff(x) for x in output]
    
    fig = plt.figure(figsize=(9,5))
    fig.suptitle(f'Lookup Table Resolution for IOC table {ioc_table_num}\nEmin={Emin:.1f} Emax={Emax:.1f} Epts={Epts} Eres={resol:.4f}')
    ax = fig.add_subplot(121)
    ax.plot(dE_x,dE_y*1000,'.',label = 'energy')
    ax.set_xlabel('Input Energy [eV]')
    ax.set_ylabel('$\Delta$E [meV]')
    #ax.set_ylim(np.min(dE_y), np.max(dE_y))
    ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
 
    ax = fig.add_subplot(122)
    ax.plot(dG_x, dG_y, label = 'gap')
    ax.set_xlabel('Input Gap [mm]')
    ax.set_ylabel('$\Delta$G [mm]')

    if write_to_ioc is False:
        print(f'\n\n\t\tJust testing. No data written to IOC for IOC table # {ioc_table_num} and EPU{EPUnum}')
    else:
        print(f'\n\n\t\tOVER-WRITING PV for IOC table # {ioc_table_num} and EPU{EPUnum}')
        
        make_table(output, ioc_table_num, EPUnum)
        






 
