class ObjectDetector:
    def __init__(self, config_path):
        torch.cuda.empty_cache()
        with open(config_path, 'r') as file:
            self.config = json.load(file)
        
        # Load the default config file
        self.cfg =  mmcv.Config.fromfile('/home/biosense/Documents/Rana/Ronit/SWIN/Swin-Transformer-Object-Detection/configs/_base_/models/cascade_mask_rcnn_r50_fpn.py')
        # self.cfg = ''
        self._update_config()
        self._save_config()

    def _update_config(self):
        # Dataset configuration
        self.cfg.dataset_type = 'CocoDataset'
        self.cfg.data_root = self.config['train_folder']

        self.cfg.data.train.type = self.cfg.dataset_type
        self.cfg.data.train.data_root = self.config['train_folder']
        self.cfg.data.train.ann_file = self.config['train_json']
        self.cfg.data.train.img_prefix = ''

        self.cfg.data.val.type = self.cfg.dataset_type
        self.cfg.data.val.data_root = self.config['val_folder']
        self.cfg.data.val.ann_file = self.config['val_json']
        self.cfg.data.val.img_prefix = ''

        self.cfg.data.test.type = self.cfg.dataset_type
        self.cfg.data.test.data_root = self.config['test_folder']
        self.cfg.data.test.ann_file = self.config['test_json']
        self.cfg.data.test.img_prefix = ''

        # Model settings
        self.cfg.model.roi_head.bbox_head.num_classes = self.config['num_classes']

        # Training settings
        self.cfg.runner.max_epochs = self.config['epochs']
        self.cfg.work_dir = './results'
        
        # Optimizer settings
        self.cfg.optimizer.lr = self.config['learning_rate']
        self.cfg.optimizer.weight_decay = self.config['weight_decay']
        
        # Learning rate scheduler settings
        self.cfg.lr_config.policy = 'step'
        self.cfg.lr_config.warmup = self.config['warmup']
        self.cfg.lr_config.warmup_iters = self.config['warmup_iters']
        self.cfg.lr_config.warmup_ratio = self.config['warmup_ratio']
        self.cfg.lr_config.step = self.config['step_lr_policy']

        # Logging settings
        self.cfg.log_config.interval = self.config['log_interval']

        # Evaluation settings
        self.cfg.evaluation.metric = 'bbox'
        self.cfg.evaluation.interval = self.config['evaluation_interval']

        # Checkpoint settings
        self.cfg.checkpoint_config.interval = self.config['checkpoint_interval']

        # Seed for reproducibility
        self.cfg.seed = 0
        set_random_seed(0, deterministic=False)
        self.cfg.gpu_ids = range(1)

        print(f'Config:\n{self.cfg.pretty_text}')

    def _save_config(self):
        # Save the updated config to a file
        self.cfg.dump('./updated_config.py')

    def train(self):
        # Build dataset and model
        mmdet.datasets.coco.CocoDataset.CLASSES = ('pt', 'gp')
        datasets = [build_dataset(self.cfg.data.train)]
        model = build_detector(self.cfg.model, train_cfg=self.cfg.get('train_cfg'), test_cfg=self.cfg.get('test_cfg'))
        model.CLASSES = datasets[0].CLASSES

        # Create directory for work
        mmcv.mkdir_or_exist(osp.abspath(self.cfg.work_dir))
        train_detector(model, datasets, self.cfg, distributed=False, validate=True)

if __name__ == '__main__':
    config_path = 'config.json'
    detector = ObjectDetector(config_path)
    detector.train()