'''
Copyright (c) 2020 by Xinfei Huang, Brandeis University

v1.0.3

Python 3 is the minimum requirement.

If the library "matplotlib" has not been installed, please go to
"https://matplotlib.org/users/installing.html" for installation instructions. 

This script numerically simulates the circuit [ input - R1 - L - (C/R2) - 0V ], where (C/R)
means that R and C are in parallel. This circuit transform sharp pulses at the input into
wider pulses at the output between L and (C/R). More details on the circuit at
"http://www.bndhep.net/Lab/Filter/Filter.html#Pulse%20Shapers". 

To use, create a txt file named "input.txt" under the same directory, an
example of which is given below. Type in the desired sets of values for
L and C in input.txt at indicated locations. The script will go through all
combinations of L and C and generate a plot of shaped pulse for each
combination. The script will generate a directory 'plots', and then store
all the plots there.

Below is an example of an "input.txt" file. Please copy everything inside "#" and paste
into an empty "input.txt" file as a template. "input.txt" should be under the same
directory as this script. 
##########
Note: This txt file will be read when shaper_sim.py is executed. The script read off parameters from this file. 

Note: Use single space to separate values. 
L (uH): .1 1 3.3 4.7 6.8 10 27
C (nF): .0033 .015 .033 .047 .1 .27 1

input pulse duration (ns): 15

R1 (Ohm): 50
R2 (Ohm): 50
##########

Theory: Let input to be X, output Y. The differential equation governing the output is:
{ d2Y/dt2 + ( R/L + 1/RC )*dY/dt + (2Y-X)/LC = 0 }
A second order differential equation is equivalent to two first order differential
equations. Use the approximation of  { dy/dt = ( y(t+dt) - y(t) )/dt } to numerically solve
the ODE.
The second order ODE system that will be numerically evaluated is:
{ dY/dt = Z }
{ dZ/dt = -( R/L + 1/RC )*Z - (2Y-X)/LC }

Some default parameters:
dt = 0.1 ns
pulse_height = 100 mV
pulse_start: t = 10 ns
'''

def main():

    print( '\nCopyright (c) 2020 by Xinfei Huang, Brandeis University' )

    import os
    

    # check cwd

    cwd = os.getcwd()

    script_dir = os.path.dirname(os.path.realpath(__file__))

    if script_dir != cwd:

        print( '\nError: Please change the current working directory to the \
directory this script is in. ' )

        return None

    new_folder = 'plots' # name of new folder for plots
    
    newdir = cwd + os.sep + new_folder
    
    if new_folder in os.listdir( cwd ):

        print( '\nError: The folder "plots" already exists. Please move, delete or \
rename the directory "plots" before running again. ' )

        return None


    # parameters (SI units)

    dt = 1e-10 # time step of numerical simulation
    t_max = 2e-6 # maximum time
    p_start = 1e-8 # the time when the pulse starts
    p_amp = .1 # the amplitude of the pulse


    # read input.txt and extract lists of L and C and other parameters

    if 'input.txt' not in os.listdir(): # error of no input.txt

        print( '\nError: No file named "input.txt" detected. Please refer to the comments \
in the source code to create one before running again. ' )

        return None

    f_input = open( 'input.txt', 'r' )

    s_input = [ ele.split() for ele in f_input.readlines() if ele != '\n' ]

    f_input.close()

    L_list = []

    C_list = []

    p_duration = None
    R1 = None
    R2 = None

    try: 

        for line in s_input:

            if line[0] == 'L':

                L_list = [ float(ele) for ele in line[2:] ]

            elif line[0] == 'C':

                C_list = [ float(ele) for ele in line[2:] ]

            elif line[0] == 'input':

                p_duration = 1e-9 * float( line[4] )

            elif line[0] == 'R1':

                R1 = float( line[2] )

            elif line[0] == 'R2':

                R2 = float( line[2] )

    except Exception:

        print( '\nError: "input.txt" has empty or invalid parameter(s). ' )

        return None

    if len(L_list)*len(C_list) == 0: 

        print( '\nError: "input.txt" has empty list of L or C. ' )

        return None

    if p_duration==None or R1==None or R2==None:

        print( '\nError: Necessary parameter(s) undefined. ' )

        return None

    print( '\nFinished reading "input.txt". ' )


    # convert L and C to [H] and [F]

    L_list = [ ele*1e-6 for ele in L_list ]

    C_list = [ ele*1e-9 for ele in C_list ]


    # time list

    t_list = [ ele*dt for ele in range( int(t_max/dt) ) ]


    # input pulse list

    x_list = [ 0 for ele in t_list ]

    if 1:

        # square pulse

        for i in range( int(p_start/dt), int( (p_start+p_duration)/dt ) ):

            x_list[i] = p_amp

    if 0:

        # step

        for i in range( int(p_start/dt), len(x_list) ):

            x_list[i] = p_amp

    
    # change directory to 'plots' folder
        
    os.mkdir( newdir )
        
    os.chdir( newdir )

    print( '\nThe "plots" folder is created. Calculation and plotting are started. ' )


    # go through combination of LC
    
    for L in L_list:

        for C in C_list:
            

            # output pulse
            # RLC is a custom function

            y_list = RLC( R1, R2, L, C, x_list, t_list )


            # generate plot
            # plot is a custom function

            plot( y_list, t_list, dt, L, C, p_amp, p_duration )
    

    # change directory back
    
    os.chdir( cwd )


    print( '\nEnd of script. ' )


    return None

    




def fwhh( y_list ):

    '''
    calculate FWHH
    unit being index
    '''

    # calculate half height
    
    y_max = max(y_list)

    hh = .5 * y_max


    # go through y_list and find the index where it cross half height
    # indeces stored in inte

    inte = []

    for i in range(1,len(y_list)):

        if (y_list[i]-hh) * (y_list[i-1]-hh) < 0:

            inte.append(i)

        elif (y_list[i]-hh) == 0:

            inte.append(i)


    # return index difference if only two interceptions occured
    # return 0 otherwise
    
    if len(inte) == 2:
        return inte[1] - inte[0]
    else:
        return 0

        


def RLC( R1, R2, L, C, x_list, t_list ):

    '''
    given input pulse and parameters, numerically calculate output of a RLC pulse shaper
    '''


    # get dt from t_list
    
    dt = t_list[1] - t_list[0]


    # initialize y_list and ydot_list

    y_list = [ 0 for ele in t_list ]

    ydot_list = list( y_list )
    

    # numerically solving for y and ydot
    
    for i in range( 1, len(t_list) ):

        dy = dt * ydot_list[i-1]

        dydot = dt * ( -ydot_list[i-1]*( R1/L + 1/(R2*C) ) \
                       + ( x_list[i-1] - ( 1 + R1/R2 )*y_list[i-1] )/(L*C) )

        y_list[i] = y_list[i-1] + dy
        
        ydot_list[i] = ydot_list[i-1] + dydot

    return y_list


def plot( y_list, t_list, dt, L, C, p_amp, p_duration ):

    '''
    generate a plot for output pulse
    '''

    import matplotlib.pyplot as plt


    # width of output pulse in ns
    # fwhh is a custom function

    fwhh_ind = fwhh( y_list )

    width = dt * fwhh( y_list ) * 1e9


    # time in ns

    t_ns_list = [ ele*(1e9) for ele in t_list ]


    # name of the png file will include the C value and the L value
    # convert 0.033 nF to '0n033' etc

    L_str = [ ele for ele in str(round(float(L*1e6),4)) ]
    ind = L_str.index( '.' )
    L_str[ind] = 'u'
    L_str = ''.join(L_str)

    C_str = [ ele for ele in str(round(float(C*1e9),4)) ]
    ind = C_str.index( '.' )
    C_str[ind] = 'n'
    C_str = ''.join(C_str)


    # plot
    
    # time range of the plot is m times fwhh
    m = 5

    plt.figure()
    if fwhh_ind != 0:
        plt.plot( t_ns_list[ :m*fwhh_ind ], \
                  [ ele*1e3 for ele in y_list ][ :m*fwhh_ind ], \
                  label = 'FWHH = %.2f ns' % width )
    else:
        plt.plot( t_ns_list, [ ele*1e3 for ele in y_list ], \
                  label = 'FWHH = %.2f ns' % width )
    plt.grid()
    plt.xlabel( 't [ns]' )
    plt.ylabel( 'output [mV]' )
    plt.legend()
    plt.title( 'Input = %.2f mV, width = %.2f ns, C = %sF, L = %sH' \
               % ( p_amp*1e3, p_duration*1e9, C_str, L_str ) )
    plt.savefig( 'C'+C_str+'_L'+L_str+'.png' )
    plt.close()


    return None







if __name__ == '__main__':
    main()