import argparse
import os

import numpy as np
import torch
import yaml
from matplotlib import colors
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 train_ln
from utils.global_functions import get_cell_names

cuda = torch.cuda.is_available()
if cuda:
    device = "cuda"
else:
    device = "cpu"

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=0.005, type=float, help="learning rate")
parser.add_argument(
    "--epochs", default=10000, type=int, help="number of epochs to train for"
)
parser.add_argument("--num_of_channels", default=43, 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=40, type=int, help="")
parser.add_argument("--num_of_trials", default=20, type=int)
parser.add_argument("--batch_size", default=24, 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='0', type=str)
parser.add_argument("--retina_index", default='01', type=str)
parser.add_argument("--log_dir", default=".", type=str)
parser.add_argument("--gamma", default=0.0, type=float)
parser.add_argument("--stopper_patience", default=20, 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('--dog', default=0, type=int)
parser.add_argument("--do_nonlin", default=1, type=int)
parser.add_argument("--fancy_nonlin", default=0, type=int)
parser.add_argument("--wandb", default=1, type=int)
parser.add_argument("--model_seed", default=16, type=int)
parser.add_argument('--performance', default='validation', type=str)
parser.add_argument('--train_again', default=1, type=int)
parser.add_argument('--hard_coded', default="7 11 6 8 3 12 13;4 1 2 5 0")
parser.add_argument('--scheduler_patience', default=5, type=int)
parser.add_argument('--skip_cells', default=1, type=int)
parser.add_argument('--fit_gaussian', default=0, type=int)
parser.add_argument('--stride', default=1, type=int)
parser.add_argument('--loss', default='mse', type=str)
parser.add_argument('--components', default=8, type=str)
parser.add_argument("--config_file",
    default="config_s4_fixed",
    type=str,
    help="config file with dataset paths, sta stuff etc",)
parser.add_argument("--stimulus_seed",
    default=0,
    type=int,
    help="stimulus seed from which the data was generated, only valid for natural movie stimulus",)
parser.add_argument('--flip_sta', default=0, type=int)
parser.add_argument('--cs', default=0, type=int)
parser.add_argument('--flip_imgs', default=0, type=int)


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
    if args.log_dir == "None":
        log_dir = os.path.dirname(os.path.abspath(__file__))
    else:
        log_dir = args.log_dir
    with open(
        f"{basepath}/data/marmoset_data/responses/{args.config_file}.yaml", "rb"
    ) as config_file:
        config_dict = yaml.unsafe_load(config_file)
    config_file = args.config_file

    lr = args.learning_rate
    stimulus_seed = args.stimulus_seed

    stopper_patience = args.stopper_patience
    scheduler_patience = args.scheduler_patience
    fit_gaussian = args.fit_gaussian == 1
    dog = args.dog == 1
    performance = args.performance
    performance_str = '' if performance == 'validation' else '_test'
    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
    stride = args.stride
    cs = args.cs == 1

    if args.hard_coded != 'No':
        train_trials = [int(x) for x in args.hard_coded.split(';')[0].split(' ')]
        valid_trials = [int(x) for x in args.hard_coded.split(';')[1].split(' ')]
        hard_coded_trials = train_trials, valid_trials
    else:
        hard_coded_trials = None

    filter_type = "ete"
    flip_sta = args.flip_sta == 1

    sta = None
    svd = None
    learn_filter = True

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

    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

    epochs = args.epochs
    if (not learn_filter) and (not fancy_nonlin):
        epochs = 1

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

    normalize_response = False
    if args.norm == 1:
        normalize_response = True
    train_again = args.train_again != 0

    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

    os.listdir(f"{basepath}/data")
    model_seed = args.model_seed

    neuronal_data_path = os.path.join(basepath, config_dict["response_path"])
    img_dir = os.path.join(basepath, config_dict["image_path"])
    config_dict["stas_path"] = os.path.join(basepath, config_dict["stas_path"])
    print("image path", img_dir)

    if cells is None:
        cells = [
            x
            for x in range(
                0, config_dict["cell_numbers"][retina_index]
            )
            if x not in config_dict["exclude_cells"][retina_index]
        ]

    for cell in cells:
        if args.skip_cells == 1:
            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.frame_movie_loader"
        dataset_config = dict(
            config=config_dict,
            basepath=basepath,
            img_dir_name=img_dir,
            neuronal_data_dir=neuronal_data_path,
            all_image_path=img_dir,
            batch_size=batch_size,
            seed=None,
            train_frac=0.8,
            subsample=subsample,
            crop=0,
            num_of_trials_to_use=num_of_trials,
            num_of_frames=num_of_channels,
            cell_index=cell,
            retina_index=retina_index,
            device=device,
            time_chunk_size=time_chunk,
            num_of_layers=1,
            cell_indices_out_of_range=True,
            oracle_correlation_threshold=None,
            normalize_responses=False,
            full_img_h=300,
            full_img_w=350,
            padding=50,
            retina_specific_crops=False,
            stimulus_seed=stimulus_seed,
            hard_coded=hard_coded_trials,
            flip_imgs=args.flip_imgs==1
        )

        flip_sta_str = '_flipped' if flip_sta else ''
        model_prefix = 'ln_models_factorized_4'
        cs_str = ''
        if cs:
            cs_str = 'cs_'
        model_dir = f"models/ln_models/{model_prefix}_{filter_type}_nm/{data_type}/retina{retina_index}/cell_{cell}/{cs_str}lr_{lr}_whole_rf_{rf_size}_ch_{num_of_channels}_l1_{l1}_l2_{l2}_g_{gamma}_bs_{batch_size}_tr_{num_of_trials}_s_{subsample}_c_{crop}_n_{args.norm}_fn_{args.fancy_nonlin}_do_n_{args.do_nonlin}{flip_sta_str}_gf_{args.fit_gaussian}_dog_{args.dog}/"

        sign = None

        if sta:
            do_sta = 1
        else:
            do_sta = 0
        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=cell_file,
            data_dir=sta_dir,
            h=img_h,
            w=img_w,
            rf_size=(num_of_channels, rf_size, rf_size),
            flip_sta=flip_sta
        )
        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=1,
            num_of_channels=num_of_channels,
            l1=l1,
            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,
            fit_gaussian=fit_gaussian,
            dog=dog,
            sign=sign,
            stride=stride,
        )
        loss_function = nn.PoissonNLLLoss
        loss_function_config = dict(
            log_input=False, reduction="sum"
        )

        optimizer = optim.Adam
        optimizer_config = dict(lr=lr)

        scheduler = optim.lr_scheduler.ReduceLROnPlateau
        scheduler_config = dict(
            mode="max",
            patience=scheduler_patience,
            threshold=0.000001,
            verbose=True,
            factor=0.5
        )

        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,
            model_config=model_config,
            dataloader_config=dataset_config,
            config_file=args.config_file,
            loss_function_config = loss_function_config,
            loss_function = loss_function,
            optimizer = optimizer,
            optimizer_config = optimizer_config,
            scheduler=scheduler,
            scheduler_config=scheduler_config,
            regularizer_start=0,
            max_coordinate = max_coordinate,
            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=True,
                 save_every_epoch=False,
                 cs=cs)

