import os.path
import pickle
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["WORLD_SIZE"] = "1"
import numpy as np
from pathlib import Path


from datasets.stas import create_cell_file_from_config_version, get_cell_sta
from models.ln_model import FactorizedModel, RankModel
from training.measures import variance_of_predictions, correlation
from training.trainers import model_step
from utils.global_functions import (
    get_cell_numbers_after_crop,
    get_cell_names,
    home,
    model_seed,
    global_config,
)
from nnfabrik import builder
import torch
import yaml
from tqdm import tqdm


home_dir = "/user/vystrcilova/"


def get_ln_model_cell_performance_dict(
    file_name,
    retina_index,
    metric="correlations",
    exclude_cells=False,
    base_path=f"{home}/models/ln_models/",
    correlation_threshold=None,
    explainable_variance_threshold=None,
    config=global_config,
):
    cell_stats = {}
    retina_index_str = f"0{retina_index+1}"
    num_of_cells = (
        get_cell_numbers_after_crop(
            retina_index,
            correlation_threshold=correlation_threshold
            if correlation_threshold is not None
            else 0,
            explained_variance_threshold=explainable_variance_threshold
            if explainable_variance_threshold is not None
            else 0,
        )
        if exclude_cells
        else global_config["cell_numbers"][retina_index_str]
    )
    cell_names_list = (
        get_cell_names(
            retina_index,
            correlation_threshold=correlation_threshold
            if correlation_threshold is not None
            else 0,
            explained_variance_threshold=explainable_variance_threshold
            if explainable_variance_threshold is not None
            else 0,
        )
        if exclude_cells
        else [x for x in range(num_of_cells)]
    )
    for cell in range(num_of_cells):
        if os.path.isfile(
            f"{base_path}/salamander/retina{retina_index+1}/cell_{cell_names_list[cell]}/{file_name}/stats/{metric}.npy"
        ):
            data = np.load(
                f"{base_path}/salamander/retina{retina_index+1}/cell_{cell_names_list[cell]}/{file_name}/stats/{metric}.npy"
            )
            cell_stats[cell_names_list[cell]] = np.max(data)
    return cell_stats

def load_ln_model(model_dir, base_path, data_type, retina_index, cell_number, filename, seed,
                  device):
    if not os.path.isfile(
            f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{filename}/weights/seed_{seed}/best_model.m"
    ):
        print(f'skipping cell {cell_number} because there is no model attached.')
        return None
    model = torch.load(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{filename}/weights/seed_{seed}/best_model.m",
        map_location=torch.device(device),
    )
    return model

def build_ln_model_and_dataloader(dataset_fn, dataloader_config, factorized,
                                  model_dict, model_config, freeze=True, device='cpu', rank_model=False):
    dataloaders = builder.get_data(dataset_fn, dataloader_config)
    if rank_model:
        model = RankModel(**model_config)
    else:
        model = FactorizedModel(**model_config)

    if 'nonlin.a' in model_dict.keys():
        model_dict['nonlinearity.a'] = model_dict['nonlin.a']
        del(model_dict['nonlin.a'])
        model_dict['nonlinearity.b'] = model_dict['nonlin.b']
        del(model_dict['nonlin.b'])
        model_dict['nonlinearity.w'] = model_dict['nonlin.w']
        del(model_dict['nonlin.w'])

    model.load_state_dict(model_dict)
    model = model.to(device)
    if freeze:
        model.eval()
        for name, parameter in model.named_parameters():
            parameter.requires_grad = False

    return dataloaders, model

def get_model_and_dataloader_for_ln(
    config,
    filename,
    cell_number,
    retina_index,
    device="cpu",
    seed=None,
    data_type="salamander",
    base_path=f"models/ln_models/",
    model_dir=home,
    data_dir = None,
    factorized=False,
    freeze=True,
    rank_model=False
):
    if '/' in filename:
        config_filename = filename.split('/')[0]
    else:
        config_filename = filename
    if seed is None:
        seed = model_seed

    model = load_ln_model(model_dir, base_path, data_type, retina_index,
                          cell_number, filename, seed, device)
    if model is None:
        return None, None, None

    with open(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml",
        "r",
    ) as config_file:
        config = yaml.unsafe_load(config_file)
    if data_dir is None:
        data_dir = home
    model_dict = model["model"]
    dataset_fn = "datasets.white_noise_loader"
    dataloader_config = config["dataloader_config"]
    if "num_of_frames" in config["model_config"].keys() and factorized:
        config["model_config"]["num_of_channels"] = config["model_config"][
            "num_of_frames"
        ]
        del config["model_config"]["num_of_frames"]
    dataloader_config[
        "train_image_path"
    ] = f'{data_dir}/{dataloader_config["config"]["training_img_dir"]}'
    dataloader_config[
        "test_image_path"
    ] = f'{data_dir}/{dataloader_config["config"]["test_img_dir"]}'
    dataloader_config[
        "neuronal_data_dir"
    ] = f'{data_dir}/{dataloader_config["config"]["neuronal_data_path"]}'
    dataloaders, model = build_ln_model_and_dataloader(dataset_fn, dataloader_config, factorized, model_dict,
                                                       model_config=config['model_config'], freeze=freeze,
                                                       device=device, rank_model=rank_model)

    return dataloaders, model, config


def get_model_and_dataloader_for_ln_on_nm(
    filename,
    cell_number,
    retina_index,
    device="cpu",
    seed=None,
    data_type="salamander",
    model_dir=home,
    base_path=f"/models/ln_models/",
    data_dir=None,
    factorized=False,
    freeze=True,
    test=False,
        rank_model=False
):
    if '/' in filename:
        config_filename = filename.split('/')[0]
    else:
        config_filename = filename
    if seed is None:
        seed = model_seed
    model = load_ln_model(model_dir, base_path, data_type, retina_index,
                          cell_number, filename, seed, device)
    if model is None:
        return None, None, None
    with open(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml",
        "r",
    ) as config_file:
        config = yaml.unsafe_load(config_file)
    model_dict = model["model"]
    dataset_fn = "datasets.frame_movie_loader"
    dataloader_config = config["dataloader_config"]

    if data_dir is not None:
        dataloader_config['basepath'] = data_dir
        dataloader_config['all_image_path'] = f'{data_dir}/{dataloader_config["config"]["image_path"]}'
        dataloader_config['img_dir_name'] = f'{data_dir}/{dataloader_config["config"]["image_path"]}'
        dataloader_config['neuronal_data_dir'] = f'{data_dir}/{dataloader_config["config"]["response_path"]}'

    dataloaders, model = build_ln_model_and_dataloader(dataset_fn, dataloader_config, factorized, model_dict,
                                                       model_config=config['model_config'], freeze=freeze,
                                                       device=device, rank_model=rank_model)
    return dataloaders, model, config


def get_ln_model_dataloader_wn_for_nm(filename,
    cell_number,
    config_dict,
    retina_index,
    device="cpu",
    seed=8,
    data_type="salamander",
    model_dir=home,
    base_path=f"/models/ln_models/",
    data_dir=None,
    factorized=False,
    freeze=True,
    test=False,
    fancy_nonlin=True,
    stimulus_seed=0,
    flip_imgs=False,
    rank_model=False

):
    if '/' in filename:
        config_filename = filename.split('/')[0]
    else:
        config_filename = filename
    if os.path.isfile(f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml"):
        with open(f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml", "r") as config_file:
            model_related_config = yaml.unsafe_load(config_file)
    else:
        return None, None, None
    if fancy_nonlin:
        model_related_config['fancy_nonlin'] = True

    model = load_ln_model(model_dir, base_path, data_type, retina_index,
                          cell_number, filename, seed, device)
    if model is None:
        return None, None, None

    dataset_fn = "datasets.frame_movie_loader"
    model_related_config["base_path"] = model_dir
    dataloader_config_from_model = model_related_config["dataloader_config"]

    new_dataloader_config = {}
    new_dataloader_config["config"] = config_dict
    new_dataloader_config["basepath"] = f"{data_dir}"
    new_dataloader_config['stimulus_seed'] = stimulus_seed
    new_dataloader_config['full_img_w'] = 350
    new_dataloader_config['full_img_h'] = 300
    new_dataloader_config['padding'] = 50
    new_dataloader_config[
        "img_dir_name"
    ] = f'{data_dir}/data/{data_type}_data/stimuli_padded_4'
    new_dataloader_config[
        "all_image_path"
    ] = f'{data_dir}/data/{data_type}_data/stimuli_padded_4/'
    new_dataloader_config["batch_size"] = 16
    new_dataloader_config["time_chunk_size"] = 100

    if test:
        new_dataloader_config["time_chunk_size"] = 50
        new_dataloader_config["batch_size"] = 16

    print(f"data_dir: {data_dir}")
    new_dataloader_config[
        "neuronal_data_dir"
    ] = f"{data_dir}/data/{data_type}_data/responses/"
    if flip_imgs:
        new_dataloader_config['flip_imgs'] = True
    print('using 5 trials to have equal number of WN and NM data')
    new_dataloader_config['num_of_trials_to_use'] = 5
    new_dataloader_config['hard_coded'] = [7, 11, 6, 1, 2, 8, 3, 0, 12, 13], [4, 14, 5, 15]

    for key, value in dataloader_config_from_model.items():
        if key not in [
            "train_image_path",
            "test_image_path",
            "overlapping",
            "conv3d",
            'movie_like',
            'use_cache',
            'config',
            'neuronal_data_dir',
            'files',
            'batch_size',
            'time_chunk_size',
            'flip_sta',
            'num_of_trials_to_use',
            'extra_frame',
            'base_path',
            'hard_coded'
        ]:
            new_dataloader_config[key] = value
    _, model = build_ln_model_and_dataloader(dataset_fn='datasets.white_noise_loader',
                                             dataloader_config=dataloader_config_from_model,
                                             model_dict=model['model'], factorized=factorized, freeze=freeze,
                                             model_config=model_related_config['model_config'],
                                             device=device, rank_model=rank_model)
    dataloaders = builder.get_data(dataset_fn, new_dataloader_config)

    return dataloaders, model, model_related_config

def get_ln_model_dataloader_nm_for_wn(filename,
    cell_number,
    config_dict,
    retina_index,
    device="cpu",
    seed=None,
    data_type="salamander",
    model_dir=home,
    base_path=f"/models/ln_models/",
    data_dir=None,
    factorized=False,
    freeze=True,
    test=False,
    flip_sta=False,
    rank_model=False
                                      ):

    if '/' in filename:
        config_filename = filename.split('/')[0]
    else:
        config_filename = filename
    if os.path.isfile(f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml"):
        with open(f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell_number}/{config_filename}/config/config.yaml", "r") as config_file:
            model_related_config = yaml.unsafe_load(config_file)
    else:
        return None, None, None

    model_related_config["base_path"] = data_dir

    dataloader_config_from_model = model_related_config["dataloader_config"]
    new_dataloader_config = {
        'train_image_path': f"{data_dir}/data/{data_type}_data/non_repeating_stimuli_flipped/" if retina_index != '04' else f'{data_dir}/data/{data_type}_data/non_repeating_stimuli_flipped_dataset4/',
        "test_image_path": f'{data_dir}/data/{data_type}_data/repeating_stimuli_flipped/' if retina_index != '04' else f'{data_dir}/data/{data_type}_data/repeating_stimuli_flipped_dataset4/',
        "neuronal_data_dir": f"{data_dir}/data/{data_type}_data/responses/", "conv3d": True, "movie_like": False,
        "flip": False,
        "time_chunk_size": 70,
        'batch_size': 16,
        "num_of_trials_to_use": 21
    }
    print(f"data_dir: {data_dir}")
    new_dataloader_config["config"] = config_dict
    new_dataloader_config['crop'] = config_dict['big_crops']['01']

    new_dataloader_config["config"]["big_crops"]["01"] = config_dict['big_crops']['01']
    new_dataloader_config["crop"] = config_dict['big_crops']['01']
    new_dataloader_config['hard_coded'] = '0 6 3 7 4 1 8 2;5 9'

    for key, value in dataloader_config_from_model.items():

        if key not in [
            "all_image_path",
            "frame_file",
            "img_dir_name",
            "full_img_w",
            "full_img_h",
            "padding",
            "stimulus_seed",
            "basepath",
            "hard_coded",
            "conv3d",
            "movie_like",
            "config",
            'neuronal_data_dir',
            'time_chunk_size',
            'batch_size',
            'num_of_trials_to_use',
            'all_image_path'
        ]:
            new_dataloader_config[key] = value

    dataloaders, model, _ = get_model_and_dataloader_for_ln_on_nm(filename, cell_number, retina_index,
                                                               device, seed, data_type, model_dir,
                                                               base_path, data_dir, factorized, freeze, test,
                                                                  rank_model=rank_model)
    dataset_fn = "datasets.white_noise_loader"

    dataloaders = builder.get_data(dataset_fn, new_dataloader_config)

    return dataloaders, model, model_related_config


def get_model_and_dataloader_for_ln_based_on_setting(config, setting, filename, cell_number, retina_index,
                                                     device, seed, data_type, base_path, model_dir, data_dir,
                                                     factorized, freeze, fixation_file=None,
                                                     stimulus_seed=None, test=True, fancy_nonlin=True,
                                                     flip_imgs=False, rank_model=False):
    if setting == 'wn':
        dataloaders, model, config = get_model_and_dataloader_for_ln(None,
                                                            filename, cell_number,
                                                            retina_index, device,
                                                            seed, data_type,
                                                            base_path,
                                                            model_dir, data_dir, factorized,
                                                            freeze,rank_model=rank_model)
    elif setting == 'nm':
        dataloaders, model, config  = get_model_and_dataloader_for_ln_on_nm(filename, cell_number, retina_index, device,
                                                                  seed, data_type, model_dir, base_path, data_dir,
                                                                  factorized, freeze, rank_model=rank_model)
    elif setting == 'wn_for_nm':
        dataloaders, model, config = get_ln_model_dataloader_wn_for_nm(filename, cell_number,
                                                               fixation_file=fixation_file,
                                                               config_dict=config, retina_index=retina_index,
                                                               device=device,
                                                               data_type=data_type,
                                                               model_dir=model_dir,
                                                               base_path=base_path,
                                                               data_dir=data_dir, factorized=factorized,
                                                               freeze=freeze,
                                                               test=test,fancy_nonlin=fancy_nonlin,
                                                               stimulus_seed=stimulus_seed,
                                                               flip_imgs=flip_imgs, rank_model=rank_model)

    elif setting == 'nm_for_wn':
        dataloaders, model, config = get_ln_model_dataloader_nm_for_wn(filename=filename,
                                                               cell_number=cell_number, config_dict=config,
                                                               retina_index=retina_index, device=device,
                                                               seed=seed, data_type=data_type, model_dir=model_dir,
                                                               base_path=base_path,
                                                               data_dir=data_dir, factorized=factorized,

                                                               freeze=freeze,
                                                               test=test, rank_model=rank_model)
    else:
        raise ValueError(f'Setting {setting} for model and dataloader not defined')
    return dataloaders, model.double() if model is not None else None, config


def get_ln_model_performance(
    file_name,
    retina_index,
    setting='nm',
    exclude_cells=False,
    correlation_threshold=None,
    explainable_variance_threshold=0,
    model_dir=None,
    device="cuda",
    rf_size=60,
    img_h=150,
    img_w=200,
    save_file=None,
    start_index=0,
    data_type="salamander",
    base_path=f"{home}/models/ln_models/",
    factorized=False,
    config=global_config,
    nm=False,
    performance='validation',
    evaluate_again=False,
    freeze=True,
    cell_names_list=None,
        fixation_file=None,
        stimulus_seed=None,
        fancy_nonlin=True,
        flip_imgs=False,
        rank_model=False

):
    cell_stats = {}
    cell_means = []
    performance_str = '' if performance == 'validation' else f'_{performance}'
    recorded_cell_mean = None
    num_of_cells = (
        get_cell_numbers_after_crop(
            retina_index,
            config=config,
            correlation_threshold=correlation_threshold
            if correlation_threshold is not None
            else 0,
            explained_variance_threshold=explainable_variance_threshold
            if explainable_variance_threshold is not None
            else 0,
        )
        if exclude_cells
        else get_cell_numbers_after_crop(retina_index, config=config)
    )
    # num_of_cells -=1
    if cell_names_list is None:
        cell_names_list = (
        get_cell_names(
            retina_index,
            config=config,
            correlation_threshold=correlation_threshold
            if correlation_threshold is not None
            else 0,
            explained_variance_threshold=explainable_variance_threshold
            if explainable_variance_threshold is not None
            else 0,
        )
        if exclude_cells
        else [x for x in range(num_of_cells)]
    )
    if model_dir is None:
        model_dir = home

    for _, cell in enumerate(cell_names_list):

        if performance == 'validation':
            if os.path.isfile(
                    f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/correlations.npy"
            ) :
                recorded_correlations = np.load(
                    f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/correlations.npy"
                )
                recorded_cell_mean = np.max(recorded_correlations)

            else:
                recorded_cell_mean = 'N/A'
            print(f'cell: {cell}, recorded validation perf: {recorded_cell_mean}')
        elif performance == 'test':
            setting_str = f'{setting}_' if setting in ['wn_for_nm', 'nm_for_wn'] else ''
            if os.path.isfile(
                    f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}correlations_test.npy"
            ):
                recorded_correlations = np.load(
                    f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}correlations_test.npy"
                )
                recorded_cell_mean = np.max(recorded_correlations)
            else:
                recorded_cell_mean = 'N/A'
            print(f'cell: {cell}, recorded test perf: {recorded_cell_mean}')
        if not evaluate_again and os.path.isfile(
            f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/correlations{performance_str}.npy"
        ):
            correlations = np.load(
                f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/correlations{performance_str}.npy"
            )
            if len(correlations) > 1:
                cell_stats[cell] = np.max(correlations)
            else:
                cell_stats[cell] = correlations[0]
            cell_means.append(cell_stats[cell])
            print(f'cell: {cell}, evaluated validation perf: {cell_stats[cell]}')

            continue
        elif not evaluate_again and os.path.isfile(
            f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_correlation.npy"
        ):
            correlations = np.load(
                f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_correlation.npy"
            )
            cell_stats[cell] = correlations.item()
            cell_means.append(correlations.item())
            print(f'cell: {cell}, evaluated validation perf: {correlations.item()}')
            continue
        else:
            cell_stats[cell] = 'N/A'
        dataloaders, model, _ = get_model_and_dataloader_for_ln_based_on_setting(config,
                                                                              setting,
                                                                              file_name, cell,
                                                                              retina_index, device, seed=None,
                                                                              data_type=data_type,
                                                                              base_path=base_path,
                                                                              model_dir=model_dir,
                                                                              data_dir=model_dir, factorized=factorized,
                                                                              freeze=freeze, fixation_file=fixation_file,
                                                                              stimulus_seed=stimulus_seed,
                                                                              test=performance=='test',
                                                                              fancy_nonlin=fancy_nonlin,
                                                                              flip_imgs=flip_imgs,
                                                                              rank_model=rank_model
                                                                              )
        if model is None:
            print(f'skipping cell {cell}')
            continue
        test_correlations = evaluate_performance(setting, model, dataloaders, cell, nm, img_h, img_w, device, model_dir,
                                                 retina_index, base_path, start_index)
        if test_correlations is None:
            print(f'skipping cell {cell}')
            continue
        cell_stats[cell] = test_correlations.item()
        cell_means.append(test_correlations.item())
        print(f"cell {cell}:", test_correlations.item())
        if performance == 'validation':
            print('recorded perf from training:', recorded_cell_mean)
    print(f"{performance} correlation: ", np.mean(cell_means))
    print(cell_stats)
    return cell_stats, cell_means

def evaluate_performance(setting, model, dataloaders,
                         cell, nm,
                         img_h, img_w, device,
                         model_dir, retina_index,
                         base_path,
                         start_index=0,
                         ):
    if model is None:
        return None
    model.double()
    model.to(device)
    model.eval()
    if setting in ['wn_for_nm', 'nm', 'nm_for_wn']:
        nm = True

    first_session_ID = list((dataloaders["train"].keys()))[0]
    print(first_session_ID)
    a_dataloader = dataloaders["train"][first_session_ID]
    img_h = int(
        (img_h - a_dataloader.dataset.crop[0] - a_dataloader.dataset.crop[1]) / 1
    )
    img_w = int(
        (img_w - a_dataloader.dataset.crop[2] - a_dataloader.dataset.crop[3]) / 1
    )
    if rf_size is not None:
        cell_file = config["sta_file"]
        cell_file = create_cell_file_from_config_version(
            original_file=cell_file, cell_index=cell, retina_index=retina_index
        )
        sta_dir = config["stas_path"]
        receptive_field, max_coordinate = get_cell_sta(
            cell_file=cell_file,
            data_dir=f'{model_dir}/{sta_dir}',
            h=img_h,
            w=img_w,
            rf_size=(a_dataloader.dataset.num_of_frames, rf_size, rf_size),
            # TODO: has t
            flip_sta=nm
        )
        size = receptive_field[0].shape

    test_variances = []
    all_outputs = []
    all_responses = []
    for images, responses in tqdm(
            dataloaders[performance][retina_index]
    ):
        images = images.double().to(device)
        responses = responses.transpose(1, 2).to(device).contiguous()
        responses = torch.flatten(responses, start_dim=0, end_dim=1)
        all_responses.append(responses.detach().cpu())
        output = model_step(
            images=images,
            model=model,
            max_coordinate=max_coordinate,
            rf_size=size,
            h=img_h,
            w=img_w,
        )
        all_outputs.append(output.detach().cpu())
        prediction_variance = variance_of_predictions(output)
        test_variances.append(float(prediction_variance.item()))

    setting_str = '' if setting in ['nm', 'wn'] else f'{setting}_'
    all_outputs = torch.cat(all_outputs)
    all_responses = torch.cat(all_responses)


    Path(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/visualizations/cell_responses/"
    ).mkdir(exist_ok=True, parents=True)

    test_correlations = correlation(
        all_outputs[start_index:], all_responses[start_index:], 1e-12
    )
    np.save(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_correlation.npy",
        test_correlations.item(),
    )
    np.save(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_responses.npy",
        all_responses.detach().cpu().numpy(),
    )
    np.save(
        f"{model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_outputs.npy",
        all_outputs.detach().cpu().numpy(),
    )
    print(
        f'saved to: {model_dir}/{base_path}/{data_type}/retina{retina_index}/cell_{cell}/{file_name}/stats/seed_{model_seed}/{setting_str}{performance}_correlation.npy')

    return test_correlations




if __name__ == "__main__":
    model_dir = '.'
    ln_rank_1_filename = 'lr_0.005_whole_rf_15_ch_30_l1_0.0_l2_0.0_g_0.0_bs_32_tr_250_s_1_c_0_n_0_fn_1_do_n_1_flipped_gf_0_dog_0'
    ln_rank_2_filename = 'cs_lr_0.005_whole_rf_15_ch_30_l1_0.0_l2_0.0_g_0.0_bs_32_tr_250_s_1_c_0_n_0_fn_1_do_n_1_flipped_gf_0_dog_0'
    result_path = f'{model_dir}/analyses/adaptation_paper/ln_model_performances'
    setting = 'wn_for_nm'
    if setting == 'nm':
        with open(
                    f"{home}/data/marmoset_data/responses/config_s4_fixed_all.yaml", "rb"
            ) as config_file:
                config = yaml.unsafe_load(config_file)
        basepath = 'models/ln_models/ln_models_factorized_4_ete_nm/'
        data_type = 'marmoset'
        fixation_file = 'data/marmoset_data/fixations/complete_fixations_seed_2022and2023_s4_fixed.txt'
        performance = 'test'
        rf_size = 15

        all_test_performances = []
        all_cell_test_performances = {}
        for retina in ['01', '04']:
            file_name = ln_rank_1_filename
            retina_perfs_test = []
            cell_stats, cell_means = get_ln_model_performance(
                setting=setting,
                retina_index=retina,
                exclude_cells=True,
                model_dir=model_dir,
                base_path=basepath,
                file_name=file_name,
                rf_size=rf_size,
                factorized=True,
                data_type=data_type,
                config=config,
                nm=True,
                explainable_variance_threshold=0.15,
                evaluate_again=True,
                performance=performance,
                device='cuda',
                stimulus_seed=None,
                fixation_file=None,
                rank_model= True if file_name.startswith('cs') else False

            )
            all_test_performances = all_test_performances + cell_means
            all_cell_test_performances[retina] = cell_stats
        with open(f'{result_path}/cs_{setting}_test_correlations_r1_r4.pkl', 'wb') as f:
            pickle.dump(all_cell_test_performances, f)

    elif setting == 'wn':
        with open(
        f"{model_dir}/data/marmoset_data/responses/config_05_all.yaml", "rb"
    ) as config_file:
            config = yaml.unsafe_load(config_file)
        basepath = 'models/ln_models/ln_models_factorized__ete'
        data_type = 'marmoset'
        fixation_file = 'data/marmoset_data/fixations/complete_fixations_seed_2022and2023_s4_fixed.txt'
        performance = 'test'
        rf_size = 15

        all_test_performances = []
        all_cell_test_performances = {}
        for retina in ['01', '04']:
            file_name = ln_rank_1_filename
            retina_perfs_test = []
            cell_stats, cell_means = get_ln_model_performance(
                setting=setting,
                retina_index=retina,
                exclude_cells=True,
                model_dir=model_dir,
                base_path=basepath,
                file_name=file_name,
                rf_size=rf_size,
                factorized=True,
                data_type=data_type,
                config=config,
                nm=False,
                explainable_variance_threshold=0.15,
                evaluate_again=True,
                performance=performance,
                device='cuda',
                stimulus_seed=None,
                fixation_file=None,
                rank_model=True if file_name.startswith('cs') else False
            )
            all_test_performances = all_test_performances + cell_means
            all_cell_test_performances[retina] = cell_stats
        with open(f'{result_path}/cs_wn_test_correlations_r1_r4.pkl', 'wb') as f:
            pickle.dump(all_cell_test_performances, f)

    elif setting == 'nm_for_wn':
        with open(
                    f"{home}/data/marmoset_data/responses/config_05_all.yaml", "rb"
            ) as config_file:
                config = yaml.unsafe_load(config_file)
        basepath = 'models/ln_models/ln_models_factorized_4_ete_nm/'
        data_type = 'marmoset'
        fixation_file = 'data/marmoset_data/fixations/complete_fixations_seed_2022and2023_s4_fixed.txt'
        performance = 'test'
        rf_size = 15

        all_test_performances = []
        all_cell_test_performances = {}
        for retina in ['01', '04']:
            file_name = ln_rank_1_filename
            retina_perfs_test = []
            cell_stats, cell_means = get_ln_model_performance(
                setting='nm_for_wn',
                retina_index=retina,
                exclude_cells=True,
                model_dir=model_dir,
                base_path=basepath,
                file_name=file_name,
                rf_size=rf_size,
                factorized=True,
                data_type=data_type,
                config=config,
                nm=False,
                explainable_variance_threshold=0.15,
                evaluate_again=False,
                performance=performance,
                device='cuda',
                stimulus_seed=None,
                fixation_file=None,
                rank_model=True if file_name.startswith('cs') else False
                # cell_names_list=[294, 67],

            )
            all_test_performances = all_test_performances + cell_means
            all_cell_test_performances[retina] = cell_stats
        with open(f'{result_path}/cs_nm_for_wn_test_correlations_r1_r4_finetuned_temp.pkl', 'wb') as f:
            pickle.dump(all_cell_test_performances, f)
        #

    elif setting == 'wn_for_nm':
        with open(
        f"{model_dir}/data/marmoset_data/responses/config_s4_fixed_all.yaml", "rb"
    ) as config_file:
            config = yaml.unsafe_load(config_file)
        basepath = 'models/ln_models/ln_models_factorized__ete'
        data_type = 'marmoset'
        fixation_file = 'data/marmoset_data/fixations/complete_fixations_seed_2022and2023_s4_fixed.txt'
        performance = 'test'
        rf_size = 15

        all_test_performances = []
        all_cell_test_performances = {}
        for retina in ['01', '04']:
            # file_name = 'lr_0.005_whole_rf_15_ch_30_l2_0.0_g_0.0_bs_32_tr_250_s_1_c_0_fn_1_do_n_1_dog_0/wn_for_nm/ff_0_pn_1'
            file_name = ln_rank_1_filename
            retina_perfs_test = []
            cell_stats, cell_means = get_ln_model_performance(
                setting=setting,
                retina_index=retina,
                exclude_cells=True,
                model_dir=model_dir,
                base_path=basepath,
                file_name=file_name,
                rf_size=rf_size,
                factorized=True,
                data_type=data_type,
                config=config,
                nm=True,
                explainable_variance_threshold=0.15,
                evaluate_again=False,
                performance=performance,
                device='cuda',
                stimulus_seed=None,
                fixation_file=None,
                rank_model=True if file_name.startswith('cs') else False

            )
            all_test_performances = all_test_performances + cell_means
            all_cell_test_performances[retina] = cell_stats
        with open(f'{result_path}/cs_wn_for_nm_test_correlations_r1_r4_finetuned_temp.pkl', 'wb') as f:
            pickle.dump(all_cell_test_performances, f)
