'''
Library of useful functions for the darmplant repo

Craig Cahillane
Oct 4, 2019
'''
import numpy as np
import scipy.constants as scc

def goodTicks(axis):
    '''
    Returns the y ticks to always by factors of 10 on the y scale, given some matplotlib axis object.
    Use if matplotlib keeps making your y-axis plot ticks spacing greater than 10.

    Input: matplotlib axis object. Like axis = plt.gca(), or axis = fig.add_subplot(111)
    Output: Correct yTicks spaced by factors of 10.
    '''
    ymin, ymax = axis.get_ylim()
    yTicks = np.array([10**x for x in np.arange(np.ceil(np.log10(ymin)), np.ceil(np.log10(ymax)))])
    return yTicks

def ward_model( ff,
                phi,
                zeta,
                E_LO,
                L,
                l_SRC,
                TS,
                TI,
                TE,
                Pbs,
                arm_loss,
                SRC_loss,
                post_SRM_loss,
                M=40.0,
                lamb=1064e-9):
    ''' Calculates the Ward DARM plant model (DC readout, dual recycled Fabry Perot IFO)
    From Ward's Thesis:
        It should be noted that this function is not exact—it depends (as usual)
        on the condition λgw >> L, but it also breaks down at the arm cavity free
        spectral range.  Moreover, this expression of the function ignores the
        effect of the finite signal recycling cavity length, which is a tiny correction

    Inputs:
    ff = frequency vector
    phi = SRC detuning in radians (nominal = pi/2)
    zeta = homodyne angle in radians (nominal = 0)
    E_LO = local oscillator amplitude 20 mA on DCPDs, sqrt(20 mA * 1e-3 / responsivity)
    L = length of the interferometer arms (nominal = 3994.5 meters)
    l_SRC = length of the SRC (nominal = 32.01 meters)
    TS = power transmission of SRM (nominal = 32%)
    TI = power transmission of ITM (nominal = 1.42%)
    TE = power transmission of ETM (nominal = 4 ppm)
    Pbs = power on the beamsplitter    (nominal = input * PRG)
    arm_loss = losses in the arms      (nominal = 38 ppm)
    SRC_loss = losses in the SRC, all grouped into SRM reflectivity losses (nominal = 1%)
    post_SRM_loss = losses from back of SRM to OMC DCPDs (nominal = T_OFI * R_OMs * T_OMC * ModeMatchOMC)

    Outputs:
    DARM TF in [watts/meters], where watts are measured at the OMC DCPDs SUM, and meters are DARM meters
    '''
    ww = 2*np.pi*ff
    w0 = 2*np.pi*scc.c/lamb

    tI = np.sqrt(TI)
    tE = np.sqrt(TE)
    tS = np.sqrt(TS)

    rI = np.sqrt(1 - TI)
    rE = np.sqrt(1 - TE - arm_loss)
    rS = np.sqrt(1 - TS - SRC_loss)

    FSR = scc.c/(2 * L)
    wa  = -FSR * np.log(rI * rE)

    SRC_Phase = -ww * l_SRC/scc.c
    beta = -np.arctan2(ww, wa)
    exp0 = np.exp(1j*(beta + SRC_Phase))

    kappa = 8 * Pbs * w0 / (M * L**2 * ww**2 * (wa**2 + ww**2))

    prefactor = (1 - post_SRM_loss) * E_LO / L * np.sqrt(2 * Pbs * w0**2 / (wa**2 + ww**2))
    numer = tS * exp0 * ((1 - rS*exp0**2) * np.cos(phi) * np.cos(zeta) - (1 + rS*exp0**2) * np.sin(phi) * np.sin(zeta))
    denom = 1 + rS**2 * exp0**4 - 2 * rS * exp0**2 * ( np.cos(2*phi) + kappa * np.sin(2*phi) / 2.0 )

    return np.sqrt(2) * prefactor * numer / denom # mysterious factor of two :(