import numpy as np
import matplotlib.pyplot as plt

#%%###########################################################################
def complex_to_polar(c):
    return np.angle(c), np.abs(c)

def cchirp(fs=int(5e7), tau=5e-5, bw=int(1e7), get_t=False):
    alpha = bw / tau
    t = np.arange(0, tau, 1 / fs)
    y = np.exp(1j * np.pi * alpha * t ** 2)
    return y, t if get_t else y

def plot_v2(x, unwrap=True):
    coef = np.fft.fft(x)
    theta, r = complex_to_polar(coef)

    plt.plot(r)
    plt.ylim(0, max(r) * 1.05)
    plt.show()
    theta = np.unwrap(theta) if unwrap else theta
    plt.plot(theta)
    plt.show()

def wplot(x):
    plt.plot(x)
    _scale_plot(plt.gcf(), plt.gca())

def mag_zoom(r, h=10, x0=0, show=True):
    N = len(r)
    N2 = N // 2
    a, b = N2 + x0 - h, N2 + x0 + h + 1
    k = np.linspace(-N2, N2 - 1, N)[a:b]
    plt.plot(k, r[a:b])
    _scale_plot(plt.gcf(), plt.gca())
    if show:
        plt.show()

def _scale_plot(fig, ax):
    xmin, xmax = ax.get_xlim()
    rng = xmax - xmin
    ax.set_xlim(xmin + .025 * rng, xmax - .025 * rng)
    fig.set_size_inches(15, 7)

#%%###########################################################################
tau, fs, bw = 200000, 50, 5
x, t = cchirp(fs=int(fs), bw=int(bw), tau=tau, get_t=1)
plot_v2(x.real, unwrap=1)