import numpy as np
import matplotlib.pyplot as plt
from thunderfish.powerspectrum import spectrogram, decibel

from modules.filehandling import LoadData
from modules.datahandling import instantaneous_frequency
from modules.filters import bandpass_filter
from modules.plotstyle import PlotStyle

ps = PlotStyle()


def main():
    # Load data
    datapath = "../data/2022-06-02-10_00/"
    data = LoadData(datapath)

    # good chirp times for data: 2022-06-02-10_00
    window_start_seconds = 3 * 60 * 60 + 6 * 60 + 43.5 + 9 + 6.20
    window_start_index = window_start_seconds * data.raw_rate
    window_duration_seconds = 0.4
    window_duration_index = window_duration_seconds * data.raw_rate

    timescaler = 1000

    raw = data.raw[
        window_start_index : window_start_index + window_duration_index, 10
    ]

    fig, (ax1, ax2) = plt.subplots(
        1, 2, figsize=(21 * ps.cm, 8 * ps.cm), sharex=True, sharey=True
    )

    # plot instantaneous frequency
    filtered1 = bandpass_filter(
        signal=raw, lowf=750, highf=1200, samplerate=data.raw_rate
    )
    filtered2 = bandpass_filter(
        signal=raw, lowf=550, highf=700, samplerate=data.raw_rate
    )

    freqtime1, freq1 = instantaneous_frequency(
        filtered1, data.raw_rate, smoothing_window=3
    )
    freqtime2, freq2 = instantaneous_frequency(
        filtered2, data.raw_rate, smoothing_window=3
    )

    ax1.plot(freqtime1 * timescaler, freq1, color=ps.g, lw=2, label="Fish 1")
    ax1.plot(freqtime2 * timescaler, freq2, color=ps.gray, lw=2, label="Fish 2")
    # ax.legend(bbox_to_anchor=(1.04, 1), borderaxespad=0)
    # # ps.hide_xax(ax1)

    # plot fine spectrogram
    spec_power, spec_freqs, spec_times = spectrogram(
        raw,
        ratetime=data.raw_rate,
        freq_resolution=150,
        overlap_frac=0.2,
    )

    ylims = [300, 1300]
    fmask = np.zeros(spec_freqs.shape, dtype=bool)
    fmask[(spec_freqs > ylims[0]) & (spec_freqs < ylims[1])] = True

    ax1.imshow(
        decibel(spec_power[fmask, :]),
        extent=[
            spec_times[0] * timescaler,
            spec_times[-1] * timescaler,
            spec_freqs[fmask][0],
            spec_freqs[fmask][-1],
        ],
        aspect="auto",
        origin="lower",
        interpolation="gaussian",
        alpha=1,
        # vmin=-100,
        # vmax=-80,
    )

    # # plot coarse spectrogram
    spec_power, spec_freqs, spec_times = spectrogram(
        raw,
        ratetime=data.raw_rate,
        freq_resolution=15,
        overlap_frac=0.3,
    )
    fmask = np.zeros(spec_freqs.shape, dtype=bool)
    fmask[(spec_freqs > ylims[0]) & (spec_freqs < ylims[1])] = True
    ax2.imshow(
        decibel(spec_power[fmask, :]),
        extent=[
            spec_times[0] * timescaler,
            spec_times[-1] * timescaler,
            spec_freqs[fmask][0],
            spec_freqs[fmask][-1],
        ],
        aspect="auto",
        origin="lower",
        interpolation="gaussian",
        alpha=1,
    )
    # ps.hide_xax(ax3)
    ax2.plot(freqtime1 * timescaler, freq1, color=ps.g, lw=2, label="_")
    ax2.plot(freqtime2 * timescaler, freq2, color=ps.gray, lw=2, label="_")

    ax2.set_xlim(75, 200)
    ax1.set_ylim(400, 1200)

    fig.supxlabel("Time [ms]", fontsize=14)
    fig.supylabel("Frequency [Hz]", fontsize=14)

    handles, labels = ax1.get_legend_handles_labels()
    ax2.legend(
        handles,
        labels,
        bbox_to_anchor=(1.04, 1),
        loc="upper left",
        ncol=1,
    )

    ps.letter_subplots(xoffset=[-0.27, -0.1], yoffset=1.05)

    plt.subplots_adjust(
        left=0.12, right=0.85, top=0.89, bottom=0.18, hspace=0.35
    )

    plt.savefig("../poster/figs/introplot.pdf")


if __name__ == "__main__":
    main()