import numpy as np
import scipy.signal
import scipy.optimize


# Daniel Burk, Michigan State University

def periodrange(range): # range is a value between 1 and 3
    Period = []
    Period.append(np.concatenate((np.arange(30.0,5.0,-5.0),np.arange(9.0,2.5,-1.5),np.arange(2.5,1.0,-0.1),[0.75,0.666,0.5,0.333,0.25,0.2,0.125]),axis=None))
    Period.append(np.concatenate((np.arange(20.0,10.0,-2.5),np.arange(9.0,4.0,-1.5),np.arange(4.0,2.0,-0.5),np.arange(2.0,0.5,-0.1),[0.5,0.333,0.25,0.2,0.125]),axis=None))
    Period.append(np.concatenate((np.arange(10.0,4.0,-1.0),np.arange(4.0,2.0,-0.5),np.arange(2.0,0.5,-0.1),[0.5,0.333,0.25,0.2,0.125,0.01,0.005]),axis=None))
    frequencies = (1/np.array(Period[range-1]))
    return(Period[range-1],frequencies)

def pazto_freq_resp(freqs, zeros, poles, scale_fac):
    b, a = scipy.signal.ltisys.zpk2tf(zeros, poles, scale_fac)
    if not isinstance(a, np.ndarray) and a == 1.0:
        a = [1.0]
    return scipy.signal.freqs(b, a, freqs * 2 * np.pi)[1]
            # list element 0 is frequencies
            # list element 1 is the complex amplitudes

def phasecalc(testresponse):  # Bring in a list of complex numbers and return the angle between 90 and 270 degrees
    testphase = []
    for t in testresponse:
        tp = np.arctan2(t.imag , t.real) * 180. / np.pi
        if tp > 90.:
            tp = tp-360.
        testphase.append(tp - 90.0) # adjust phase to better match what is seen in the real world calibrations
    return(testphase)


# This is the function definition that is used in the scipy.minimize function.
def minimize(_var):     # Uses data found in frequencies, and in response.
                        # Make sure phase and response tables use the same subset of frequencies.
    p1r, p1i, p3r, p4r, p5r,z1r,z2r,z3r, scale_fac = _var
    new_resp = pazto_freq_resp(
        freqs=frequencies,
        zeros=np.array([0.0 + 0.0 * 1j,
                        0.0 + 0.0 * 1j,
                        z1r + 0.0 * 1j,
                        z2r + 0.0 * 1j,
                        z3r + 0.0 * 1j], dtype=np.complex128),
        poles=np.array([p1r + p1i * 1j,
                        p1r - p1i * 1j,
                        p3r + 0.0 * 1j,
                        p4r + 0.0 * 1j,
                        p5r + 0.0 * 1j], dtype=np.complex128),
        scale_fac=scale_fac)
    return ((np.abs(new_resp) - np.abs(response)) ** 2).sum()



def getpaz(frequencies,response,Phasedeg):
    # frequencies = np.array([])
    # response = np.array9[],dtype=np.float32)
    # Phasedeg = []
    evaluation = 1.0E+09 # For evaluating how close the solution is to the original curve
    paz = [] # The new poles and zeros for the channels

    for z in range(0,32): # iterate 32 times to find the solution that best describes the phase response.
        initial_x=[]
        X0=np.random.random(9)
        #                                Using the minimize function, find the poles & zeros solution that best describes
        #                                the instrument response as found in responses, on frequencies breakpoint "frequencies"
        out = scipy.optimize.minimize(
            fun=minimize,
            method="BFGS",
            x0=X0,
            options={"eps": 1e-10, "maxiter": 1e8})

        x = out.x
        new_poles = np.array([-abs(x[0]) + abs(x[1]) * 1j,
                              -abs(x[0]) - abs(x[1]) * 1j,
                              -abs(x[2]) + 0.0 * 1j,
                              -abs(x[3]) + 0.0 * 1j,
                              -abs(x[4]) + 0.0 * 1j],
                              dtype=np.complex128)

        new_zeros = np.array([ 0.0 + 0.0 * 1j,
                               0.0 + 0.0 * 1j,
                              x[5] + 0.0 * 1j,
                              x[6] + 0.0 * 1j,
                              x[7] + 0.0 * 1j], dtype=np.complex128)
        new_scale_fac = x[8]
        #              Create the response curve that results from this theoretical new poles and zeroes solution
        inverted_response = pazto_freq_resp(freqs=frequencies, zeros=new_zeros, poles=new_poles,scale_fac=new_scale_fac)
        inphase = phasecalc(inverted_response) # phase from inverted response, listed in degrees
        curvefit = np.sqrt(((np.array(Phasedeg) - np.array(inphase))**2).mean()) # rmse
        if (curvefit) < evaluation:
            final_iteration = z
            best_poles=new_poles
            best_zeros=new_zeros
            best_scale_fac=new_scale_fac
            print(f'Iteration # {z}: Phase misfit drops to {curvefit}')
            evaluation = curvefit
    return(best_poles,best_zeros,best_scale_fac,evaluation,z)

def paztest_fixed():
    #################################### test with def #################################
    Component = 'LM.NE8K.MHZ'
    Caldate   = '05/15/2019'
    #              Frequency breakpoints within the passband of the seismometer to simulate
    frequencies = np.array([0.05, 0.0571, 0.0667, 0.080, 0.111, 0.133, 0.167,
                   0.222, 0.250, 0.286, 0.333, 0.400, 0.500, 0.526,
                   0.555, 0.588, 0.625, 0.666, 0.714, 0.770, 0.833,
                   0.910, 1.000, 1.111, 1.250, 1.429, 1.667, 2.000,
                   3.000, 4.000, 5.000, 8.000])

    #               here are the gain values for the seismometer at the above frequencies
    response = np.array([3.00, 4.48, 7.11, 12.28, 32.81, 56.56, 110.00, 258.43,
            366.09, 542.60, 852.84, 1451.12, 2764.14, 3201.37, 3734.65,
            4390.91, 5205.66, 6225.33, 7508.61, 9123.91, 11134.60,
            13556.01, 16267.16, 18911.45, 20951.61, 22004.53, 22120.93,
            21630.53, 20132.44, 18990.64, 17947.77, 15053.22]) #,dtype=np.float32)

    #           phase delay of the light beam vs. ground motion in degrees at above frequencies
    Phasedeg = [-6.660, -7.714, -8.880, -10.800, -14.800, -18.000, -22.200,
                -30.000, -33.300, -37.543, -44.400, -52.560, -66.600,
                -69.158, -73.800, -78.353, -83.475, -89.520, -96.429,
                -104.677, -114.600, -126.654, -140.760, -157.600, -175.500,
                -193.886, -211.560, -228.024, -254.865, -269.640, -280.080,
                -300.528]

    best_poles,best_zeros,best_scale_fac,evaluation,final_iteration = getpaz(frequencies,response,Phasedeg)

    print("n========================================")
    print(f"Inverted values for {Component} on caldate of {Caldate}:")
    print("Poles =n", best_poles)
    print("Zeros =n", best_zeros)
    print("Scale_fac = ", best_scale_fac)
    print(f"Evaluated misfit of phase = {evaluation} on iteration # {final_iteration}n")
    print("========================================")


##############################  RUN CODE AS A DEF ######################
#Component = 'LM.NE8K.MHZ'
paztest_fixed()