import argparse
import os
from pathlib import Path

import numpy as np
import torch
import yaml
from matplotlib import colors
from nnfabrik import builder
from torch import nn, optim

from datasets.stas import (create_cell_file_from_config_version, get_cell_sta,
                           get_cell_svd, separate_time_space_sta, get_filter_sign)
from training.trainers import model_step, save_checkpoint, train_step, train_ln
from utils.global_functions import dataset_seed, get_cell_names

cuda = torch.cuda.is_available()
if cuda:
    device = "cuda"
else:
    device = "cpu"
divnorm = colors.TwoSlopeNorm(vcenter=0.0)


def get_receptive_field_slice(file, cell_index, time_frame_index):
    receptive_fields = np.load(file)
    receptive_field = receptive_fields[cell_index, time_frame_index]
    return receptive_field


parser = argparse.ArgumentParser()
parser.add_argument("--learning_rate", default=5e-3, type=float, help="learning rate")
parser.add_argument(
    "--epochs", default=500, type=int, help="number of epochs to train for"
)
parser.add_argument("--num_of_channels", default=25, type=int)
parser.add_argument(
    "--data_path",
    default=".",
    type=str,
    help="path to the data, if None, the root of the project/data is considered",
)
parser.add_argument("--rf_size", default=20, type=int, help="")
parser.add_argument("--num_of_trials", default=250, type=int)
parser.add_argument("--batch_size", default=12, type=int)
parser.add_argument("--image_width", default=200, type=int)
parser.add_argument("--image_height", default=150, type=int)
parser.add_argument("--l1", default=0.0, type=float)
parser.add_argument("--l2", default=0.0, type=float)
parser.add_argument("--crop", default=0, type=int, nargs=4)
parser.add_argument("--subsample", default=1, type=int)
parser.add_argument("--time_chunk", default=100, type=int)
parser.add_argument("--cell_index", default="2", type=str)
parser.add_argument("--retina_index", default='04', type=str)
parser.add_argument("--log_dir", default="None", type=str)
parser.add_argument("--gamma", default=0.0, type=float)
parser.add_argument("--stopper_patience", default=10, type=int)
parser.add_argument("--norm", default=0, type=int)
parser.add_argument("--sta", default=0, type=int)
parser.add_argument("--svd", default=0, type=int)
parser.add_argument("--gaussian_fit", default=0, type=int)
parser.add_argument('--scheduler_patience', default=5, type=int)
parser.add_argument("--do_nonlin", default=1, type=int)
parser.add_argument("--fancy_nonlin", default=1, type=int)
parser.add_argument("--wandb", default=0, type=int)
parser.add_argument("--model_seed", default=8, type=int)
parser.add_argument("--performance", default="validation", type=str)
parser.add_argument("--flip_sta", default=0, type=int)
parser.add_argument('--dog', default=0, type=int)
parser.add_argument('--hard_coded', default='0 6 3 7 4 1 8 2;5 9', type=str)
parser.add_argument('--cs', default=1, type=int)
parser.add_argument(
    "--config_file",
    default="config_05_ds4",
    type=str,
    help="config file with dataset paths, sta stuff etc",
)

if __name__ == "__main__":
    args = parser.parse_args()
    if args.data_path is None:
        basepath = os.path.dirname(os.path.abspath(__file__))
    else:
        basepath = args.data_path
    with open(
        f"{basepath}/data/marmoset_data/responses/{args.config_file}.yaml", "rb"
    ) as cf:
        config_dict = yaml.unsafe_load(cf)

    config_file = args.config_file
    hard_coded = args.hard_coded
    lr = args.learning_rate
    epochs = args.epochs
    stopper_patience = args.stopper_patience
    scheduler_patience = args.scheduler_patience
    num_of_channels = args.num_of_channels
    rf_size = args.rf_size
    num_of_trials = args.num_of_trials
    batch_size = args.batch_size
    img_w = args.image_width
    img_h = args.image_height
    data_type = config_dict["data_type"]
    time_chunk = args.time_chunk
    cs = args.cs == 1

    performance = args.performance
    flip_sta = args.flip_sta == 1

    filter_type = "ete"
    do_sta = False
    do_svd = False
    do_gaussian_fit = False
    dog = args.dog==1
    sta = None
    svd = None
    learn_filter = True
    if args.sta == 1:
        filter_type = "sta"
        do_sta = True
        learn_filter = False
        if args.gaussian_fit == 1:
            do_gaussian_fit = True

    do_svd = False
    if args.svd == 1:
        do_svd = True
        do_sta = True
        learn_filter = False
        filter_type = "svd"

    do_nonlin = True
    if args.do_nonlin == 0:
        do_nonlin = False

    fancy_nonlin = False
    if args.fancy_nonlin == 1:
        fancy_nonlin = True
        assert performance != "test"

    l1 = args.l1
    l2 = args.l2
    gamma = args.gamma
    retina_index = args.retina_index

    normalize_response = False
    if args.norm == 1:
        normalize_response = True

    wandb_log = True
    if args.wandb == 0:
        wandb_log = False

    if args.cell_index != "None":
        cells = [int(args.cell_index)]
    else:
        cells = None

    max_coordinate = None

    crop = args.crop
    subsample = args.subsample
    if args.log_dir == "None":
        log_dir = os.path.dirname(os.path.abspath(__file__))
    else:
        log_dir = args.log_dir

    model_seed = args.model_seed

    neuronal_data_path = os.path.join(basepath, config_dict["neuronal_data_path"])
    training_img_dir = os.path.join(basepath, config_dict["training_img_dir"])
    test_img_dir = os.path.join(basepath, config_dict["test_img_dir"])

    if not fancy_nonlin and not learn_filter:
        epochs = 1
    if cells is None:
        cells = [
            x
            for x in range(config_dict["cell_numbers"][retina_index])
        ]

    for cell in cells:
        if cell in config_dict["exclude_cells"][retina_index]:
            continue
        if cell not in get_cell_names(retina_index=retina_index, config=config_dict, correlation_threshold=0,
                                      explained_variance_threshold=0.15):
            print(f'skiping cell {cell} not reliable enough')
            continue
        dataset_fn = "datasets.white_noise_loader"
        dataset_config = dict(
            config=config_dict,
            neuronal_data_dir=neuronal_data_path,
            train_image_path=training_img_dir,
            test_image_path=test_img_dir,
            batch_size=batch_size,
            crop=crop,
            subsample=subsample,
            seed=dataset_seed,
            num_of_trials_to_use=num_of_trials,
            use_cache=True,
            movie_like=False,
            conv3d=True,
            num_of_frames=num_of_channels,
            cell_index=cell,
            retina_index=retina_index,
            normalize_responses=normalize_response,
            cell_indices_out_of_range=False,
            num_of_layers=1,
            time_chunk_size=time_chunk,
            retina_specific_crops=False,
            extra_frame=0,
            hard_coded=hard_coded
        )

        dataloaders = builder.get_data(dataset_fn, dataset_config)
        print(dataloaders)

        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])
            / subsample
        )
        img_w = int(
            (img_w - a_dataloader.dataset.crop[2] - a_dataloader.dataset.crop[3])
            / subsample
        )
        inputs, targets = next(iter(a_dataloader))
        if do_gaussian_fit:
            gaussian_fit_str = "_gf_1"
        else:
            gaussian_fit_str = ""
        if flip_sta:
            flip_str = '_flip_1'
        else:
            flip_str = ''
        model_prefix = 'ln_models_factorized_'

        cs_str = ''
        if cs:
            cs_str = 'cs_'
        model_dir = f"models/ln_models/{model_prefix}_{filter_type}/{data_type}/retina{retina_index}/cell_{cell}/{cs_str}lr_{lr}_whole_rf_{rf_size}_ch_{num_of_channels}_l2_{l2}_g_{gamma}_bs_{batch_size}_tr_{num_of_trials}_s_{subsample}_c_{crop}_fn_{args.fancy_nonlin}_do_n_{args.do_nonlin}{gaussian_fit_str}_dog_{args.dog}{flip_str}/"


        cell_file = config_dict["sta_file"]
        cell_file = create_cell_file_from_config_version(
            original_file=cell_file, cell_index=cell, retina_index=retina_index
        )
        sta_dir = config_dict["stas_path"]
        half_rf_size = int(rf_size / 2)
        receptive_field, max_coordinate = get_cell_sta(
            cell_file=f'{cell_file}',
            data_dir=f'{basepath}/{sta_dir}/',
            h=img_h,
            w=img_w,
            rf_size=(num_of_channels, rf_size, rf_size),
            flip_sta=flip_sta
        )

        receptive_field = a_dataloader.dataset.transform_image(receptive_field)
        sta = receptive_field
        sign = get_filter_sign(sta)
        if do_svd:
            svd = get_cell_svd(sta)
        else:
            sta = separate_time_space_sta(sta)
        size = sta[0].shape

        input_shape = size
        model_config = dict(
            input_shape=input_shape,
            num_of_neurons=targets.shape[1],
            num_of_channels=num_of_channels,
            l1=l1,
            dog=dog,
            gamma=gamma,
            fancy_nonlin=fancy_nonlin,
            do_sta=do_sta,
            do_svd=do_svd,
            do_nonlin=do_nonlin,
            learn_filter=learn_filter,
            sta=sta,
            svd=svd,
            seed=model_seed,
            sign=sign,
            fit_gaussian=do_gaussian_fit
        )

        loss_function = nn.PoissonNLLLoss
        loss_function_config = dict(
            log_input=False, reduction="sum"
        )

        optimizer = optim.Adam
        optimizer_config = dict(lr=lr, weight_decay=args.l2)

        scheduler = optim.lr_scheduler.ReduceLROnPlateau#(

        scheduler_config = dict(
            mode="max",
            patience=scheduler_patience,
            threshold=0.000001,
            verbose=True,
            factor=0.5
        )

        correlations = []
        train_correlations = []
        train_variances = []
        valid_losses = []
        train_losses = []
        valid_variances = []
        prediction_variances = []
        all_valid_prediction_variances = []
        true_train_prediction_variances = []
        true_valid_prediction_variances = []
        penalties = []

        weight_dir = f"{log_dir}/{model_dir}/weights/seed_{model_seed}/"
        stats_dir = f"{log_dir}/{model_dir}/stats/seed_{model_seed}/"
        config_dir = f"{log_dir}/{model_dir}/config/"

        Path(weight_dir).mkdir(exist_ok=True, parents=True)
        Path(stats_dir).mkdir(exist_ok=True, parents=True)
        Path(config_dir).mkdir(exist_ok=True, parents=True)

        max_avg_valid_corr = -1

        config = dict(
            epochs=epochs,
            base_path=basepath,
            img_w=img_w,
            img_h=img_h,
            rf_size=rf_size,
            model_dir=model_dir,
            log_dir=log_dir,
            size=size,
            scheduler_config=scheduler_config,
            optimizer_config=optimizer_config,
            loss_function_config=loss_function_config,
            model_config=model_config,
            dataloader_config=dataset_config,
            config_file=args.config_file,
            loss_function=loss_function,
            optimizer=optimizer,
            scheduler=scheduler,
            regularizer_start=0,
            max_coordinate=max_coordinate,
            flip_sta=flip_sta,
            scheduler_patience=scheduler_patience,
            stopper_patience=stopper_patience
        )

        train_ln(dataset_fn=dataset_fn,
                 dataloader_config=dataset_config,
                 model_config=model_config,
                 dataset_config_dict=config_dict,
                 model_dir=model_dir,
                 device=device,
                 max_coordinate=max_coordinate,
                 img_h=img_h,
                 img_w=img_w,
                 wandb_log=wandb_log,
                 config=config,
                 filter_type=filter_type,
                 cuda=device == 'cuda',
                 performance=performance,
                 nm=False,
                 save_every_epoch=False,
                 cs=cs
                 )