import sys
sys.path.append('/playpen-raid/Author/LucidAtlas/')
import matplotlib.pyplot as plt
plt.rcParams['figure.dpi']= 300
import argparse
from pipeline.load import *
import torch.utils.data as data_utils
import model.networks.basics.workspace as ws
import pipeline

if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="Train a LucidAtlas autodecoder")
    arg_parser.add_argument(
        "--experiment",
        "-e",
        dest="experiment_directory",
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/airway_namlss_v1_0123_full.json', #'/playpen-raid/Author/LucidAtlas/configs/airways/airway_mlp_v1_0123.json', #"/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_lucidatlas_v14_0122.json", #
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/airway_namlsslip_v1_0123_full.json',
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_mlp_v1_0124.json",
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_namlss_v1_0124_part.json",
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_lucidatlas_v14_0124_part.json",
        #default="/playpen-raid/Author/LucidAtlas/configs/OASISBrain/brain_mlp_reg.json",
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v1/airway_mlp.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v1/airway_namlss.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v1/airway_lucidatlas_v15_part.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/ToyData/lucidatlas_full.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/ToyData/toy_namlss.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/OASISBrain/v2/brain_namlss.json',
        default='/playpen-raid/Author/LucidAtlas/configs/OASISBrain/v2/brain_lucidatlas_part.json',
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v2/airway_namlss.json',
        #default="/playpen-raid/Author/LucidAtlas/configs/airways/v2/airway_lucidatlas_part.json",
        #default='/playpen-raid/Author/LucidAtlas/configs/airways/v3/airway_gamlss.json',
        #default="/playpen-raid/Author/LucidAtlas/configs/airways/v3/airway_lucidatlas_full.json",
        #default="/playpen-raid/Author/LucidAtlas/configs/airways/v3/airway_mlp.json",
        help="The experiment directory. This directory should include "
             + "experiment specifications in 'specs.json', and logging will be "
             + "done in this directory as well.",
    )
    arg_parser.add_argument(
        "--checkpoint",
        "-c",
        dest="checkpoint",
        default="latest",
        help="The checkpoint weights to use. This can be a number indicated an epoch "
        + "or 'latest' for the latest weights (this is the default)",
    )

    arg_parser.add_argument(
        "--train",
        dest="whether_train",
        default=0,
        type=int,
        help="whether to train from scratch",
    )

    arg_parser.add_argument(
        "--test",
        dest="whether_test",
        default=1,
        type=int,
        help="whether to test",
    )

    arg_parser.add_argument(
        "--vis",
        dest="whether_vis",
        default=1,
        type=int,
        help="whether to vis",
    )


    args = arg_parser.parse_args()

    if bool(args.whether_train):
        pipeline.train_with_val(args.experiment_directory)
    if bool(args.whether_test):
        pipeline.pred_and_eval(args.experiment_directory, which_set='test')
        pipeline.pred_and_eval(args.experiment_directory, which_set='train')
        pipeline.ind_pred_and_eval_from_t0(args.experiment_directory)
        pipeline.pred_and_eval_single_local(args.experiment_directory, which_set='test')
    if bool(args.whether_vis):
        pipeline.visualize_local(args.experiment_directory)
    if bool(args.whether_test):
        pipeline.pred_and_eval_single_indp(args.experiment_directory, which_set='test')
        pipeline.pred_and_eval_single_global(args.experiment_directory, which_set='test')
        #pipeline.ind_pred_and_eval_timeline_pairs(args.experiment_directory)
    if bool(args.whether_vis):
        pipeline.visualize_global(args.experiment_directory)
        pipeline.visualize_indp(args.experiment_directory)
        pipeline.visualize_correlation(args.experiment_directory)
    print('1')




