import numpy as np
import matplotlib.pyplot as plt

def plot(x, y, labels, name="plot.png", **kwargs):
    fig, ax = plt.subplots(1,1, dpi=150)
    ax.clear()
    colours = ["xkcd:marine blue", "xkcd:crimson", "xkcd:greenish blue", "xkcd:electric purple"]
    if len(y.shape) == 1:
        ax.semilogy(x, y, color=colours[0], label=labels[0], **kwargs)
    elif len(x.shape) > 1:
        for i in range(y.shape[1]):
            ax.semilogy(x[:,i], y[:,i], label=labels[i], color=colours[i], **kwargs)
    elif len(y.shape) > 1:
        for i in range(y.shape[1]):
            ax.semilogy(x, y[:,i], label=labels[i], color=colours[i], **kwargs)
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Magnitude")
    ax.legend()
    ax.set_xlim([150,300])
    ax.grid(linestyle=":")


def coupling(freq, darm_bkg, darm_inj, witness_bkg, witness_inj, PSD=True):
    darm_threshold  = 2
    witness_threshold = 2.5
    if PSD:
        array_test = np.ones(freq.shape)
        array_test = np.ones(freq.shape)
        check_darm = np.where(np.sqrt(np.abs(darm_inj/darm_bkg)) > darm_threshold, array_test, array_test*0)      
        check_witness = np.where(np.sqrt(np.abs(witness_inj/witness_bkg)) > witness_threshold, array_test, array_test*0) 
        upper_limit = np.trim_zeros(np.where(np.logical_and(check_darm == 0,check_witness == 1), np.sqrt(np.abs((darm_inj)/(witness_inj - witness_bkg))), 0))
        up_lim_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 0, check_witness == 1),  freq, 0))
        upper_limit = upper_limit[upper_limit != 0]
        up_lim_freq = up_lim_freq[up_lim_freq != 0]
        lower_limit = np.trim_zeros(np.where(np.logical_and(check_darm == 1,check_witness == 0), np.sqrt(np.abs((darm_inj-darm_bkg)/(witness_inj))), 0))
        low_lim_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 1, check_witness == 0),  freq, 0))
        lower_limit = lower_limit[lower_limit != 0]
        low_lim_freq = low_lim_freq[low_lim_freq != 0]
        coupling = np.trim_zeros(np.where(np.logical_and(check_darm == 1,check_witness == 1),np.sqrt(np.abs((darm_inj-darm_bkg)/(witness_inj - witness_bkg))), 0))
        coupling_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 1, check_witness == 1),  freq, 0))
        coupling = coupling[coupling != 0]
        coupling_freq = coupling_freq[coupling_freq != 0]
        proj_coupling = np.trim_zeros(np.where(np.logical_and(check_darm == 1,check_witness == 1), witness_bkg * np.abs((darm_inj-darm_bkg)/(witness_inj - witness_bkg)), 0))
        proj_coupling_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 1, check_witness == 1),  freq, 0))
        proj_coupling = proj_coupling[proj_coupling != 0]
        proj_coupling_freq = proj_coupling_freq[proj_coupling_freq != 0]
        proj_upper_limit = np.trim_zeros(np.where(np.logical_and(check_darm == 0,check_witness == 1), witness_bkg * np.abs((darm_inj)/(witness_inj - witness_bkg)), 0))
        proj_up_lim_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 0, check_witness == 1),  freq, 0))
        proj_upper_limit = proj_upper_limit[proj_upper_limit != 0]
        proj_up_lim_freq = proj_up_lim_freq[proj_up_lim_freq != 0]
        proj_lower_limit = np.trim_zeros(np.where(np.logical_and(check_darm == 1,check_witness == 0), witness_bkg*np.abs((darm_inj-darm_bkg)/(witness_inj)), 0))
        proj_low_lim_freq = np.trim_zeros(np.where(np.logical_and(check_darm == 1, check_witness == 0),  freq, 0))
        proj_lower_limit = proj_lower_limit[proj_lower_limit != 0]
        proj_low_lim_freq = proj_low_lim_freq[proj_low_lim_freq != 0]
        coupling_f = np.sqrt(np.abs((darm_inj - darm_bkg)/(witness_inj - witness_bkg)))
        fig, [ax1, ax2] = plt.subplots(2,1,dpi=300, sharex=True, figsize=(8,6))
        ax1.semilogy(coupling_freq, coupling, label="coupling factor", color="xkcd:marine blue", marker=".", linestyle="none")
        ax1.semilogy(up_lim_freq, upper_limit, label="upper limit", color="xkcd:crimson", marker="x", linestyle="none")
        ax1.semilogy(low_lim_freq, lower_limit, label="lower limit", color="xkcd:greenish blue", marker="x", linestyle="none")
        ax1.set_xlim([200,250])
        ax1.set_ylim([5,30])
        ax1.grid(linestyle=":", which="major", color="xkcd:medium grey")
        ax1.grid(linestyle=":",which="minor", color="xkcd:light grey")
        #ax1.set_xlabel("Frequency [Hz]")
        ax1.set_ylabel("Magnitude")
        ax1.legend()
        ax1.grid(linestyle=":")
        ax1.set_title("Coupling")
        ax2.semilogy(proj_coupling_freq, proj_coupling, label="coupling factor", color="xkcd:marine blue", marker=".", linestyle="none")
        ax2.semilogy(proj_up_lim_freq, proj_upper_limit, label="upper limit", color="xkcd:crimson", marker="x", linestyle="none")
        ax2.semilogy(proj_low_lim_freq, proj_lower_limit, label="lower limit", color="xkcd:greenish blue", marker="x", linestyle="none")
        ax2.semilogy(freq, darm_bkg, label="DARM", color="xkcd:electric purple")
        ax2.semilogy(freq, witness_bkg*81, label="Estimate (CF = 9)", color="xkcd:electric blue", linestyle='--')
        ax2.set_xlim([200,250])
        ax2.set_ylim([.5e-4, 1e-2])
        ax2.set_xlabel("Frequency [Hz]")
        ax2.set_ylabel('Magnitude')
        ax2.legend()
        ax2.grid(linestyle=":", which="major", color="xkcd:medium grey")
        ax2.grid(linestyle=":",which="minor", color="xkcd:light grey")
        ax2.set_title("Projection to DARM")
        fig.savefig("DARM_PROJECTION_IMC_REFL_QPDA2_P_PITCH_0.6_ESTIMATE_1.png")
        #plot(freq, np.array([darm_bkg, witness_bkg * coupling_f**2]).T, labels=["DARM Injection", "Witness Projection"])
    

data = np.loadtxt("QPDA2_P_DARM_jitter_coupling.txt")
freq = data[:,0]
DARM_inj_06 = data[:,1] 
DARM_inj_03  = data[:,2] 
DARM_bkg = data[:,3] 
REFL_QPD_inj_06  = data[:,4] 
REFL_QPD_inj_03 = data[:,5] 
REFL_QPD_bkg = data[:,6] 

coupling(freq, DARM_bkg, DARM_inj_06, REFL_QPD_bkg, REFL_QPD_inj_06)