import matplotlib.pyplot as plt
import numpy as np


def rect(x):
    r"""
    Rectangle function:
        rect(x) = {1, if |x|<= 0.5; 0, otherwise}
    """
    return np.where(np.abs(x) >= 0.5, 0, 1)


def chirp(t, T, Kr):
    """
    Create a chirp signal :
        S_{tx}(t) = rect(t/T) * exp(1j*pi*Kr*t^2)
    """
    return rect(t / T) * np.exp(1j * np.pi * Kr * t ** 2)


Ns = 1000
t = np.linspace(-1, 1, Ns)
yrect = rect(t)

plt.figure()
plt.plot(t, yrect, '-r', linewidth=2)
plt.axis([-1, 1, -0.1, 1.1])
plt.grid()
plt.xlabel('Time/s')
plt.ylabel('Amplitude')
plt.title('rect')
plt.show()

Kr1 = 5
Kr2 = 20
T1 = 2.0
T2 = 4.0
t = np.linspace(-3, 3, Ns)
ychirp1 = chirp(t, T1, Kr=Kr1)
ychirp2 = chirp(t, T2, Kr=Kr1)

ychirp3 = chirp(t, T1, Kr=Kr2)
ychirp4 = chirp(t, T2, Kr=Kr2)

plt.figure()
plt.subplot(221)
plt.grid()
plt.plot(t, ychirp1, '-r')
plt.xlabel('Time/s')
plt.ylabel('Amplitude')
plt.title("T=" + str(T1) + "s, K=" + str(Kr1) + "Hz/s")
plt.subplot(222)
plt.grid()
plt.plot(t, ychirp2, '-b')
plt.xlabel('Time/s')
plt.ylabel('Amplitude')
plt.title("T=" + str(T2) + "s, K=" + str(Kr1) + "Hz/s")
plt.subplot(223)
plt.grid()
plt.plot(t, ychirp3, '-r')
plt.xlabel('Time/s')
plt.ylabel('Amplitude')
plt.title("T=" + str(T1) + "s, K=" + str(Kr2) + "Hz/s")
plt.subplot(224)
plt.grid()
plt.plot(t, ychirp4, '-b')
plt.xlabel('Time/s')
plt.ylabel('Amplitude')
plt.title("T=" + str(T2) + "s, K=" + str(Kr2) + "Hz/s")
plt.tight_layout()
plt.show()

# ==================compute spectral

ychirp1_fft = np.fft.fft(ychirp1)
ychirp2_fft = np.fft.fft(ychirp2)
ychirp3_fft = np.fft.fft(ychirp3)
ychirp4_fft = np.fft.fft(ychirp4)

ychirp1_fft = np.fft.fftshift(ychirp1_fft)
ychirp2_fft = np.fft.fftshift(ychirp2_fft)
ychirp3_fft = np.fft.fftshift(ychirp3_fft)
ychirp4_fft = np.fft.fftshift(ychirp4_fft)

f = np.fft.fftfreq(t.shape[-1])
# print(f)
f = np.fft.fftshift(f)
# print(f)

plt.figure()
plt.subplot(221)
plt.grid()
plt.plot(f, np.abs(ychirp1_fft), '-r')
plt.xlabel('Frequency/Hz')
plt.ylabel('Amplitude')
plt.title("T=" + str(T1) + "s, K=" + str(Kr1) + "Hz/s")
plt.subplot(222)
plt.grid()
plt.plot(f, np.abs(ychirp2_fft), '-b')
plt.xlabel('Frequency/Hz')
plt.ylabel('Amplitude')
plt.title("T=" + str(T2) + "s, K=" + str(Kr1) + "Hz/s")
plt.subplot(223)
plt.grid()
plt.plot(f, np.abs(ychirp3_fft), '-r')
plt.xlabel('Frequency/Hz')
plt.ylabel('Amplitude')
plt.title("T=" + str(T1) + "s, K=" + str(Kr2) + "Hz/s")
plt.subplot(224)
plt.grid()
plt.plot(f, np.abs(ychirp4_fft), '-b')
plt.xlabel('Frequency/Hz')
plt.ylabel('Amplitude')
plt.title("T=" + str(T2) + "s, K=" + str(Kr2) + "Hz/s")
plt.tight_layout()
plt.show()
