import os
import argparse

file_dir = os.path.dirname(__file__)  # the directory that options.py resides in


class MonodepthOptions:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description="vddepth options")

        # PATHS
        self.parser.add_argument("--data_path",
                                 type=str,
                                 help="path to the training data",
                                 default='../data/DrivingStereo')
        self.parser.add_argument("--log_dir",
                                 type=str,
                                 help="log directory",
                                 default='log')

        # TRAINING options
        self.parser.add_argument("--one_optim",
                                 help="if set, we use naive optimisation",
                                 action="store_true")
        self.parser.add_argument("--ViT",
                                 help="if set, use monovit depth network",
                                 action="store_true")

        self.parser.add_argument("--split",
                                 type=str,
                                 help="which training split to use",
                                 choices=["eigen_zhou", "eigen_full", "odom", "benchmark"],
                                 default="eigen_zhou")
        self.parser.add_argument("--num_layers",
                                 type=int,
                                 help="number of resnet layers",
                                 default=18,
                                 choices=[18, 34, 50, 101, 152])
        self.parser.add_argument("--dataset",
                                 type=str,
                                 help="dataset to train on",
                                 default="DrivingStereo",
                                 choices=["kitti", "kitti_odom", "kitti_depth", "kitti_test", "waymo_da", "DrivingStereo"])
        self.parser.add_argument("--png",
                                 help="if set, trains from raw KITTI png files (instead of jpgs)",
                                 action="store_true")
        self.parser.add_argument("--height",
                                 type=int,
                                 help="input image height",
                                 default=192)
        self.parser.add_argument("--width",
                                 type=int,
                                 help="input image width",
                                 default=512)
        self.parser.add_argument("--scales",
                                 nargs="+",
                                 type=int,
                                 help="scales used in the loss",
                                 default=[0, 1, 2, 3])
        self.parser.add_argument("--min_depth",
                                 type=float,
                                 help="minimum depth",
                                 default=0.1)
        self.parser.add_argument("--max_depth",
                                 type=float,
                                 help="maximum depth",
                                 default=80.0)
        self.parser.add_argument("--use_stereo",
                                 help="if set, uses stereo pair for training",
                                 action="store_true")
        self.parser.add_argument("--frame_ids",
                                 nargs="+",
                                 type=int,
                                 help="frames to load",
                                 default=[0, -1, 1])
        self.parser.add_argument("--boarder",
                                type=int,
                                help="when large boarder value, the borader of image gradient gets thicker",
                                default=1)
        self.parser.add_argument("train_test_split",
                                 help="",
                                 action="store_true")

        # OPTIMIZATION options
        self.parser.add_argument("--optim",
                                 type=str,
                                 help="choose optimizer between sgd or adam",
                                 default='sgd')
        self.parser.add_argument("--batch_size",
                                 type=int,
                                 help="batch size",
                                 default=1)
        self.parser.add_argument("--learning_rate",
                                 type=float,
                                 help="learning rate",
                                 default=1e-4)
        self.parser.add_argument("--momentum",
                                 help="momentum for sgd",
                                 type=float,
                                 default=0.9)
        self.parser.add_argument("--num_epochs",
                                 type=int,
                                 help="number of epochs",
                                 default=1)
        self.parser.add_argument("--grad_clip",
                                 help="gradient clip",
                                 action="store_true")
        self.parser.add_argument("--max_grad_norm",
                                 help="max_grad_norm. using this with grad_clip",
                                 type=float,
                                 default=1.)


        self.parser.add_argument("--pytorch_random_seed",
                                 default=42,
                                 type=int)
        self.parser.add_argument("--depth_loss",
                                help="using depth loss",
                                action="store_true")
        self.parser.add_argument("--filter_type",
                                help="choose filter type",
                                default="mean",
                                choices=['mean', 'median'])
            
        # SYSTEM options
        self.parser.add_argument("--no_cuda",
                                 help="if set disables CUDA",
                                 action="store_true")
        self.parser.add_argument("--num_workers",
                                 type=int,
                                 help="number of dataloader workers",
                                 default=4)
        self.parser.add_argument("--val_num_workers",
                                 type=int,
                                 help="number of dataloader workers",
                                 default=4)

        # LOADING option
        self.parser.add_argument("--model",
                                 type=str,
                                 help="choose models",
                                 choices=['monodepth2', 'sgdepth', 'hrdepth', 'monovit', 'litemono', 'newcrfs', 'adabins'])
        self.parser.add_argument("--load_weights_folder",
                                 type=str,
                                 help="name of model to load",
                                 default='weights/mono_640x192')
        self.parser.add_argument("--label_producer_folder",
                                 type=str,
                                 help="name of model to produce pseudo label",
                                 default="./weights/mono+stereo_1024x320")
        self.parser.add_argument("--models_to_load",
                                 nargs="+",
                                 type=str,
                                 help="models to load",
                                 default=["encoder", "depth", "pose_encoder", "pose"])

        # LOGGING options
        self.parser.add_argument("--log_frequency",
                                 type=int,
                                 help="number of batches between each tensorboard log",
                                 default=250)
        self.parser.add_argument("--save_frequency",
                                 type=int,
                                 help="number of epochs between each save",
                                 default=1)
        self.parser.add_argument("--save_intermediate_models",
                                 help="if set, save the model each time we log to tensorboard",
                                 action='store_true')
        self.parser.add_argument('--save',
                                 help='decide save the model or not after adatation',
                                 action='store_true')
        self.parser.add_argument('--save_metric',
                                 help='decide save metric or not',
                                 action='store_true')


        # EVALUATION options
        self.parser.add_argument("--eval_stereo",
                                 help="if set evaluates in stereo mode",
                                 action="store_true")
        self.parser.add_argument("--robust_test",
                                 help="if set evaluates using robust images",
                                 action="store_true")

        self.parser.add_argument("--eval_mono",
                                 help="if set evaluates in mono mode",
                                 action="store_true")
        self.parser.add_argument("--disable_median_scaling",
                                 help="if set disables median scaling in evaluation",
                                 action="store_true")
        self.parser.add_argument("--pred_depth_scale_factor",
                                 help="if set multiplies predictions by this number",
                                 type=float,
                                 default=1)
        self.parser.add_argument("--ext_disp_to_eval",
                                 type=str,
                                 help="optional path to a .npy disparities file to evaluate")
        self.parser.add_argument("--eval_split",
                                 type=str,
                                 default="sunny",
                                 choices=["eigen", "eigen_benchmark", "eigen_zhou", "benchmark", "odom_9",
                                          "odom_10", "cityscapes", "nuScenes_val", "sunny", "rainy", "foggy", "cloudy", "all_weather",
                                          "waymo_da"],
                                 help="which split to run eval on")
        self.parser.add_argument("--benchmarking",
                                 help="if set, run in benchmarking mode."
                                      "The runtime is faster, but the results are slightly more random",
                                 action="store_true")
        self.parser.add_argument("--adapt_data_list",
                                 type=str,
                                 help='adaptation dataset',
                                 default='test_files.txt')
        self.parser.add_argument("--save_pred_disps",
                                 help="if set saves predicted disparities",
                                 action="store_true")
        self.parser.add_argument("--no_eval",
                                 help="if set disables evaluation",
                                 action="store_true")
        self.parser.add_argument("--eval_eigen_to_benchmark",
                                 help="if set assume we are loading eigen results from npy but "
                                      "we want to evaluate using the new benchmark.",
                                 action="store_true")
        self.parser.add_argument("--eval_out_dir",
                                 help="if set will output the disparities to this folder",
                                 type=str)
        self.parser.add_argument("--post_process",
                                 help="if set will perform the flipping post processing "
                                      "from the original monodepth paper",
                                 action="store_true")
        
        # Visualize while evaluation
        self.parser.add_argument("--visualize",
                                 help="visualize the evaluation results",
                                 action="store_true")
        self.parser.add_argument("--error_range",
                                 type=float,
                                 help="the range of the error to visualize, from 0 to error_range, in meters",
                                 default=200)
        self.parser.add_argument("--vis_name",
                                 help="saved error figure name",
                                 type=str,
                                 default='diff')
        
        # Custom options
        self.parser.add_argument("--tta",
                                 help="Whether to adapt in test time",
                                 action="store_true")
        self.parser.add_argument("--train_bn",
                                 help="If set, turn on training mode",
                                 action="store_true")
        self.parser.add_argument("--tta_params",
                                 help="Which parameters to adapt",
                                 type=str,
                                 choices=['even', 'norm', 'dynamic'],
                                 default='norm')
        self.parser.add_argument("--edge_loss",
                                help="Wheather to use edge loss",
                                action="store_true")
        self.parser.add_argument("--lamb",
                                 help="lambda. decide magnitude of edge loss",
                                 type=float,
                                 default=0.2)
        self.parser.add_argument("--filter_size",
                                 help="decide size of filter of Laplacian like edge detector",
                                 type=int,
                                 default=1)
        self.parser.add_argument("--device",
                                 help="pick where the data is load",
                                 default='cpu')
        self.parser.add_argument("--save_results",
                                 help="Whether to save results of the sequence",
                                 action="store_true")
        self.parser.add_argument("--save_disp",
                                 help="Where to save disparity results of the sequence",
                                 action="store_true")
        self.parser.add_argument("--pan_size",
                                help="decide the size of panoptic segmentation model",
                                type=str,
                                default="tiny")
        self.parser.add_argument("--ada_ratio",
                                help="decide adaptation ratio among encoder",
                                type=float,
                                default=0.1)
        self.parser.add_argument("--eval_last_seq",
                                help="decide scale and evaluate last sequence of the dataset sequence",
                                action="store_true")
        self.parser.add_argument("--edge_order",
                                help="decide first or second order grdient",
                                choices=['first', 'second'])
        
    def parse(self):
        self.options = self.parser.parse_args()
        return self.options
