import numpy as np
import pickle
import torch
import yaml

from training.measures import correlation
from utils.global_functions import (get_exclude_cells_based_on_thresholds,
    home,
    global_config,
)


def calculate_retina_reliability(
    retina_index,
    control_stimuli_responses=None,
    response_dir=f"{home}/data/responses/",
    config_dict=None,
    fev_threshold=None,
):
    if config_dict is None:
        config_dict = global_config
    if control_stimuli_responses is None:
        control_stimuli_responses, _ = get_control_and_running_stimuli(
            retina_index, response_dir=response_dir, config=config_dict
        )
    if fev_threshold is not None:
        get_exclude_cells_based_on_thresholds(
            retina_index,
            config=config_dict,
            explainable_variance_threshold=fev_threshold,
        )
    # control_stimuli_responses = np.delete(control_stimuli_responses, get_exclude_cells_based_on_thresholds(retina_index, explainable_variance_threshold=fev_threshold), axis=0)
    print("cells:", fev_threshold, control_stimuli_responses.shape[0])
    avg_response = np.mean(control_stimuli_responses, axis=2)
    correlations = []
    for trial in range(control_stimuli_responses.shape[2]):
        avg_response = np.mean(
            np.delete(control_stimuli_responses, trial, axis=2), axis=2
        )
        corr = correlation(
            torch.tensor(control_stimuli_responses[:, :, trial], dtype=torch.float),
            torch.tensor(avg_response, dtype=torch.float),
            eps=1e-8,
            dim=1,
        )
        correlations.append(corr[:, 0].numpy())
    correlations = np.asarray(correlations)
    means = np.mean(correlations, axis=0)
    print(list(means))
    print(np.mean(means))
    return means


def calculate_r_squared(control_stimuli):
    predictions = np.array(
        [
            control_stimuli[:, :, x]
            for x in range(control_stimuli.shape[2])
            if x < int(control_stimuli.shape[-1] / 2)
        ]
    )  # x % 2 == 0])
    responses = np.array(
        [
            control_stimuli[:, :, x]
            for x in range(control_stimuli.shape[2])
            if x >= int(control_stimuli.shape[-1] / 2)
        ]
    )  # % 2 == 1])
    responses = np.moveaxis(responses, 1, 0)
    responses = np.moveaxis(responses, 1, 2)
    predictions = np.moveaxis(predictions, 1, 0)
    predictions = np.moveaxis(predictions, 1, 2)

    prediction = np.mean(predictions, axis=2)
    response = np.mean(responses, axis=2)
    variances_explained, eves = [], []

    corr = correlation(
        torch.tensor(response, dtype=torch.float),
        torch.tensor(prediction, dtype=torch.float),
        eps=1e-12,
        dim=1,
    )
    print("means corr cells:", corr)
    print("mean corr:", torch.mean(corr))
    return corr[:, 0]


def explained_variance(responses, predictions):
    mse = np.mean(
        (
            np.moveaxis(np.tile(predictions, (responses.shape[2], 1, 1)), 0, 2)
            - responses
        )
        ** 2,
        axis=(1, 2),
    )
    responses = np.transpose(responses)
    total_variance, explainable_var = [], []

    for n in range(45):
        response = responses[:, :, n]
        obs_var = np.mean((np.var(response, axis=0, ddof=1)), axis=0)  # obs variance
        tot_var = np.var(response, axis=(0, 1), ddof=1)  # total variance
        total_variance.append(tot_var)
        explainable_var.append(tot_var - obs_var)  # explainable variance

    total_variance = np.array(total_variance)
    explainable_var = np.array(explainable_var)
    var_explained = total_variance - mse
    eve = var_explained / explainable_var
    return var_explained, eve


def calculate_repeated_variance(control_stimuli_responses):
    variance = np.mean(np.var(control_stimuli_responses, axis=(2), ddof=1), axis=1)
    return variance


def calculate_total_variance(control_stimuli_responses):
    total_variance = np.var(control_stimuli_responses, axis=(1, 2), ddof=1)
    return total_variance


def get_control_and_running_stimuli(
    retina_index, response_dir=f"{home}/data/responses/", config=global_config
):
    if isinstance(retina_index, int):
        file = config["files"][retina_index]
    else:
        file = [x for x in config['files'] if retina_index in x][0]
    with open(f"{response_dir}/{file}", "rb") as pkl:
        neural_data = pickle.load(pkl)
    running_stimuli = neural_data["train_responses"]
    control_stimuli = neural_data["test_responses"]
    return control_stimuli, running_stimuli


def calculate_explainable_variance_for_retina(
    control_stimuli, running_stimuli, retina_index, config=global_config
):
    if control_stimuli is None:
        control_stimuli, running_stimuli = get_control_and_running_stimuli(
            retina_index, config=config
        )
    control_variance = calculate_repeated_variance(control_stimuli)
    total_variance = calculate_total_variance(control_stimuli, running_stimuli)
    explainable_variance = (total_variance - control_variance) / total_variance
    return explainable_variance


if __name__ == "__main__":
    # Specify retina index for which to compute reliability
    retina_index = '04'
    stimulus = 'wn'

    if stimulus == 'wn':
        with open(
                f"data/marmoset_data/responses/config_05_all.yaml") as f:
            config = yaml.unsafe_load(f)
    elif stimulus == 'nm':
        with open(f"data/marmoset_data/responses/config_s4_fixed_ds4.yaml") as f:
            config = yaml.unsafe_load(f)
    else:
        raise ValueError

    response_dir = f"responses/"
    control_stimuli, running_stimuli = get_control_and_running_stimuli(
        retina_index, response_dir=response_dir, config=config
    )
    r_squared = calculate_r_squared(control_stimuli)
    explainable_variance = calculate_explainable_variance_for_retina(
        control_stimuli=control_stimuli,
        running_stimuli=running_stimuli,
        retina_index=retina_index,
    )
    print(np.nanmean(explainable_variance))

    reliabilities = calculate_retina_reliability(
        retina_index=retina_index,
        control_stimuli_responses=control_stimuli,
        response_dir=response_dir,
        fev_threshold=None,
    )

    print(np.mean(reliabilities))
