import os
import yaml
import json
import argparse
import logging
from datetime import datetime
from pathlib import Path
import torch
from torch.utils.data import DataLoader


from train.train import train_model
from utils.logging_utils import init_logging
from dataset.preprocess_multi_processing import run_parallel
from train.train import evaluate_on_val
from train.test import run_test_inference
from dataset.lidar_dataset import LidarFusionDataset
from model.model import SimpleMLP


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default="config.yaml", help="Path to YAML configuration.")
    parser.add_argument("--mode", choices=["preprocess", "train", "val", "test"], default="train")
    parser.add_argument("--weights_path", default="best_model.pth", help="Path to model weights.")
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text())
    now = datetime.now()
    out_dir = Path(cfg["logging"]["save_dir"] + now.strftime("_%m_%d_%H_%M_%S"))
    os.makedirs(out_dir, exist_ok=True)
    init_logging(out_dir / "run.log")
    mode = args.mode
    if mode == "preprocess":
        split_file = os.path.join(cfg["dataset"]["root"], cfg["dataset"]["split_file"])
        with open(split_file) as f:
            split_dict = json.load(f)
        all_zones = split_dict["train"] + split_dict["val"] + split_dict["test"]
        run_parallel(cfg, all_zones, num_workers=cfg["dataset"]["pre-processing_num_workers"])
        
    elif mode == "train":
        train_model(cfg, out_dir)
        
    elif mode == "val":
        model = SimpleMLP(
            input_dim=cfg['dataset']['n_classes'] * 2,
            hidden_dims=cfg['model']['hidden_dims'],
            n_classes=cfg['dataset']['n_classes']).to(cfg['training']['device'])
        model.load_state_dict(torch.load(args.weights_path))
        val_set = LidarFusionDataset(cfg, split="val", shuffle_zones=False)
        val_loader = DataLoader(val_set, batch_size=1, num_workers=0, shuffle=False)
        miou, mf1, ious = evaluate_on_val(model, val_loader, cfg)
    elif mode == "test":
        model = SimpleMLP(
            input_dim=cfg['dataset']['n_classes'] * 2,
            hidden_dims=cfg['model']['hidden_dims'],
            n_classes=cfg['dataset']['n_classes']).to(cfg['training']['device'])
        model.load_state_dict(torch.load(args.weights_path))
        test_set = LidarFusionDataset(cfg, split="test", shuffle_zones=False)
        test_loader = DataLoader(test_set, batch_size=1, num_workers=0, shuffle=False)
        run_test_inference(model, test_loader, cfg, out_dir)
    
if __name__ == "__main__":
    main()