#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# training executable for MASt3R
# --------------------------------------------------------
import sys
sys.path.append('.')
sys.path.append('submodules/mast3r')
from mast3r.model import AsymmetricMASt3R
from mast3r.losses import ConfMatchingLoss, MatchingLoss, APLoss, Regr3D, InfoNCE, Regr3D_ScaleShiftInv
from mast3r.datasets import ARKitScenes, BlendedMVS, Co3d, MegaDepth, ScanNetpp, StaticThings3D, Waymo, WildRGBD

import mast3r.utils.path_to_dust3r  # noqa
# add mast3r classes to dust3r imports
import dust3r.training
dust3r.training.AsymmetricMASt3R = AsymmetricMASt3R
dust3r.training.Regr3D = Regr3D
dust3r.training.Regr3D_ScaleShiftInv = Regr3D_ScaleShiftInv
dust3r.training.MatchingLoss = MatchingLoss
dust3r.training.ConfMatchingLoss = ConfMatchingLoss
dust3r.training.InfoNCE = InfoNCE
dust3r.training.APLoss = APLoss

import dust3r.datasets
dust3r.datasets.ARKitScenes = ARKitScenes
dust3r.datasets.BlendedMVS = BlendedMVS
dust3r.datasets.Co3d = Co3d
dust3r.datasets.MegaDepth = MegaDepth
dust3r.datasets.ScanNetpp = ScanNetpp
dust3r.datasets.StaticThings3D = StaticThings3D
dust3r.datasets.Waymo = Waymo
dust3r.datasets.WildRGBD = WildRGBD
from src.datasets.scannet import Scannet
from src.datasets.scannetpp import Scannetpp
from src.datasets.megadepth import MegaDepth
dust3r.datasets.Scannet = Scannet
dust3r.datasets.Scannetpp = Scannetpp
dust3r.datasets.MegaDepth = MegaDepth

from src.model import LSM_MASt3R
dust3r.training.LSM_MASt3R = LSM_MASt3R
from src.losses import GaussianLoss
dust3r.training.GaussianLoss = GaussianLoss

from dust3r.training import get_args_parser as dust3r_get_args_parser  # noqa
from dust3r.training import train  # noqa

import yaml


def get_args_parser():
    parser = dust3r_get_args_parser()
    parser.prog = 'LSM_MASt3R training'
    
    # Load the configuration
    with open("configs/model_config.yaml", "r") as f:
        config = yaml.safe_load(f)
    
    # Convert the config dict to a string of keyword arguments
    config_str = ", ".join(f"{k}={v}" for k, v in config.items())
    
    # Set the default model string with parameters
    parser.set_defaults(model=f"LSM_MASt3R({config_str})")
    
    return parser


if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()
    train(args)