import argparse
import json
import logging
from pathlib import Path

import torch
import yaml

from train.test import infer_deepchoice, test_with_baselines
from train.train import build_model, build_split_loader, list_split_batches, synchronize_model_dims, train_deepchoice


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")


def _load_config(config_path):
    cfg = yaml.safe_load(Path(config_path).read_text())
    return cfg


def _to_jsonable(value):
    if isinstance(value, dict):
        return {key: _to_jsonable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_jsonable(item) for item in value]
    if hasattr(value, "tolist"):
        return value.tolist()
    return value

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default="configs/config_deepchoice_base.yaml")
    parser.add_argument("--mode", choices=["train", "test", "infer"], default="train")
    parser.add_argument("--weights_path", default=None)
    parser.add_argument("--split", default=None)
    parser.add_argument("--limit_files", type=int, default=None)
    parser.add_argument("--inference_output", default=None)
    parser.add_argument("--metrics_output", default=None)
    parser.add_argument("--comparison_config", default=None)
    parser.add_argument("--comparison_weights_path", default=None)
    parser.add_argument("--comparison_field_name", default="best_mlp")
    args = parser.parse_args()

    cfg = _load_config(args.config)
    synchronize_model_dims(cfg)

    if args.mode == "train":
        train_deepchoice(cfg, cfg["training"]["output_dir"])
        return

    model = build_model(cfg)
    if args.weights_path is None:
        raise ValueError("--weights_path is required for --mode test and --mode infer")
    weights_path = args.weights_path
    state_dict = torch.load(weights_path, map_location=cfg["training"]["device"], weights_only=False)
    model.load_state_dict(state_dict)

    comparison_model = None
    comparison_cfg = None
    if args.comparison_config is not None and args.comparison_weights_path is not None:
        comparison_cfg = _load_config(args.comparison_config)
        synchronize_model_dims(comparison_cfg)
        comparison_model = build_model(comparison_cfg)
        comparison_state_dict = torch.load(
            args.comparison_weights_path,
            map_location=comparison_cfg["training"]["device"],
            weights_only=False,
        )
        comparison_model.load_state_dict(comparison_state_dict)

    split_name = args.split or cfg["test"]["split"]
    split_paths = list_split_batches(cfg, split_name, limit=args.limit_files)
    if not split_paths:
        raise FileNotFoundError(f"No test batches found for split '{split_name}' in {cfg['data']['batches_root']}")
    loader = build_split_loader(
        cfg,
        split_name,
        shuffle=False,
        limit=args.limit_files,
        file_batch_size=cfg.get("test", {}).get("file_batch_size", 1),
    )

    if args.mode == "test":
        metrics = test_with_baselines(model, loader, cfg, n_classes=int(cfg["model"]["num_classes"]))
        logging.info(
            "Test split=%s | files=%s | loss=%.4f | mIoU=%.4f | mF1=%.4f | majority mIoU=%.4f mF1=%.4f | hard_vote mIoU=%.4f mF1=%.4f | anyview mIoU=%.4f mF1=%.4f | IoUs=%s",
            split_name,
            len(split_paths),
            metrics["loss"],
            metrics["miou"],
            metrics["mf1"],
            metrics["baselines"]["majority"]["miou"],
            metrics["baselines"]["majority"]["mf1"],
            metrics["baselines"]["hard_vote"]["miou"],
            metrics["baselines"]["hard_vote"]["mf1"],
            metrics["baselines"]["anyview"]["miou"],
            metrics["baselines"]["anyview"]["mf1"],
            metrics["ious"],
        )
        if args.metrics_output is not None:
            metrics_path = Path(args.metrics_output)
            metrics_path.parent.mkdir(parents=True, exist_ok=True)
            metrics_path.write_text(json.dumps(_to_jsonable(metrics), indent=2), encoding="utf-8")
        return

    inference_output = args.inference_output or str(Path(cfg["training"]["output_dir"]) / f"inference_{split_name}")
    result = infer_deepchoice(
        model,
        loader,
        cfg,
        inference_output,
        comparison_model=comparison_model,
        comparison_config=comparison_cfg,
        comparison_field_name=args.comparison_field_name,
    )
    logging.info(
        "Inference split=%s | files=%s | samples=%s | accuracy=%.4f | model mIoU=%.4f mF1=%.4f | majority mIoU=%.4f mF1=%.4f | hard_vote mIoU=%.4f mF1=%.4f | anyview mIoU=%.4f mF1=%.4f | output=%s | las_tiles=%s",
        split_name,
        len(split_paths),
        result["num_samples"],
        result["accuracy"],
        result["metrics"]["model"]["miou"],
        result["metrics"]["model"]["mf1"],
        result["metrics"]["majority"]["miou"],
        result["metrics"]["majority"]["mf1"],
        result["metrics"]["hard_vote"]["miou"],
        result["metrics"]["hard_vote"]["mf1"],
        result["metrics"]["anyview"]["miou"],
        result["metrics"]["anyview"]["mf1"],
        inference_output,
        len(result["las_paths"]),
    )
    if args.metrics_output is not None:
        metrics_path = Path(args.metrics_output)
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.write_text(json.dumps(_to_jsonable(result), indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()