from nnfabrik import builder
import os
import pickle
import argparse
import yaml
from tqdm import tqdm
import numpy as np
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

from power_spectrum_analysis import get_power_spectrum


parser = argparse.ArgumentParser()
parser.add_argument('--cell', default=0, type=int)
parser.add_argument('--retina_index', default='01', type=str)
parser.add_argument('--cell_stimulus', default='nm', type=str)


if __name__ == '__main__':
    args = parser.parse_args()
    basepath = '.'
    config_file = 'config_fixed_all'
    with open(
        f"{basepath}/data/marmoset_data/responses/{config_file}.yaml", "rb"
    ) as config_file:
        config_dict = yaml.unsafe_load(config_file)
    img_dir = os.path.join(basepath, config_dict["image_path"])
    neuronal_data_path = os.path.join(basepath, config_dict["response_path"])
    batch_size = 1
    subsample = 1
    num_of_trials = 20
    cell = args.cell
    num_of_channels = 30
    retina_index = args.retina_index
    device = 'cpu'
    stimulus  = 'nm'
    time_chunk = 4
    dataset_fn = 'datasets.frame_movie_loader'
    dataloader_config = dict(config=config_dict,
            basepath=basepath,
            img_dir_name=img_dir,
            neuronal_data_dir=neuronal_data_path,
            all_image_path=img_dir,
            batch_size=batch_size,
            seed=None,
            train_frac=0.8,
            subsample=subsample,
            crop=0,
            num_of_trials_to_use=num_of_trials,
            num_of_frames=num_of_channels,
            cell_index=None,
            retina_index=retina_index,
            device=device,
            time_chunk_size=time_chunk,
            num_of_layers=1,
            cell_indices_out_of_range=True,
            oracle_correlation_threshold=None,
            normalize_responses=False,
            full_img_h=800,
            full_img_w=1000,
            padding=0,
            img_h=800,
            img_w=1000,
            retina_specific_crops=False,
            stimulus_seed=2022,
            hard_coded=([0, 1, 2, 3, 4, 5, 6,], [7, 8, 9]))
    dataloader = builder.get_data(dataset_fn, dataloader_config)
    path = f'{basepath}/analyses/adaptation_paper/filters/'

    with open(f'{path}/cs_sc_{args.cell_stimulus}_r{int(retina_index)}_v2', 'rb') as file:
        all_cell_filters = pickle.load(file)
    filter_name = f'{args.cell_stimulus}_cell_{cell}'

    avg_filter = all_cell_filters[cell][0]
    avg_filter = np.repeat(np.repeat(avg_filter, 4, axis=0), 4, axis=1)
    print(avg_filter.shape)
    avg_filter = torch.tensor(avg_filter).transpose(-1, 0).unsqueeze(0).unsqueeze(0).to('cuda')

    frames = []

    limit = 2000
    for i, (images, responses) in enumerate(
            tqdm(
                dataloader["train"][dataloader_config["retina_index"]]
            , total=limit)
    ):
        if i == limit:
            break

        output = torch.nn.functional.conv3d(images.double().to('cuda'), avg_filter)

        frames.append(output[0, 0, 0, :600, :600].detach().cpu().numpy())
        frames.append(output[0, 0, 1, :600, :600].detach().cpu().numpy())
        frames.append(output[0, 0, 2, :600, :600].detach().cpu().numpy())
        frames.append(output[0, 0, 3, :600, :600].detach().cpu().numpy())

    frame_width = frames[0].shape[0]
    frame_height = frames[0].shape[1]

    kfreq_x = np.fft.fftshift(np.fft.fftfreq(frame_width)) * frame_width
    kfreq_y = np.fft.fftshift(np.fft.fftfreq(frame_height)) * frame_height
    kfreq_2d = np.meshgrid(kfreq_x, kfreq_y)

    k_norm = np.sqrt(kfreq_2d[0] ** 2 + kfreq_2d[1] ** 2)
    kbins = np.arange(0.5, max(frame_height, frame_width) // 2 + 1) - 0.052
    kvals = 0.5 * (kbins[1:] + kbins[:-1])

    spectra = []

    for frame in frames:
        spectra.append(get_power_spectrum(frame, k_norm, kbins))

    spectra = np.array(spectra)
    mean_spectrum = spectra.mean(axis=0)
    std_spectrum = spectra.std(axis=0)

    output_folder = f'{basepath}/analyses/model_paper/spectra/retina_{retina_index}/stim_{stimulus}_{filter_name}/'

    Path(output_folder).mkdir(exist_ok=True, parents=True)
    print('saving to', output_folder)
    np.save(f'{output_folder}/spectra.npy', spectra)
    np.save(f'{output_folder}/filter.npy', avg_filter.detach().cpu().numpy())
    print('spectra saved')
    with sns.plotting_context("talk"):
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.loglog(kvals, mean_spectrum, c="tab:blue")

        ax.fill_between(kvals, mean_spectrum - std_spectrum, mean_spectrum + std_spectrum, alpha=0.25, color="tab:blue")
        ax.grid(which="both", axis="x")
        ax.grid(which="major", axis="y")
        ax.set_xlabel("$k$")
        ax.set_ylabel("norm. power")
        plt.savefig(f'{output_folder}/rags.png', dpi=300)
        plt.show()

    plt.imshow(frames[0], cmap='gray', vmin=-1, vmax=1)
    plt.savefig(f'{output_folder}/cropped_example_frame.png', dpi=300)


