import matplotlib.pyplot as plt
import numpy as np

def importData(filename):
    dataFile = open(filename)
    data = dataFile.read().splitlines()
    dataFile.close()

    f = []
    p = []
    m = []

    for line in data[16:]:
        line = line.split('    ')
        f += [float(line[0])]
        m += [float(line[1])]
        p += [float(line[2])]

    return f, m, p

def findUGF(m, p):
    ugfI = np.argmin(np.abs(np.array(m)-1))
    return f[ugfI], p[ugfI]

def findGM(m, p, maxI):
    pi = np.argmin(np.abs(np.array(p[:b])+180))
    return pi, 20*np.log10(m[pi])

# no daughter
b = 450
f, m, p = importData('TFAG4395A_15-09-2016_140247.txt')
ugf, phaseMargin = findUGF(m, p)
pi, gm = findGM(m, p, b)

# with daughter
f2, m2, p2 = importData('TFAG4395A_15-09-2016_140553.txt')
ugf2, phaseMargin2 = findUGF(m2, p2)
pi2, gm2 = findGM(m2, p2, b)

plt.figure(figsize=(16,9))
plt.subplot(211)
plt.title('2016-09-15 H1 IMC PDH Open Loop Gain TF')
plt.loglog(f, m, lw=2, label='w/o 200 kHz Pole')
plt.loglog(f2, m2, color='g', lw=2, label='w/ 200 kHz Pole')
plt.axvline(x=ugf, lw=2, linestyle='--', color='b')
plt.axvline(x=ugf2, lw=2, linestyle='--', color='g')
plt.axvline(x=f[pi], lw=1, linestyle='-', color='b')
plt.axvline(x=f2[pi2], lw=1, linestyle='-', color='g')
plt.legend(loc='best')
plt.ylim([1e-3, 10])
plt.grid(which='both')
plt.ylabel('Magnitude [V/V]')
plt.annotate(r'UGF (w/o) = %.1f kHz' % (ugf/1e3), xy=(0.01, 0.28), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.annotate(r'UGF (w/) = %.1f kHz' % (ugf2/1e3), xy=(0.01, 0.21), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.annotate(r'GM (w/o) = %.1f dB' % gm, xy=(0.01, 0.14), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.annotate(r'GM (w/) = %.1f dB' % gm2, xy=(0.01, 0.07), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.subplot(212)
plt.axvline(x=ugf, lw=2, linestyle='--', color='b')
plt.axvline(x=ugf2, lw=2, linestyle='--', color='g')
plt.semilogx(f, p, lw=2)
plt.semilogx(f2, p2, color='g', lw=2)
plt.grid(which='both')
plt.ylim([-180, 180])
plt.yticks(np.arange(-180, 181, 45))
plt.xlabel('frequency [Hz]')
plt.ylabel('Phase [deg]')
plt.annotate(r'PM (w/o) = %.1f$^\circ$' % (180+phaseMargin), xy=(0.01, 0.95), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.annotate(r'PM (w/) = %.1f$^\circ$' % (180+phaseMargin2), xy=(0.01, 0.88), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')

plt.savefig('20160915_h1imcpdh_olp.png')

g1 = np.array(m) * np.exp(2*np.pi*np.array(p)*1j/360)
g2 = np.array(m2) * np.exp(2*np.pi*np.array(p2)*1j/360)

maxPeak = np.max(np.abs(1/(g1+1)))
maxPeak2 = np.max(np.abs(1/(g2+1)))

plt.figure(figsize=(16,9))
plt.subplot(211)
plt.title('2016-09-15 H1 IMC PDH Loop Suppression (1 / 1+G)')
plt.loglog(f, np.abs(1/(g1+1)), lw=2, label='w/o 200 kHz Pole')
plt.loglog(f2, np.abs(1/(g2+1)), color='g', lw=2, label='w/ 200 kHz Pole')
plt.legend(loc='best')
plt.grid(which='both')
plt.ylim([1e-1, 10])
plt.ylabel('Magnitude [V/V]')
plt.annotate(r'Max gain peaking (w/o) = %.1f' % maxPeak, xy=(0.01, 0.95), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.annotate(r'Max gain peaking (w/) = %.1f' % maxPeak2, xy=(0.01, 0.88), xycoords='axes fraction', horizontalalignment='left', verticalalignment='top', fontsize='14')
plt.subplot(212)
plt.semilogx(f, np.angle(1/(g1+1), deg=True), lw=2)
plt.semilogx(f2, np.angle(1/(g2+1), deg=True), color='g', lw=2)
plt.grid(which='both')
plt.xlabel('frequency [Hz]')
plt.ylim([-180, 180])
plt.yticks(np.arange(-180, 181, 45))
plt.ylabel('Phase [deg]')

plt.savefig('20160915_h1imcpdh_loopsupp.png')

plt.figure(figsize=(16,9))
plt.subplot(211)
plt.title('2016-09-15 H1 IMC PDH Closed Loop Gain (G / 1+G)')
plt.loglog(f, np.abs(g1/(g1+1)), lw=2, label='w/o 200 kHz Pole')
plt.loglog(f2, np.abs(g2/(g2+1)), color='g', lw=2, label='w/ 200 kHz Pole')
plt.legend(loc='best')
plt.grid(which='both')
plt.ylabel('Magnitude [V/V]')
plt.ylim([1e-3, 10])
plt.subplot(212)
plt.semilogx(f, np.angle(g1/(g1+1), deg=True), lw=2)
plt.semilogx(f2, np.angle(g2/(g2+1), deg=True), color='g', lw=2)
plt.grid(which='both')
plt.xlabel('frequency [Hz]')
plt.ylim([-180, 180])
plt.yticks(np.arange(-180, 181, 45))
plt.ylabel('Phase [deg]')

plt.savefig('20160915_h1imcpdh_clg.png')

plt.show()