from scipy import *
from pylab import *
from pyvirgo.controls import *
from pyvirgo.controls.plots import *
from pyvirgo.controls.locking import *

# load optical matrix (W/m)
execfile('optmat.py')

# compute rms
def rms(f, x):
    "Compute integrated RMS depending on frequency cut-off."
    r = zeros(shape(x))
    df = f[1] - f[0]
    r[len(x) - 1] = x[len(x) - 1]*df
    for i in range(1, len(x)):
        r[len(x) - i - 1] = r[len(x) - i] +  x[len(x) - i - 1]*df
    return r

######## free motion
Z = load('/home/vajente/virgo/AdV/free_motion/free.txt')
f0 = Z[:,0]
z0 = Z[:,1]*1e-6

f0 = f0[1:]
z0 = z0[1:]

FF = zeros((ndof, ndof, nf), dtype=complex)

PLOT = False
FREE = False

####################################################################
######## DARM loop #################################################
####################################################################

# optical transfer function (W/m)
G = opt[dof['DARM'], sig['ASY_DC'], :] 
# mechanics
mech = TF([], [], [0.6], [500], 13e-6, 0)
M = mech.fresp(f)

# try design
comp = TF([-39], [10], [0.6], [5], 1, 0)
C = comp.fresp(f)

filt = TF([1, 10, 30, 315, 600], [0, 1, 1, 0, 0], \
          [0, 0.5, 300, 874], [0.7, 1, 1, 0.7], \
           1, 100)
F = filt.fresp(f)

delay = TF([], [], [], [], 1, 0)
delay.delay = 600e-6
D = delay.fresp(f)
gain = 1 / abs((D*C*G*M*F)[amin(find(f>60))])

if PLOT:
    figure()
    subplot(211)
    loglog(f, abs(gain*D*C*M*F*G))
    ylabel('Abs')
    title('DARM open loop transfer function')
    grid(True)
    axis(xmin=0.01, xmax=5000)
    subplot(212)
    semilogx(f, 180/pi* angle(gain*D*C*M*F*G))
    grid(True)
    ylabel('Phase')
    xlabel('Frequency [Hz]')
    axis(xmin=0.01, xmax=5000)
    savefig('plots/loops_darm_oltf.eps')
    savefig('plots/loops_darm_oltf.png')

if FREE:
    # simulate effect on free motion below 10 Hz
    olt = 1e7*f0**-2
    zc = z0 / olt
    r0 = rms(f0, z0)
    rc = rms(f0, zc)
    
    figure()
    loglog(f0, z0, 'r-')
    loglog(f0, zc, 'b-')
    loglog(f0, rc, 'b:')
    loglog(f0, r0, 'r:')
    grid(True)
    xlabel('Frequency [Hz]')
    ylabel('Residual motion [m/rHz]')
    legend(('Open loop', 'Closed loop'))
    savefig('plots/loops_darm_residual.eps')
    savefig('plots/loops_darm_residual.png')

FF[dof['DARM'], dof['DARM'], :] = gain * F * D * M * C

####################################################################
######## CARM (SSFS) loop ##########################################
####################################################################
# optical transfer function (W/m)
G = opt[dof['CARM'], sig['SYM_SB1_P'], :]

# high frequency rough compensation
comp = TF([25500, 50000], [7, 750], [25500, 50000], [10, 1],  1, 1)
C = comp.fresp(f)

# low frequency rough compensation
comp2 = TF([4.33], [6], [0.55], [10], 1, 1000)
C2 = comp2.fresp(f)

filt = TF([], [0], [0], [0], 1, 23000)
F = filt.fresp(f)

delay = TF([],[],[],[],1,1)
delay.delay = 370e-9
D = delay.fresp(f)

gain = 10/ abs((D*G*C2*C*F)[amin(find(f>2300))])

if PLOT:
    figure()
    subplot(211)
    loglog(f, abs(gain*D*C*C2*F*G))
    ylabel('Abs')
    title('SSFS open loop transfer function')
    grid(True)
    axis(xmin=0.01, xmax=5000)
    subplot(212)
    semilogx(f, 180/pi* angle(gain*D*C*C2*F*G))
    grid(True)
    ylabel('Phase')
    xlabel('Frequency [Hz]')
    axis(xmin=0.01, xmax=5000)
    savefig('plots/loops_ssfs_oltf.eps')
    savefig('plots/loops_ssfs_oltf.png')

FF[dof['CARM'], dof['CARM'], :] = gain*D*C*C2*F

####################################################################
######## PRCL loop #################################################
####################################################################

# optical transfer function
G = opt[dof['PRCL'], sig['SYM_3P1_P'], :]
# mechanics
mech = TF([], [], [0.6], [500], 13e-6, 0)
M = mech.fresp(f)

filt = TF([3, 10], [0, 1], [0, 300, 500], [0.7, 1, 1], 1, 1)
F = filt.fresp(f)

delay = TF([],[],[],[],1,1)
delay.delay = 600e-6
D = delay.fresp(f)

gain = 1/ abs((G*M*D*F)[amin(find(f>40))])

if PLOT:
    figure()
    subplot(211)
    loglog(f, abs(gain*D*M*F*G))
    ylabel('Abs')
    title('PRCL open loop transfer function')
    grid(True)
    axis(xmin=0.01, xmax=5000)
    subplot(212)
    semilogx(f, 180/pi* angle(gain*D*M*F*G))
    grid(True)
    ylabel('Phase')
    xlabel('Frequency [Hz]')
    axis(xmin=0.01, xmax=5000)
    savefig('plots/loops_prcl_oltf.eps')
    savefig('plots/loops_prcl_oltf.png')

if FREE:
    # simulate effect on free motion below 10 Hz
    olt = 2e4*f0**-2
    zc = z0 / olt
    r0 = rms(f0, z0)
    rc = rms(f0, zc)
    
    figure()
    loglog(f0, z0, 'r-')
    loglog(f0, zc, 'b-')
    loglog(f0, rc, 'b:')
    loglog(f0, r0, 'r:')
    grid(True)
    xlabel('Frequency [Hz]')
    ylabel('Residual motion [m/rHz]')
    legend(('Open loop', 'Closed loop'))
    savefig('plots/loops_prcl_residual.eps')
    savefig('plots/loops_prcl_residual.png')

FF[dof['PRCL'], dof['PRCL'], :] = gain*D*M*F

####################################################################
######## MICH loop #################################################
####################################################################

# optical transfer function
G = opt[dof['MICH'], sig['SYM_2M1_Q'], :]
# mechanics
mech = TF([], [], [0.6], [500], 13e-6, 0)
M = mech.fresp(f)

filt = TF([3, 7], [0, 1], [0, 100, 200], [0.7, 1, 1], 1, 1)
F = filt.fresp(f)

delay = TF([],[],[],[],1,1)
delay.delay = 600e-6
D = delay.fresp(f)

gain = -1/ abs((G*M*D*F)[amin(find(f>20))])

if PLOT:
    figure()
    subplot(211)
    loglog(f, abs(gain*D*M*F*G))
    ylabel('Abs')
    title('MICH open loop transfer function')
    grid(True)
    axis(xmin=0.01, xmax=5000)
    subplot(212)
    semilogx(f, 180/pi* angle(gain*D*M*F*G))
    grid(True)
    ylabel('Phase')
    xlabel('Frequency [Hz]')
    axis(xmin=0.01, xmax=5000)
    savefig('plots/loops_mich_oltf.eps')
    savefig('plots/loops_mich_oltf.png')


if FREE:
    # simulate effect on free motion below 10 Hz
    olt = 5e3*f0**-2
    zc = z0 / olt
    r0 = rms(f0, z0)
    rc = rms(f0, zc)
    
    figure()
    loglog(f0, z0, 'r-')
    loglog(f0, zc, 'b-')
    loglog(f0, rc, 'b:')
    loglog(f0, r0, 'r:')
    grid(True)
    xlabel('Frequency [Hz]')
    ylabel('Residual motion [m/rHz]')
    legend(('Open loop', 'Closed loop'))
    savefig('plots/loops_mich_residual.eps')
    savefig('plots/loops_mich_residual.png')

FF[dof['MICH'], dof['MICH'], :] = gain*D*M*F

####################################################################
######## SREC loop #################################################
####################################################################

# optical transfer function
G = opt[dof['SREC'], sig['SYM_2M1_P'], :]
# mechanics
mech = TF([], [], [0.6], [500], 13e-6, 0)
M = mech.fresp(f)

filt = TF([3, 7], [0, 1], [0, 100, 200], [0.7, 1, 1], 1, 1)
F = filt.fresp(f)

delay = TF([],[],[],[],1,1)
delay.delay = 600e-6
D = delay.fresp(f)

gain = 1/ abs((G*M*D*F)[amin(find(f>20))])

if PLOT:
    figure()
    subplot(211)
    loglog(f, abs(gain*D*M*F*G))
    ylabel('Abs')
    title('SREC open loop transfer function')
    grid(True)
    axis(xmin=0.01, xmax=5000)
    subplot(212)
    semilogx(f, 180/pi* angle(gain*D*M*F*G))
    grid(True)
    ylabel('Phase')
    xlabel('Frequency [Hz]')
    axis(xmin=0.01, xmax=5000)
    savefig('plots/loops_srec_oltf.eps')
    savefig('plots/loops_srec_oltf.png')

if FREE:
    # simulate effect on free motion below 10 Hz
    olt = 5e3*f0**-2
    zc = z0 / olt
    r0 = rms(f0, z0)
    rc = rms(f0, zc)
    
    figure()
    loglog(f0, z0, 'r-')
    loglog(f0, zc, 'b-')
    loglog(f0, rc, 'b:')
    loglog(f0, r0, 'r:')
    grid(True)
    xlabel('Frequency [Hz]')
    ylabel('Residual motion [m/rHz]')
    legend(('Open loop', 'Closed loop'))
    savefig('plots/loops_srec_residual.eps')
    savefig('plots/loops_srec_residual.png')

FF[dof['SREC'], dof['SREC'], :] = gain*D*M*F

#######################################################
# Build optical matrix                                #
#######################################################

GG = zeros((ndof, ndof, nf), dtype=complex)

GG[dof['DARM'], dof['DARM'], :] = opt[dof['DARM'], sig['ASY_DC'], :]
GG[dof['DARM'], dof['CARM'], :] = opt[dof['CARM'], sig['ASY_DC'], :]
GG[dof['DARM'], dof['MICH'], :] = opt[dof['MICH'], sig['ASY_DC'], :]
GG[dof['DARM'], dof['PRCL'], :] = opt[dof['PRCL'], sig['ASY_DC'], :]
GG[dof['DARM'], dof['SREC'], :] = opt[dof['SREC'], sig['ASY_DC'], :]
                             
GG[dof['CARM'], dof['DARM'], :] = opt[dof['DARM'], sig['SYM_SB1_P'], :]
GG[dof['CARM'], dof['CARM'], :] = opt[dof['CARM'], sig['SYM_SB1_P'], :]
GG[dof['CARM'], dof['MICH'], :] = opt[dof['MICH'], sig['SYM_SB1_P'], :]
GG[dof['CARM'], dof['PRCL'], :] = opt[dof['PRCL'], sig['SYM_SB1_P'], :]
GG[dof['CARM'], dof['SREC'], :] = opt[dof['SREC'], sig['SYM_SB1_P'], :]
                             
GG[dof['PRCL'], dof['DARM'], :] = opt[dof['DARM'], sig['SYM_3P1_P'], :]
GG[dof['PRCL'], dof['CARM'], :] = opt[dof['CARM'], sig['SYM_3P1_P'], :]
GG[dof['PRCL'], dof['MICH'], :] = opt[dof['MICH'], sig['SYM_3P1_P'], :]
GG[dof['PRCL'], dof['PRCL'], :] = opt[dof['PRCL'], sig['SYM_3P1_P'], :]
GG[dof['PRCL'], dof['SREC'], :] = opt[dof['SREC'], sig['SYM_3P1_P'], :]
                             
GG[dof['MICH'], dof['DARM'], :] = opt[dof['DARM'], sig['SYM_2M1_Q'], :]
GG[dof['MICH'], dof['CARM'], :] = opt[dof['CARM'], sig['SYM_2M1_Q'], :]
GG[dof['MICH'], dof['MICH'], :] = opt[dof['MICH'], sig['SYM_2M1_Q'], :]
GG[dof['MICH'], dof['PRCL'], :] = opt[dof['PRCL'], sig['SYM_2M1_Q'], :]
GG[dof['MICH'], dof['SREC'], :] = opt[dof['SREC'], sig['SYM_2M1_Q'], :]
                             
GG[dof['SREC'], dof['DARM'], :] = opt[dof['DARM'], sig['SYM_2M1_P'], :]
GG[dof['SREC'], dof['CARM'], :] = opt[dof['CARM'], sig['SYM_2M1_P'], :]
GG[dof['SREC'], dof['MICH'], :] = opt[dof['MICH'], sig['SYM_2M1_P'], :]
GG[dof['SREC'], dof['PRCL'], :] = opt[dof['PRCL'], sig['SYM_2M1_P'], :]
GG[dof['SREC'], dof['SREC'], :] = opt[dof['SREC'], sig['SYM_2M1_P'], :]



#######################################################
# Compute closed loop transfer functions              #
#######################################################

CLTF = zeros((ndof, ndof, nf), dtype=complex)
for i in range(nf):
    CLTF[:,:,i] = inv(matrix(eye(5)) + matrix(GG[:,:,i]) * matrix(FF[:,:,i]))


######################################################
# Noise projections                                  #
######################################################

# build sensing noise vector (W)
noise = load('advirgo1_noise_dd_hp.txt')
NDARM = zeros((ndof, nf), dtype=double)
NDARM[dof['DARM'], :] = transpose(noise[:, sig['ASY_DC']]) 

NCARM = zeros((ndof, nf), dtype=double)
NCARM[dof['CARM'], :] = transpose(noise[:, sig['SYM_SB1_P']])

NMICH = zeros((ndof, nf), dtype=double)
NMICH[dof['MICH'], :] = transpose(noise[:, sig['SYM_2M1_Q']])

NPRCL = zeros((ndof, nf), dtype=double)
NPRCL[dof['PRCL'], :] = transpose(noise[:, sig['SYM_3P1_P']])

NSREC = zeros((ndof, nf), dtype=double)
NSREC[dof['SREC'], :] = transpose(noise[:, sig['SYM_2M1_P']])

# DARM error signal (total)
EDARM = zeros((ndof, nf), dtype=complex)
ECARM = zeros((ndof, nf), dtype=complex)
EPRCL = zeros((ndof, nf), dtype=complex)
EMICH = zeros((ndof, nf), dtype=complex)
ESREC = zeros((ndof, nf), dtype=complex)

for i in range(nf):
    EDARM[:,i:i+1] = matrix(CLTF[:,:,i]) * transpose(matrix(NDARM[:,i])) 
    ECARM[:,i:i+1] = matrix(CLTF[:,:,i]) * transpose(matrix(NCARM[:,i])) 
    EMICH[:,i:i+1] = matrix(CLTF[:,:,i]) * transpose(matrix(NMICH[:,i])) 
    EPRCL[:,i:i+1] = matrix(CLTF[:,:,i]) * transpose(matrix(NPRCL[:,i])) 
    ESREC[:,i:i+1] = matrix(CLTF[:,:,i]) * transpose(matrix(NSREC[:,i])) 

# calibrate in h
CAL = zeros((ndof, ndof, nf), dtype=complex)
for i in range(nf):
    CAL[:,:,i] = matrix(CLTF[:,:,i]) * matrix(GG[:,:,i])

calib = 1 / abs(CAL[dof['DARM'], dof['DARM'], :]) / 3000


HDARM = abs(EDARM[dof['DARM'],:]) * calib
HCARM = abs(ECARM[dof['DARM'],:]) * calib
HMICH = abs(EMICH[dof['DARM'],:]) * calib
HPRCL = abs(EPRCL[dof['DARM'],:]) * calib
HSREC = abs(ESREC[dof['DARM'],:]) * calib


# load also sensitivity
h = load('/home/vajente/virgo/AdV/bench_nsns.txt')


figure()
loglog(h[:,0], h[:,1], 'r-', linewidth=2)
loglog(f, abs(HDARM), 'r-')
loglog(f, abs(HCARM), 'g-')
loglog(f, abs(HPRCL), 'b-')
loglog(f, abs(HMICH), 'c-')
loglog(f, abs(HSREC), 'k-')
legend(('Design','DARM','CARM','PRCL','MICH','SREC'), 'upper right')
grid(True)
xlabel('Frequency [Hz]')
ylabel('Strain sensitivity [1/rHz]')
axis(xmin=10, xmax=5000, ymin=1e-30, ymax=1e-20)
savefig('plots/loop_noises_dd.eps')
savefig('plots/loop_noises_dd.png')


# Auxiliary dofs

figure()
loglog(f, abs(abs(EDARM[dof['CARM'],:])), 'r-')
loglog(f, abs(abs(ECARM[dof['CARM'],:])), 'g-')
loglog(f, abs(abs(EPRCL[dof['CARM'],:])), 'b-')
loglog(f, abs(abs(EMICH[dof['CARM'],:])), 'c-')
loglog(f, abs(abs(ESREC[dof['CARM'],:])), 'k-')
legend(('DARM','CARM','PRCL','MICH','SREC'), 'lower left')
grid(True)
xlabel('Frequency [Hz]')
ylabel('Error signal [W]')
axis(xmin=5, xmax=5000)
title('CARM error signal')
savefig('plots/carm_noises.eps')
savefig('plots/carm_noises.png')

figure()
loglog(f, abs(abs(EDARM[dof['PRCL'],:])), 'r-')
loglog(f, abs(abs(ECARM[dof['PRCL'],:])), 'g-')
loglog(f, abs(abs(EPRCL[dof['PRCL'],:])), 'b-')
loglog(f, abs(abs(EMICH[dof['PRCL'],:])), 'c-')
loglog(f, abs(abs(ESREC[dof['PRCL'],:])), 'k-')
legend(('DARM','CARM','PRCL','MICH','SREC'), 'lower left')
grid(True)
xlabel('Frequency [Hz]')
ylabel('Error signal [W]')
axis(xmin=5, xmax=5000)
title('PRCL error signal')
savefig('plots/prcl_noises.eps')
savefig('plots/prcl_noises.png')

figure()
loglog(f, abs(abs(EDARM[dof['MICH'],:])), 'r-')
loglog(f, abs(abs(ECARM[dof['MICH'],:])), 'g-')
loglog(f, abs(abs(EPRCL[dof['MICH'],:])), 'b-')
loglog(f, abs(abs(EMICH[dof['MICH'],:])), 'c-')
loglog(f, abs(abs(ESREC[dof['MICH'],:])), 'k-')
legend(('DARM','CARM','PRCL','MICH','SREC'), 'lower left')
grid(True)
xlabel('Frequency [Hz]')
ylabel('Error signal [W]')
axis(xmin=5, xmax=5000)
title('MICH error signal')
savefig('plots/mich_noises.eps')
savefig('plots/mich_noises.png')

figure()
loglog(f, abs(abs(EDARM[dof['SREC'],:])), 'r-')
loglog(f, abs(abs(ECARM[dof['SREC'],:])), 'g-')
loglog(f, abs(abs(EPRCL[dof['SREC'],:])), 'b-')
loglog(f, abs(abs(EMICH[dof['SREC'],:])), 'c-')
loglog(f, abs(abs(ESREC[dof['SREC'],:])), 'k-')
legend(('DARM','CARM','PRCL','MICH','SREC'), 'lower left')
grid(True)
xlabel('Frequency [Hz]')
ylabel('Error signal [W]')
axis(xmin=5, xmax=5000)
title('SREC error signal')
savefig('plots/srec_noises.eps')
savefig('plots/srec_noises.png')