import pandas as pd
import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

from IPython import embed


def main():
    df = pd.read_csv("../../derived_data/whitenoise_trials.csv", index_col=0, sep=";")
    pos_df = pd.read_csv("../../derived_data/receptivefield_positions.csv", sep=";", index_col=0)
    baseline_df = pd.read_csv("../../derived_data/figure2_baseline_properties.csv", sep=";", index_col=0)

    datasets = pos_df.dataset_id.unique()
    positions = []
    delays = []
    phases = []
    for dataset in datasets:
        pos = pos_df.receptor_pos_relative[pos_df.dataset_id == dataset].values[0]
        trials = df[df.dataset_id == dataset]
        if len(trials) < 1:
            continue
        delay = np.mean(trials.delay)
        phase = baseline_df[baseline_df.dataset_id == dataset].phase_shifted.values[0]
        positions.append(pos)
        delays.append(delay)
        phases.append(phase)

    delays = np.array(delays)
    positions = np.array(positions)
    phases = np.array(phases)
    slope, intercept, *params = stats.linregress(positions, delays * 1000)
    xRange = np.arange(0., 1.1, 0.1)
    yfit = intercept + slope * xRange

    fig = plt.figure(figsize=(6, 3), constrained_layout=True)
    ax = fig.add_subplot(121)

    ax.scatter(positions, delays * 1000, label="data")
    ax.plot(xRange, yfit, color="tab:orange", label=f"r:{params[0]:.2f}, p:{params[1]:.3f}", ls="-")

    ax.set_xlabel("receptor position [rel.]")
    ax.set_xlim([0, 1.0])
    ax.set_xticks(np.arange(0, 1.1, 0.2))
    ax.set_xticks(np.arange(0, 1.1, 0.1), minor=True)

    ax.set_ylabel("delay [ms]")
    ax.set_yticks(np.arange(0.0, 6.1, 1.0))
    ax.set_yticks(np.arange(0.0, 6.1, 0.5), minor=True)

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.legend(fontsize=7)

    pslope, pintercept, *pparams = stats.linregress(phases, delays * 1000)
    xRange = np.arange(0., 4 * np.pi + 0.1, 0.1)
    pyfit = pintercept + pslope * xRange

    ax = fig.add_subplot(122)
    ax.scatter(phases, delays * 1000, label="data")
    ax.plot(xRange, pyfit, color="tab:orange", label=f"r:{pparams[0]:.2f}, p:{pparams[1]:.3f}", ls="-")
    ax.set_xlabel("phase [rad]")
    ax.set_xlim([0, 4 * np.pi])
    ax.set_xticks(np.arange(0, 4 * np.pi + 0.1, np.pi))
    ax.set_xticks(np.arange(0, 4 * np.pi + 0.1, 0.5 * np.pi), minor=True)
    ax.set_xticklabels([f"{i}$\pi$" for i in range(5)])
    ax.set_ylabel("delay [ms]")
    ax.set_yticks(np.arange(0.0, 6.1, 1.0))
    ax.set_yticks(np.arange(0.0, 6.1, 0.5), minor=True)

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.legend(fontsize=7)

    fig.savefig("whitenoise_delays.pdf")


if __name__ == "__main__":
    main()