import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["WORLD_SIZE"] = "1"
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="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")
parser.add_argument('--filename', default='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',
                    type=str)
parser.add_argument("--performance", default="validation")
parser.add_argument("--directory_prefix", default="ln_models_factorized_4_ete_nm")
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-3, type=float)
parser.add_argument("--config_file", default="config_05_all")
parser.add_argument("--loss", default="poisson", type=str)
parser.add_argument("--flip_sta", default=0, type=int)
parser.add_argument('--learn_filter', default=0, type=int)

if __name__ == "__main__":
    args = parser.parse_args()
    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)

    wandb_log = True
    model_dir = args.model_dir
    filename = args.filename
    performance = args.performance
    directory_prefix = args.directory_prefix
    data_type = "marmoset"
    setting = "nm_for_wn"
    config_dict = config_dict
    seed = args.seed
    lr = args.lr
    cell = args.cell
    retina = args.retina
    freeze_filter = args.freeze_filter == 1
    device = "cuda" if torch.cuda.is_available() else 'cpu'

    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,
        performance=performance,
        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=True,
        log_dir=model_dir,
        learn_filter=args.learn_filter ==1,
        rank_model=True


    )
