"""
High-level inference for the JKTSV DINOv3 geolocation model.

Example
-------
    from inference import GeoTagPredictor

    predictor = GeoTagPredictor("nadh0708/JKTSV-modelD")   # HF repo id
    # ...or a local checkpoint:
    predictor = GeoTagPredictor("model/modelD_40e_2.pth")

    result = predictor.predict("street.jpg")
    # {'lat': -6.21, 'lon': 106.84}

    results = predictor.predict(["a.jpg", "b.jpg"])         # batched
    # [{'lat': ..., 'lon': ...}, ...]

The model was trained on Google Street View perspective crops of Jakarta roads
(8 headings, 0-315 deg). Predictions are only meaningful for Jakarta street
imagery.
"""

from __future__ import annotations

from typing import Union

import torch
from PIL import Image
from torchvision import transforms

from modeling_geotag import DinoGeoRegressor

_IMAGENET_MEAN = [0.485, 0.456, 0.406]
_IMAGENET_STD = [0.229, 0.224, 0.225]

ImageInput = Union[str, Image.Image]


class GeoTagPredictor:
    def __init__(
        self,
        model_id_or_path: str,
        device: str | torch.device | None = None,
        filename: str = "pytorch_model.bin",
    ):
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.model = DinoGeoRegressor.from_pretrained(
            model_id_or_path, filename=filename, device=self.device
        )
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
        ])

    def _load(self, image: ImageInput) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        return Image.open(image).convert("RGB")

    @torch.no_grad()
    def predict(
        self, images: Union[ImageInput, list[ImageInput]]
    ) -> Union[dict, list[dict]]:
        """Predict (lat, lon) for one image or a list of images."""
        single = not isinstance(images, (list, tuple))
        batch = [images] if single else list(images)

        pixel_values = torch.stack([self.transform(self._load(im)) for im in batch])
        pixel_values = pixel_values.to(self.device)

        lonlat = self.model.predict_lonlat(pixel_values).cpu()
        results = [
            {"lat": float(lat), "lon": float(lon)}
            for lon, lat in lonlat.tolist()
        ]
        return results[0] if single else results


def _cli() -> None:
    import argparse
    import json

    parser = argparse.ArgumentParser(description="Geolocate Jakarta street imagery.")
    parser.add_argument("model", help="HF repo id or local checkpoint path")
    parser.add_argument("images", nargs="+", help="image file path(s)")
    parser.add_argument("--device", default=None)
    parser.add_argument(
        "--filename", default="pytorch_model.bin",
        help="weights filename inside the HF repo",
    )
    args = parser.parse_args()

    predictor = GeoTagPredictor(args.model, device=args.device, filename=args.filename)
    out = predictor.predict(args.images)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    _cli()