import argparse
import os
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import torch
from scipy.stats import binned_statistic
from tqdm import tqdm
import pickle

parser = argparse.ArgumentParser()
parser.add_argument("--animal", default="marmoset", type=str)
parser.add_argument("--dataset", default="20220412_SN_252MEA6010_le_s4", type=str)
parser.add_argument("--stimulus", default="frozennoise4x4bw1blink12750run2550freeze_85Hz_seed1", type=str)
parser.add_argument("--upsample", default=1, type=int)
parser.add_argument("--stimulus_smoothing", default=0.0, type=float)
parser.add_argument("--total_trials", default=2, type=int)
parser.add_argument("--square_crop", default=False, action="store_true", dest="square_crop")


def get_power_spectrum2(image, norm_ks, k_bins):
    fft_amp = np.abs(np.fft.fftshift(np.fft.fft2(image))) ** 2
    fft_amp /= fft_amp.size
    amps = binned_statistic(
        norm_ks.flatten(), fft_amp.flatten(),
        statistic="mean",
        bins=k_bins,
    )[0]
    amps_sum = amps.sum()
    if amps_sum != 0:
        amps /= amps_sum
    else:
        print("Warning: The sum of the amplitude bins is zero.")
        amps[:] = 0  # Or handle as needed
    return amps


def get_power_spectrum(image, norm_ks, k_bins):
    fft_amp = np.abs(np.fft.fftshift(np.fft.fft2(image))) ** 2
    amps = binned_statistic(
        norm_ks.flatten(), fft_amp.flatten(),
        statistic="mean",
        bins=k_bins,
    )[0]
    amps /= amps.sum()
    return amps


def get_nm_frames(data_dir):
    img_dir_name = f'{data_dir}/data/marmoset_data/stimuli'
    frame_file = "_img_"
    full_img_w = 1000
    full_img_h = 800
    images = os.listdir(img_dir_name)
    images = [frame for frame in images if frame_file in frame]
    all_frames = np.zeros((len(images), full_img_w, full_img_h), dtype=np.float16)
    i = 0
    for img_file in tqdm(sorted(images)):
        img = np.fromfile(os.path.join(img_dir_name, img_file), dtype=np.uint8)
        img = img.reshape((1000, 800))
        all_frames[i] = ((img / 255) - 0.5)/0.5
        i += 1
    all_frames = [x[200:-200, 100:-100].T for x in all_frames]
    return all_frames

def create_gaussian(size, sigma_x, sigma_y):
    y = np.arange(0, size, 1)
    x = np.arange(0, size, 1)
    x, y = np.meshgrid(x, y)
    mu = size//2, size//2
    # Compute the Gaussian function
    gaussian = np.exp(-(((x - mu[0]) ** 2) / (2 * sigma_x ** 2) + ((y - mu[1]) ** 2) / (2 * sigma_y ** 2)))
    # Normalize the filter
    gaussian /= np.sum(gaussian)
    return gaussian

def create_dog(size, sigma1_x, sigma1_y, sigma2_x, sigma2_y, amplitude2=0.5):
    y = np.arange(0, size)
    x = np.arange(0, size)
    x, y = np.meshgrid(x, y)
    mu = size//2, size//2

    gaussian1 = np.exp(-(((x - mu[0]) ** 2) / (2 * sigma1_x ** 2) + ((y - mu[1]) ** 2) / (2 * sigma1_y ** 2)))

    gaussian2 = np.exp(-(((x - mu[0]) ** 2) / (2 * sigma2_x ** 2) + ((y - mu[1]) ** 2) / (2 * sigma2_y ** 2)))
    gaussian2 = gaussian2/np.sum(gaussian2)
    gaussian1 = gaussian1/np.sum(gaussian1)
    dog_filter = gaussian1 - amplitude2*gaussian2

    dog_filter /= np.max(np.abs(dog_filter))

    return dog_filter


def get_wn_frames(data_dir):
    arr = np.load(f'{data_dir}/data/marmoset_data/non_repeating_stimuli_flipped/trial_000/all_images.npy')
    frames = [arr[:, :, i] for i in range(arr.shape[-1])]
    frames = [np.kron(x, np.ones((4,4))) for x in frames]
    frames = [x[:, 100:-100].T for x in frames]
    return frames

def apply_gaussian(image, gaussian):
    gaussian_filtered = ft_filtering(image, gaussian)
    return gaussian_filtered

def apply_dog(image, dog):
    """Apply Gaussian and Difference of Gaussian filters to the image."""
    dog_filtered = ft_filtering(image, dog)
    return dog_filtered

def apply_filter(image, filter):
    filtered_img = torch.nn.functional.conv2d(torch.tensor(image, dtype=torch.float32, device='cuda').unsqueeze(0).unsqueeze(0),
                                              torch.tensor(filter, dtype=torch.float32,  device='cuda').unsqueeze(0).unsqueeze(0), padding=28)
    return filtered_img[0,0].detach().cpu().numpy().astype(np.float16)

def load_filters(filter_names, path, size=None, sigma_x=None, sigma_y=None, sigma2_x=None, sigma2_y=None, amplitude2=None):
    filters = []
    for name in filter_names:
        if name == 'None':
            filters.append('None')
        elif name == 'gauss':
            filter = create_gaussian(size, sigma_x, sigma_y)
            filters.append(filter)
        elif name == 'dog':
            filter = create_dog(size, sigma_x, sigma_y, sigma2_x, sigma2_y, amplitude2=amplitude2)
            filters.append(filter)
        else:
            f = torch.tensor(np.load(f'{path}/{name}.npy'))
            filters.append(f.numpy())
    return filters

def ft_filtering(image, filter):
    image = torch.from_numpy(image).to(torch.float32).unsqueeze(0).unsqueeze(0)
    filter = torch.from_numpy(filter).to(torch.float32).unsqueeze(0).unsqueeze(0)

    f_img = torch.fft.fft2(image)
    f_filter = torch.fft.fft2(filter, s = image.shape[-2:])
    f_result = f_img * f_filter
    f_result = torch.fft.ifft2(f_result).real

    return f_result[0,0].numpy().astype(np.float16)


def load_sc_filters(all_filters_file, cells):
    spat_filters = []
    y = np.linspace(0, np.pi, 15)
    w_s = np.sin(y)
    with open(all_filters_file, 'rb') as file:
        filters = pickle.load(file)
    for key, (spat, temp) in filters.items():
        if key in cells:
            spat_filters.append((spat.T*w_s).T*w_s)
    return spat_filters


parser = argparse.ArgumentParser(description="Pretty print json")
parser.add_argument("--cell_stimulus", default='nm', type=str)
parser.add_argument("--cell", default="0", type=str)
parser.add_argument('--retina_index', default='1', type=str)
parser.add_argument('--data_dir', default='.', type=str)
if __name__ == '__main__':
    args = parser.parse_args()
    cell_types = {'0': 'm_off',
                  '1': 'p_off',
                  '2': 'l_off',
                  '4': 'p_on'}

    cell_stimulus = args.cell_stimulus
    data_dir = args.data_dir
    cells = [int(args.cell)]
    retina = args.retina_index
    with open(f'{data_dir}/analyses/adaptation_paper/cell_type_dict.pkl', 'rb') as file:
        cell_type_dict = pickle.load(file)
    retina_cell_types = cell_type_dict[f'0{retina}']
    cell_type = None

    filter_names =  'wn_mean_m_off', 'nm_mean_m_off', 'wn_mean_p_off', 'nm_mean_p_off', 'wn_mean_p_off',  'nm_mean_l_off', 'wn_mean_l_off', 'nm_mean_p_on', 'wn_mean_p_on'
    for sigma, sigma2 in [(6.9, 7)]:
        for amplitude2 in [0.9996, 0.99965]:
                if sigma >= sigma2:
                    continue
                size = 60
                path = f'{data_dir}/analyses/model_paper/spectra/retina_1/'
                filter_names = ['dog'] # 'dog', 'gauss', 'None'
                filters = load_filters(filter_names, path=path, size=size,
                                       sigma_x=sigma, sigma_y=sigma, sigma2_y=sigma2, sigma2_x=sigma2, amplitude2=amplitude2)
                for stimulus in ['nm']:
                    filtering = 'DoG' # GS, F
                    for filter_name, f in zip(filter_names, filters):
                        print(filter_name,)

                        frame_width = 600
                        frame_height = 600

                        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.02
                        kvals = 0.5 * (kbins[1:] + kbins[:-1])

                        if stimulus == 'nm':
                            frames = get_nm_frames(data_dir)
                        elif stimulus == 'wn':
                            frames = get_wn_frames(data_dir)
                        else:
                            raise ValueError(f"Stimulus types only nm or wn not {stimulus}")

                        if filtering == 'GS':
                            frames = [apply_gaussian(x, f) for x in tqdm(frames)]
                        elif filtering == 'DoG':
                            frames = [apply_dog(x, f) for x in tqdm(frames)]
                        elif filtering == 'F':
                            frames = [apply_filter(x, f) for x in tqdm(frames)]

                        spectra = []

                        for frame in tqdm(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)

                        if filtering == 'DoG':
                            sigma_str = f'_s1_{sigma}_s2_{sigma2}_a_{amplitude2}'
                        elif filtering == 'GS':
                            sigma_str = f'_s1_{sigma}'
                        elif filtering == 'F':
                            sigma_str = filter_name
                        elif filtering == 'A':
                            sigma_str = f'A_{filter_name}'
                        else:
                            sigma_str = 'no'

                        output_folder = f'{data_dir}/analyses/model_paper/spectra/retina_{retina}/stim_{stimulus}_f_{filtering}_{sigma_str}/'
                        Path(output_folder).mkdir(exist_ok=True, parents=True)
                        print('saving to', output_folder)

                        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)
                        np.save(f'{output_folder}/spectra.npy', spectra)
                        np.save(f'{output_folder}/filter.npy', f)


