import math
import os
import pickle
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import gridspec
import torch
from tqdm import tqdm


def get_filter_sign(f):
    max_value = np.max(f)
    min_value = np.min(f)
    if np.abs(min_value) > np.abs(max_value):
        return -1
    else:
        return 1

def crop_cell_rf(whole_image, rf_size, rf_center, img_h=150, img_w=200):
    image = whole_image[
        :,
        (rf_center[0] - min(rf_size[0] // 2, rf_center[0])) : (
            rf_center[0]
            + min((rf_size[0] // 2) + (rf_size[0] % 2), img_h - rf_center[0])
        ),
        (rf_center[1] - min(rf_size[1] // 2, rf_center[1])) : (
            rf_center[1]
            + min((rf_size[1] // 2) + (rf_size[1] % 2), img_w - rf_center[1])
        ),
    ]

    left = rf_center[0] - math.floor(rf_size[0] / 2)
    right = img_h - (rf_center[0] + math.ceil(rf_size[0] / 2))
    top = rf_center[1] - math.floor(rf_size[1] / 2)
    bottom = img_w - (rf_center[1] + math.ceil(rf_size[1] / 2))

    new_img = np.zeros((whole_image.shape[0],) + rf_size)
    if left < 0:
        if top < 0:
            new_img[:, -1 * left :, -1 * top :] = image
        elif bottom < 0:
            new_img[:, -1 * left :, :bottom] = image
        else:
            new_img[:, -1 * left :, :] = image
    elif right < 0:
        if top < 0:
            new_img[:, :right, -1 * top :] = image
        elif bottom < 0:
            new_img[:, :right, :bottom] = image
        else:
            new_img[:, :right, :] = image
    elif top < 0:
        new_img[:, :, -1 * top :] = image
    elif bottom < 0:
        new_img[:, :, :bottom] = image

    else:
        new_img = image
    return new_img

def get_cell_sta(cell_file, data_dir, rf_size, h=150, w=200, flip_sta=False):
    cell_rf = np.load(f"{data_dir}/{cell_file}")
    if flip_sta:
        cell_rf = np.flip(cell_rf, axis=1).copy()
    temporal_variances = np.var(cell_rf, axis=0)
    max_coordinate = np.unravel_index(np.argmax(temporal_variances), shape=temporal_variances.shape)
    cell_rf = cell_rf[cell_rf.shape[0] - rf_size[0]:, :, :]
    images = crop_cell_rf(cell_rf, rf_size[1:], max_coordinate, img_w=w, img_h=h)
    if (images.shape[1] < rf_size[1]) or (images.shape[2] < rf_size[2]):
        raise ValueError(f'For cell with STA in file {cell_file} the STA center was either not found or is too close '
                         f'to the edge of the image for filter size {rf_size}')
    return images, max_coordinate


def separate_time_space_sta(sta):
    t, h, w = sta.shape
    temporal_variances = np.var(sta, axis=0)
    max_coordinate = np.unravel_index(np.argmax(temporal_variances), (h, w))
    time_pos = np.argmax(np.abs(sta[:, max_coordinate[0], max_coordinate[1]]))
    mis = sta[time_pos, max_coordinate[0], max_coordinate[1]]

    spat_filter = np.sign(mis) * sta[time_pos]
    spat_filter /= np.linalg.norm(spat_filter)

    temp_filter = sta[:, max_coordinate[0], max_coordinate[1]]
    temp_filter /= np.linalg.norm(temp_filter)
    return spat_filter, temp_filter


def get_cell_svd(sta):
    t, h, w = sta.shape
    sta = sta.reshape(sta.shape[0], -1)
    temp, s, spat = np.linalg.svd(sta)
    temp = temp[:, 0]
    spat = spat[0]
    spat = spat.reshape((h, w))
    return spat, temp


def create_cell_file_from_config_version(original_file, cell_index, retina_index:str):
    new_file = original_file.replace("01", retina_index)
    new_file = f"{new_file}_{cell_index}.npy"
    return new_file

def crop_around_receptive_field(max_coordinate, images, rf_size, h, w):
    assert max_coordinate is not None
    if len(images.shape) == 4:
        images = torch.permute(images, (-2, -1, 0, 1))
        images = images[
            max_coordinate[0]
            - min(rf_size[0] // 2, max_coordinate[0]) : max_coordinate[0]
            + min(rf_size[0] // 2 + rf_size[0] % 2, h - max_coordinate[0]),
            max_coordinate[1]
            - min(rf_size[1] // 2, max_coordinate[1]) : max_coordinate[1]
            + min(rf_size[1] // 2 + rf_size[1] % 2, w - max_coordinate[1]),
        ]
        images = torch.permute(images, (-2, -1, 0, 1))
    elif len(images.shape) == 5:
        images = images[
            :,
            :,
            :,
            max_coordinate[0]
            - min(rf_size[0] // 2, max_coordinate[0]) : max_coordinate[0]
            + min(rf_size[0] // 2 + rf_size[0] % 2, h - max_coordinate[0]),
            max_coordinate[1]
            - min(rf_size[1] // 2, max_coordinate[1]) : max_coordinate[1]
            + min(rf_size[1] // 2 + rf_size[1] % 2, w - max_coordinate[1]),
        ]
    return images

def calculate_initial_stas(responses, images, cell_index, time_bins, height, width):
    sta = np.zeros((time_bins, height, width))
    spike_count = 0
    cell_responses = responses[cell_index]

    for i, response in enumerate(cell_responses):
        if i > time_bins and response > 0:
            sta += response * images[i - time_bins : i]
            spike_count += response
    if spike_count > 0:
        sta = sta / spike_count
    return sta

def save_sta(
    response_file,
    image_files,
    image_dir,
    save_dir,
    save_file,
    time_bins,
    height,
    width,
):
    with open(response_file, "rb") as f:
        responses = pickle.load(f)
    responses = responses["train_responses"]
    responses = np.moveaxis(responses, 1, 2)

    print("loading images")
    all_stas = np.zeros((responses.shape[0],) + (time_bins, height, width))
    print("calculating stas")
    for i, image_file in enumerate(sorted(image_files)[:10]):
        print(image_file)
        image_responses = responses[:, :, i]
        print(f"responses_shape: {image_responses.shape}")
        image = np.load(f"{image_dir}/{image_file}/all_images.npy")
        image = np.moveaxis(image, -1, 0)
        print(f"image shape: {image.shape}")
        height, width = image.shape[1:]
        for cell in tqdm(range(responses.shape[0])):
            sta = calculate_initial_stas(
                image_responses,
                image,
                cell_index=cell,
                time_bins=time_bins,
                height=height,
                width=width,
            )
            all_stas[cell] += sta

    all_stas /= 25
    for cell in range(responses.shape[0]):
        np.save(f'{save_dir}/{save_file}_cell_{cell}', all_stas)

        show_sta(all_stas[cell], cell=cell)

    # all_stas /= len(image_files)



def show_sta(sta, cell, vmin=-1, vmax=1):
    fig = plt.figure(figsize=(20, 24))
    gs = gridspec.GridSpec(
        10,
        5,
        width_ratios=[1, 1, 1, 1, 1],
        wspace=0.05,
        hspace=0.05,
        top=0.95,
        bottom=0.05,
        left=0.17,
        right=0.845,
    )
    for i in range(sta.shape[0]):
        x = math.floor(i / 5)
        y = i % 5
        ax = plt.subplot(gs[x, y])
        ax.set_axis_off()
        ax.imshow(sta[i], cmap="gray", vmin=vmin, vmax=vmax)
    ax = plt.subplot(gs[0, 2])
    ax.set_title(f"Cell {cell}")
    fig.subplots_adjust(hspace=0.05, wspace=0.05)
    # plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    home = './'
    time_bins = 40
    image_files = os.listdir(
        f"{home}/data/marmoset_data/non_repeating_stimuli_frozennoise4x4bw1blink12750run2550freeze_85Hz_seed1/"
    )
    image_files = [x for x in image_files if x.startswith("trial")]
    save_sta(
        f"{home}/data/marmoset_data/responses/05_SS_frozennoise4x4bw1blink12750run2550freeze_85Hz_seed1.pkl",
        image_files=image_files,
        image_dir=f"{home}/data/marmoset_data/non_repeating_stimuli_frozennoise4x4bw1blink12750run2550freeze_85Hz_seed1/",
        save_dir=f"{home}/data/marmoset_data/stas/",
        save_file=f"cell_data_02_NC_stas",
        time_bins=time_bins,
        height=150,
        width=200,
    )


