#!/usr/bin/env python3
"""Run annotation pipeline for GTA-UAV-LR dataset.

GTA-UAV-LR: synthetic dataset from GTA V engine.
- drone/images/: 33763 images, 512x384, RGB PNG
- satellite/:    14640 images, 256x256, RGBA PNG (alpha = map boundary)
- Total: 48403 images
- 6 flight heights: 100, 200, 300, 400, 500, 600 meters

Usage:
    python scripts/run_gta_uav.py
    python scripts/run_gta_uav.py --source db       # only satellite (14.6K)
    python scripts/run_gta_uav.py --source drone     # only drone (33.8K)
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

import numpy as np
import torch

from src.conf.hardware_conf import HardwareConfig
from src.conf.input_conf import InputConfig
from src.conf.models_conf import ModelsConfig
from src.conf.pipeline_conf import PipelineConfig
from src.conf.seg_conf import SegConfig
from src.augmentor.io_utils import setup_logging
from src.main import run_pipeline
from scripts.seg_classes import UNIFIED_PROMPTS


INPUT_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
OUTPUT_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-aug"


def main() -> None:
    parser = argparse.ArgumentParser(description="Annotate GTA-UAV-LR")
    parser.add_argument("--source", choices=["db", "drone", "all"], default="all",
                        help="Process only db (satellite), drone, or all (default)")
    parser.add_argument("--stages", nargs="+",
                        default=["depth", "edges", "segmentation", "chmv2"],
                        help="Stages to run")
    args = parser.parse_args()

    import gin
    gin.clear_config()

    source = None if args.source == "all" else args.source
    if source == "drone":
        source = "query"

    pipeline_conf = PipelineConfig(
        input_root=INPUT_ROOT,
        output_root=OUTPUT_ROOT,
        stages=args.stages,
        save_npy=False,
        save_vis=True,
        save_safetensors=True,
        cleanup_npy=True,
        seg_fix_dark_water=True,
        seg_reclassify_wetland=True,
        resume=True,
        source=source,
        log_level="INFO",
    )

    hw_conf = HardwareConfig(
        profile_name="rtx4090",
        total_ram_gb=24.0,
        reserve_gb=2.0,
        use_fp16=True,
        batch_size=None,
        num_workers=4,
    )

    # GTA-UAV: satellite 256x256, drone 512x384
    # Use 256 for satellite, 512 for drone (non-square → resize to square)
    input_conf = InputConfig(image_size=256, query_image_size=512)

    # GTA V synthetic scenes: urban, suburban, rural, coastal, mountainous
    # 11 base classes + pool (swimming pools common in GTA suburbs)
    seg_conf = SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)

    models_conf = ModelsConfig(weights_dir=str(_PROJECT_ROOT / "in" / "weights"))

    setup_logging(
        pipeline_conf.log_level,
        log_file=Path(OUTPUT_ROOT) / "pipeline.log",
    )

    torch.manual_seed(42)
    np.random.seed(42)

    run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)


if __name__ == "__main__":
    main()