import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
from nerfstudio.process_data.colmap_utils import colmap_to_json
from quidient.data.dataset import ARKitProcessor
from examples.ert.recon_with_ert import RunERT
import types
from pathlib import Path
import wx

def process_arkit_data(data_path):

  processed_root = str(Path(data_path).parent / "processed")
  
  dataset_conf = types.SimpleNamespace()
  dataset_conf.data = types.SimpleNamespace()
  dataset_conf.data.data_root = data_path
  dataset_conf.data.data_type = "arkit"
  dataset_conf.data.load_step_size = 1
  dataset_conf.data.pose_estimator = "arkit"
  dataset_conf.data.process_step_size = 1
  dataset_conf.data.processed_root = processed_root

  dataset_conf.flags = types.SimpleNamespace()
  dataset_conf.flags.read_depth = True
  dataset_conf.flags.read_rgb = True
  dataset_conf.flags.video = False

  ert_conf = types.SimpleNamespace()
  ert_conf.exp_config_path = None
  ert_conf.training = types.SimpleNamespace()
  ert_conf.training.experiment_name = "garden_ref"
  ert_conf.training.vis = "viewer+tensorboard"
  ert_conf.training.cache_device = "cpu"
  ert_conf.training.steps_per_save = 1000
  ert_conf.training.max_num_iterations = 9000

  ert_conf.data = types.SimpleNamespace()
  ert_conf.data.processed_root = processed_root
  ert_conf.data.name = "mipnerf360"
  ert_conf.data.scene_name = "garden"

  ert_conf.flags = types.SimpleNamespace()
  ert_conf.flags.finetune_3dgs_with_3dgrt = False
  ert_conf.flags.latest_checkpoint = True
  ert_conf.flags.auto_quit_viewer = False
  ert_conf.flags.automesh = True
  ert_conf.flags.enable_profiling = True
  ert_conf.flags.adaptive_control = True
  ert_conf.flags.refine_cams = False
  ert_conf.flags.train_with_depth = False
  ert_conf.flags.enable_dd_loss = False
  ert_conf.flags.enable_normal_consistency_loss = False
  ert_conf.flags.enable_uncertainty = False
  ert_conf.flags.blender_data = False
  ert_conf.flags.blendernerf_data = False
  ert_conf.flags.force_optimize = False
  ert_conf.flags.generate_report = True
  ert_conf.flags.normal_init = False
  ert_conf.flags.autoscale_poses = True
  ert_conf.flags.refine_radiometry = False

  ert_conf.param = types.SimpleNamespace()
  ert_conf.param.load_checkpoint = "None"
  ert_conf.param.seed_point_source = "colmap"
  ert_conf.param.mode = "3dgrt"
  ert_conf.param.strategy = "default"
  ert_conf.param.min_alpha = 0.01
  ert_conf.param.warmup_length = 500
  ert_conf.param.stop_split_at = 8000
  ert_conf.param.refine_every = 100
  ert_conf.param.reset_alpha_every = 30
  ert_conf.param.cull_alpha_thresh = 0.03
  ert_conf.param.cull_scale_thresh = 0.05
  ert_conf.param.densify_grad_thresh = 0.0003
  ert_conf.param.densify_norm_thresh = 2e-9
  ert_conf.param.densify_size_thresh = 0.05
  ert_conf.param.max_gaussian_count = 3000000
  ert_conf.param.lambda_normal = 1.0
  ert_conf.param.lambda_normal_smooth = 1.0
  ert_conf.param.lambda_depth = 0.2
  ert_conf.param.bound_extent = 1.8
  ert_conf.param.downscale_factor = 8.0
  ert_conf.param.train_split_fraction = 1.0
  ert_conf.param.radiometry_refine_warmup = 2000
  ert_conf.param.seed_point_source = "depth"

  ert_conf.results = types.SimpleNamespace()
  ert_conf.results.output_dir = None
  ert_conf.results.report_path = None
  ert_conf.results.voxel_size = -1.0
  ert_conf.results.depth_trunc = 4.0

  processor = ARKitProcessor(dataset_conf)
  processor.process()
  ert = RunERT(ert_conf.exp_config_path)
  depth_dict = ert.process_depth(ert_conf.param, ert_conf.flags, ert_conf.data)
  colmap_to_json(Path(processed_root + "/sparse/0"), Path(processed_root), image_id_to_depth_path=depth_dict)
