import numpy as np
from matplotlib import pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

import kid_readout.analysis.noise_archive

if not globals().has_key('arch'):
    arch = kid_readout.analysis.noise_archive.load_archive('/home/data/archive/StarCryo_3x3_0813f5_HPD_Dark_1.npy')
    arch['round_temp'] = np.round(arch.end_temp*1000/10)*10

arch = arch[(arch.power_dbm > -114) & (arch.power_dbm <= -98)]
arch = arch[np.abs(arch.end_temp -0.235) > 0.01]
arch = arch[(arch.Q_i > 0) & (arch.Q_i < 4e5)]
idxs = np.unique(arch['ridx'])

pset = arch[abs(arch.power_dbm +101.6) < 1]
icolors = plt.cm.spectral(np.linspace(0.1,0.9,8))
fig = plt.figure()
ax = fig.add_subplot(111)
for idx in range(8):
    this = pset[pset.ridx == idx]
    this = this.sort(['end_temp'])
    ax.semilogy(this.end_temp*1000,this.noise_250_hz,'.-',lw=2,color=icolors[idx],label=('res #%d' % idx))
    ax.semilogy(this.end_temp*1000,this.noise_30_khz,'--',color=icolors[idx])
ax.legend(prop=dict(size='small'))
ax.grid()
ax.set_xlim(0,700)
ax.set_ylim(1e-19,1e-15)
ax.set_title(('StarCryo_3x3_0813f5_HPD_Dark_1 @ %.1f dBm' %(this.power_dbm.iat[0])),size='medium')
ax.set_xlabel('Temperature (mK)')
ax.set_ylabel('Noise Power (1/Hz)')
ax.text(0.95,0.1,"Solid = Device\nDashed = Amplifier",fontdict=dict(size='small'), ha='right',va='bottom',
            transform=ax.transAxes)
fig.savefig('/home/data/archive/StarCryo_3x3_0813f5_HPD_Dark_1_alll.pdf',bbox_inches='tight')
    

for idx in range(8):
    pdf = PdfPages('/home/data/archive/%s_Resonator_%d.pdf' % ('StarCryo_3x3_0813f5_HPD_Dark_1',idx))
    res = arch[arch.ridx==idx]
    f0max = res.f_0.max()
    powers = np.unique(res.power_dbm).tolist()
    powers.sort()
    rtemps = np.unique(res.round_temp).tolist()
    rtemps.sort()
    pcolors = plt.cm.spectral(np.linspace(0.1,0.9,len(powers)))
    tcolors = plt.cm.spectral(np.linspace(0.1,0.9,len(rtemps)))
    
    title = 'StarCryo_3x3_0813f5_HPD_Dark_1\nResonator %d ~%.6f MHz' % (idx,f0max)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    for power in powers:
        c = pcolors[powers.index(power)]
        this = res[res.power_dbm == power]
        this = this.sort(['end_temp'])
        f0 = this.f_0
        frac_f0 = f0/f0.max() - 1
        ax.plot(this.end_temp*1000,frac_f0,'.-',color=c,label=('%.1f dBm' % power))
    ax.grid()
    ax.set_xlim(0,700)
    ax.legend(prop=dict(size='small'))
    ax.set_title(title,size='small')
    ax.set_xlabel('Temperature (mK)')
    ax.set_ylabel('Fractional frequency shift')
    pdf.savefig(fig,bbox_inches='tight')
    ax.set_ylim(-1e-4,0)
    pdf.savefig(fig,bbox_inches='tight')
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for power in powers:
        c = pcolors[powers.index(power)]
        this = res[res.power_dbm == power]
        this = this.sort(['end_temp'])
        ax.plot(this.end_temp*1000,this.Q_i,'.-',color=c,label=('%.1f dBm' % power))
    ax.grid()
    ax.legend(prop=dict(size='small'))
    ax.set_xlim(0,700)
    ax.set_title(title,size='small')
    ax.set_xlabel('Temperature (mK)')
    ax.set_ylabel('$Q_i$')
    pdf.savefig(fig,bbox_inches='tight')
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for power in powers:
        c = pcolors[powers.index(power)]
        this = res[res.power_dbm == power]
        this = this.sort(['end_temp'])
        ax.plot(this.end_temp*1000,this.Q_e_real,'.-',color=c,label=('%.1f dBm' % power))
    ax.grid()
    ax.legend(prop=dict(size='small'))
    ax.set_xlim(0,700)
    ax.set_title(title,size='small')
    ax.set_xlabel('Temperature (mK)')
    ax.set_ylabel('$Q_c$')
    pdf.savefig(fig,bbox_inches='tight')
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for power in powers:
        c = pcolors[powers.index(power)]
        this = res[res.power_dbm == power]
        this = this.sort(['end_temp'])
        ax.plot(this.end_temp*1000,this.Q,'.-',color=c,label=('%.1f dBm' % power))
    ax.grid()
    ax.legend(prop=dict(size='small'))