import os
import torch

from training.cross_stimulus_training import (
    train_ln_cross_stimulus,
)
import yaml
import argparse

print(torch.cuda.is_available())
print(torch.cuda.current_device())
parser = argparse.ArgumentParser(description="Pretty print json")

parser.add_argument("--seed", default=8, type=int)
parser.add_argument("--model_dir", default=".")
parser.add_argument("--data_dir", default=".")
parser.add_argument('--filename', default="cs_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")
parser.add_argument("--performance", default="validation")
parser.add_argument("--directory_prefix", default="")
parser.add_argument("--retina", default='01')
parser.add_argument('--cell', default='0')
parser.add_argument("--freeze_filter", default=0, type=int)
parser.add_argument("--change_nonlin", default=1, type=int)
parser.add_argument("--lr", default=1e-5, type=float)
parser.add_argument("--config_file", default="config_s4_fixed_all")
parser.add_argument("--loss", default="poisson", type=str)
parser.add_argument("--flip_sta", default=0, type=int)

if __name__ == "__main__":
    args = parser.parse_args()

    wandb_log = True
    model_dir = args.model_dir
    data_dir = args.data_dir

    with open(
        f"{data_dir}/data/marmoset_data/responses/{args.config_file}.yaml", "rb"
    ) as config_file:
        config_dict = yaml.unsafe_load(config_file)

    filename = args.filename
    performance = args.performance
    directory_prefix = "ln_models_factorized__ete"
    data_type = "marmoset"
    setting = "wn_for_nm"
    config_dict = config_dict
    seed = args.seed
    lr = args.lr
    cell = args.cell
    retina = args.retina
    freeze_filter = args.freeze_filter == 1
    # fixation_file='complete_fixations_seed_2022and2023_s4.txt',
    # stimulus_seed=0,
    device = "cuda"

    scheduler = None
    optimizer = None
    if args.loss == "mse":
        loss_function = torch.nn.MSELoss()
    else:
        loss_function = torch.nn.PoissonNLLLoss(log_input=False, reduction="sum")

    train_ln_cross_stimulus(
        model_dir,
        filename,
        directory_prefix,
        data_dir,
        seed,
        config_dict,
        retina=retina,
        cell_index=cell,
        fixation_file=None,
        stimulus_seed=None,
        lr=lr,
        data_type=data_type,
        setting=setting,
        device=device,
        change_nonlin=args.change_nonlin == 1,
        stopper_patience=50,
        wandb_log=wandb_log,
        freeze_filter=freeze_filter,
        flip_sta=True,
        save_every_epoch=False,
        flip_imgs=False,
        log_dir=model_dir,
        fancy_nonlin=True,
        rank_model=True

    )
