import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml

from common.classification_training import train_classifier
from common.seed import set_seed


def main() -> int:
    parser = argparse.ArgumentParser(description="Train soup density classifier from pot crop images.")
    parser.add_argument("--config", default="configs/density_classifier.yaml")
    parser.add_argument("--dry-run", action="store_true", help="Validate config and write a training plan only.")
    args = parser.parse_args()
    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
    set_seed(int(config.get("seed", 42)))
    output_dir = Path(config["output_dir"])
    report_dir = Path(config["report_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    report_dir.mkdir(parents=True, exist_ok=True)
    plan = {
        "task": "density_classifier",
        "model_name": config["model_name"],
        "num_classes": config["num_classes"],
        "use_steam_feature": config["use_steam_feature"],
        "problem_type": config["problem_type"],
        "metrics": ["exact_accuracy", "within_1_accuracy", "large_error_rate"],
        "note": "Dry-run only validates config and writes this plan. Remove --dry-run to train when pot crop images and manifest are ready.",
    }
    (report_dir / "train_plan.json").write_text(json.dumps(plan, indent=2), encoding="utf-8")
    if args.dry_run:
        print(json.dumps(plan, indent=2))
        return 0
    report = train_classifier(
        manifest=config["manifest"],
        image_dir=config["image_dir"],
        target_column="label_id",
        output_dir=config["output_dir"],
        report_dir=config["report_dir"],
        model_name=config["model_name"],
        num_classes=int(config["num_classes"]),
        epochs=int(config["epochs"]),
        batch_size=int(config["batch_size"]),
        lr=float(config["lr"]),
        device_name=config.get("device", "auto"),
        labels={int(key): value for key, value in config["labels"].items()},
    )
    print(json.dumps(report, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())