import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.animation import FuncAnimation
from matplotlib import rcParams
from IPython.display import HTML

from matplotlib.ticker import EngFormatter

plt.style.use('seaborn-poster')
rcParams['font.size'] = 18

#plt.ion()

fT = 10e6
Av0 = 100e3
fH0 = fT / Av0

f = np.logspace(1, 8, 1000)

def Amag(f, beta_inv):
    A = Av0 / (1 + 1j*f/fH0)
    Av = A / (1 + A/beta_inv)
    return abs(Av)

def nstages(G):
    fig, ax = plt.subplots(figsize=(10,7))
    ax.set(xlim=(min(f)/2, max(f)*2),
           ylim=(0.1, Av0),
           xlabel='Hz',
           ylabel='V / V',
           title='Frequency response of N stages for\nfixed system gain = %.0f' % G)
    
    ax.loglog(f, Amag(f, 1e6), '-k')
    ax.hlines(1, min(f), max(f), linewidth=1)

    ax.grid(True, linewidth=0.8)

    formatter1 = EngFormatter(places=0)
    ax.xaxis.set_major_formatter(formatter1)
    ax.yaxis.set_major_formatter(formatter1)

    # first line
    line = ax.plot(f, Amag(f, G), color='b', lw=2)[0]
    N_text = ax.text(2*fT/G, 2*G, 'N = 1', size='xx-large')
    g_text = ax.text(0.01*fT/G, 0.7*G, 'g = %i'%G, size='large', va='top')

    def amp(frame):
        n = frame + 1
        g = G**(1.0/n)
        a = Amag(f, g)**n
        line.set_ydata(a)
        N_text.set_text('N = %i' % n)
        g_text.set_text('g = %.1f' % g)
        return (line,)

    anim = FuncAnimation(fig, amp, interval=500, frames=7, repeat=True)
    return anim

HTML(nstages(10).to_html5_video())

HTML(nstages(100).to_html5_video())

HTML(nstages(1000).to_html5_video())

HTML(nstages(10000).to_html5_video())

#
# TODO: this is just for a fixed system gain A
# plot for a series of system gains


# -3dB frequency vs. N stages for fixed system gain A and individual stage GBW

def f3(N, A, fT):
    g = A**(1.0/N)
    return (fT/g) * np.sqrt(2**(1.0/N)-1)

n = np.arange(1, 16)

A = 100
fT = 10e6

bw = f3(n, A, fT)/(fT/A)


plt.close('all')
fig, ax = plt.subplots()

ax.plot(n, bw, 'o')
ax.set(xticks=n, xlim=(0, n[-1]+1))

yt = np.around(bw, 1)
ax.set(yticks=bw[(0, 1, 2, 3, 4, 5, 8),])

formatter1 = EngFormatter(places=1)
ax.yaxis.set_major_formatter(formatter1)

ax.plot((0, 9), (bw[8], bw[8]), 'k', linewidth=1)

ax.set(
    title='Bandwidth increase using N stages\nwith system gain = %.0f' % A,
    xlabel='N stages',
    ylabel='$BW_N \, / \, BW_1$')

ax.text(9, (f3(9, A, fT)*A/fT - 0.5), '9', va='top', ha='center')

plt.show()

def plot_closedloop():
    fig, ax = plt.subplots(figsize=(10,7))
    ax.set(xlim=(min(f)/2, max(f)*2),
           ylim=(0.1, 2*Av0),
           ylabel='gain (V/V)',
           xlabel='Hz',
           title=r'Closed-loop gain and $1/\beta$ gain prediction')
    
    # open-loop amplifier, set 1/beta ~0
    ax.loglog(f, Amag(f, 1e9), '-k')
    ax.hlines(1, min(f), max(f), linewidth=1)

    ax.grid(True, linewidth=0.8)

    formatter1 = EngFormatter(places=0)
    ax.xaxis.set_major_formatter(formatter1)
    ax.yaxis.set_major_formatter(formatter1)
    
    # line and text
    bi = 1e9
    line_beta = ax.plot((f[0], f[-1]), [bi, bi], 'r')
    line_a = ax.plot(f, Amag(f, bi), color='b', lw=2)[0]
    
    A_text = ax.text(0.5*fT/bi, 0.7*Amag(0.5*fT/bi, bi), '$|A|$', size='large', color='b', va='top')
    beta_text = ax.text(7*fT/bi, 0.7*bi, r'$1 / \beta', size='large', color='r', va='top')

    bis = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1e3, 2e3, 5e3, 10e3, 20e3, 50e3, 100e3]
    bis = [1, 10, 100, 1e3, 3e3, 10e3, 20e3, 50e3, 80e3, 100e3]
    #bis = np.logspace(0, np.log10(Av0), 20)
    bis = np.concatenate((bis, np.flip(bis,0)))

    def amp(frame):
        bi = bis[frame]
        a = Amag(f, bi)
        line_beta[0].set_ydata([bi, bi])
        line_a.set_ydata(a)
        
        A_text.set_position((0.5*fT/bi, 0.7*Amag(0.5*fT/bi, bi)))
        beta_text.set_position((7*fT/bi, 0.7*bi))
        beta_text.set_text(r'$1 / \beta$')
        
    
        return (line_a,)

    anim = FuncAnimation(fig, amp, interval=500, frames=len(bis), repeat=True, blit=True)
    return anim

HTML(plot_closedloop().to_html5_video())