import streamlit as st
import numpy as np
import matplotlib.pyplot as plt


def main():
    st.set_page_config(layout="wide")
    st.title(
        "Procesamiento Digital de Señales - UNTREF \n ## Spectral Leakage en la DFT"
    )
    st.divider()
    signal_expression, signal_settings = st.columns(2)

    with signal_settings:
        st.markdown("### Parámetros de la señal")
        num_samples = st.slider(
            "Cantidad de Muestras", min_value=4, max_value=64, value=8, step=1
        )

        freq = st.slider(
            r"Frecuencia angular $\omega_0$",
            min_value=-np.pi,
            max_value=np.pi,
            value=0.0,
            step=np.pi / 64,
        )
    with signal_expression:
        freq_samples = int(np.round(64 * freq / np.pi))
        st.markdown("### Señal analizada:")
        with st.columns(3)[1]:
            st.markdown(f"# $$e^{{j\\frac{{{freq_samples}}}{{64}}\pi n}}$$")

    samples = np.arange(num_samples)
    signal = np.exp(1j * (freq) * samples)

    dtft_samples = 2048
    dtft_result = np.fft.fftshift(np.fft.fft(signal, dtft_samples)) / num_samples
    dtft_freqs = np.fft.fftshift(np.fft.fftfreq(dtft_samples))

    dft_result = np.fft.fftshift(np.fft.fft(signal)) / num_samples
    dft_freqs = np.fft.fftshift(np.fft.fftfreq(len(signal)))

    dft_freqs_rad = 2 * np.pi * dft_freqs
    dtft_freqs_rad = 2 * np.pi * dtft_freqs

    # Plot the signals
    fig = plt.figure(figsize=(10, 4))
    plt.axvline(np.pi, color="black", linestyle="--")
    plt.axvline(-np.pi, color="black", linestyle="--")
    plt.axvline(freq, color="black", alpha=0.5, linestyle=":", label=r"$\omega_0$")
    plt.plot(dtft_freqs_rad, np.abs(dtft_result), color="red", label="DTFT")
    plt.stem(dft_freqs_rad, np.abs(dft_result), markerfmt="C0o", label="DFT")


    plt.xlabel("Frecuencia angular [rad]")
    plt.ylabel("Magnitud de la transformada normalizada")
    plt.ylim(0, 1.1)
    plt.xlim(-np.pi - 0.1, np.pi + 0.1)
    plt.xticks(
        np.linspace(-np.pi, np.pi, 9),
        [
            r"$-\pi$",
            r"$-\frac{3}{4}\pi$",
            r"$-\frac{\pi}{2}$",
            r"$-\frac{\pi}{4}$",
            r"0",
            r"$\frac{\pi}{4}$",
            r"$\frac{\pi}{2}$",
            r"$\frac{3}{4}\pi$",
            r"$\pi$",
        ],
    )
    plt.grid()
    plt.legend(loc="upper right", framealpha=1, ncol=3, bbox_to_anchor=(0.65, 1.06))

    plt.tight_layout()

    st.pyplot(fig, use_container_width=True)


if __name__ == "__main__":
    main()