#!/usr/bin/env python3
"""Run standalone GeoText-1652 image/text retrieval on an RGB GeoTIFF.

Example:
    IMAGE_URL="https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/geoembeddings-demo/building-detection_sm.tif"
    python infer.py \
      --image-url "$IMAGE_URL" \
      --query "a building" \
      --query "a sports field"
"""

from __future__ import annotations

import argparse
from io import BytesIO
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import requests
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from rasterio.io import MemoryFile
from ruamel.yaml import YAML
from torchvision import transforms

MODEL_REPOSITORY = "truemanv5666/GeoText1652_model"
CHECKPOINT_NAME = "geotext_official_checkpoint.pth"
DEFAULT_IMAGE_URL = (
    "https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/"
    "geoembeddings-demo/building-detection_sm.tif"
)
CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--image-url", default=DEFAULT_IMAGE_URL, help="Public RGB GeoTIFF/COG URL."
    )
    parser.add_argument(
        "--query",
        action="append",
        required=True,
        help="Text query to score. Pass this option repeatedly to rank multiple queries.",
    )
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=Path(
            os.environ.get(
                "GEOTEXT_SOURCE_DIR", Path(__file__).resolve().parent / "GeoText-1652"
            )
        ),
        help="Checkout of https://github.com/MultimodalGeo/GeoText-1652.",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path(os.environ.get("GEOTEXT_MODEL_DIR", Path(__file__).resolve().parent)),
        help="Directory holding the Hugging Face checkpoint. Downloaded on first use.",
    )
    parser.add_argument("--device", default="auto", choices=("auto", "cuda", "cpu"))
    parser.add_argument("--output", type=Path, help="Optional JSON output file.")
    return parser.parse_args()


def resolve_device(requested: str) -> torch.device:
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    if requested == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but no CUDA device is available.")
    return torch.device(requested)


def ensure_checkpoint(model_dir: Path) -> Path:
    checkpoint = model_dir / CHECKPOINT_NAME
    if not checkpoint.exists():
        model_dir.mkdir(parents=True, exist_ok=True)
        snapshot_download(repo_id=MODEL_REPOSITORY, local_dir=model_dir)
    if not checkpoint.exists():
        raise FileNotFoundError(f"Expected checkpoint was not downloaded: {checkpoint}")
    return checkpoint


def fetch_rgb_geotiff(url: str) -> Tuple[Image.Image, Dict[str, Any]]:
    """Fetch RGB GeoTIFF/COG, JPEG, or PNG imagery from a public URL.

    The historical function name is kept so existing server imports remain
    compatible. JPEG and PNG inputs have no CRS or bounds, so their geospatial
    metadata fields are returned as ``None``.
    """
    response = requests.get(url, timeout=180)
    response.raise_for_status()

    try:
        with MemoryFile(response.content) as memory_file:
            with memory_file.open() as dataset:
                if dataset.count < 3:
                    raise ValueError(
                        f"GeoText requires RGB imagery; input has {dataset.count} band(s)."
                    )
                rgb = dataset.read([1, 2, 3], masked=True).astype(np.float32).filled(np.nan)
                metadata = {
                    "crs": str(dataset.crs),
                    "bounds": list(dataset.bounds),
                    "width": dataset.width,
                    "height": dataset.height,
                    "bands_used": [1, 2, 3],
                    "format": "geotiff",
                    "is_geospatial": True,
                }
        low, high = np.nanpercentile(rgb, (2, 98))
        if not np.isfinite(low) or not np.isfinite(high) or high <= low:
            raise ValueError("Unable to create an RGB image from the GeoTIFF pixel values.")
        stretched = np.nan_to_num((rgb - low) / (high - low), nan=0.0, posinf=1.0, neginf=0.0)
        stretched = np.clip(stretched, 0, 1)
        return Image.fromarray((stretched.transpose(1, 2, 0) * 255).astype(np.uint8)), metadata
    except Exception:
        pass

    try:
        with Image.open(BytesIO(response.content)) as source:
            image_format = source.format
            if image_format not in ("JPEG", "PNG"):
                raise ValueError(f"unsupported image format: {image_format}")
            image = source.convert("RGB").copy()
    except Exception as error:
        raise ValueError("image_url could not be decoded as GeoTIFF, JPEG, or PNG") from error
    return image, {
        "crs": None,
        "bounds": None,
        "width": image.width,
        "height": image.height,
        "bands_used": ["red", "green", "blue"],
        "format": image_format.lower(),
        "is_geospatial": False,
    }


def tile_image(
    image: Image.Image, tile_size: int, overlap: int
) -> List[Dict[str, Any]]:
    """Split an image into overlapping source-pixel tiles without padding."""
    if tile_size <= 0:
        raise ValueError("tile_size must be greater than zero")
    if overlap < 0 or overlap >= tile_size:
        raise ValueError("tile_overlap must satisfy 0 <= tile_overlap < tile_size")

    width, height = image.size
    stride = tile_size - overlap
    tiles = []
    tile_index = 0
    for row in range(0, height, stride):
        for column in range(0, width, stride):
            tile_width = min(tile_size, width - column)
            tile_height = min(tile_size, height - row)
            tiles.append(
                {
                    "column": column,
                    "height": tile_height,
                    "image": image.crop((column, row, column + tile_width, row + tile_height)),
                    "index": tile_index,
                    "row": row,
                    "width": tile_width,
                }
            )
            tile_index += 1
    return tiles


def load_model(source_dir: Path, model_dir: Path, device: torch.device):
    method_dir = source_dir / "Method"
    if not method_dir.is_dir():
        raise FileNotFoundError(f"GeoText source was not found at {method_dir}")
    sys.path.insert(0, str(method_dir))

    from models.model_re_bbox import Bench  # noqa: PLC0415
    from models.tokenization_bert import BertTokenizer  # noqa: PLC0415

    with (method_dir / "configs" / "re_bbox.yaml").open() as file:
        config = YAML(typ="safe").load(file)
    config["vision_config"] = str(method_dir / "configs" / "config_swinB_384.json")
    config["text_config"] = str(method_dir / "configs" / "config_bert.json")
    config["text_encoder"] = str(model_dir / "bert")

    tokenizer = BertTokenizer.from_pretrained(config["text_encoder"])
    model = Bench(config)
    model.load_pretrained(str(ensure_checkpoint(model_dir)), config, is_eval=True)
    return model.to(device).eval(), tokenizer, config


def _image_tensor(image: Image.Image, config: Dict[str, Any], device: torch.device):
    transform = transforms.Compose(
        [
            transforms.Resize(
                (config["image_res"], config["image_res"]), interpolation=Image.BICUBIC
            ),
            transforms.ToTensor(),
            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
        ]
    )
    return transform(image).unsqueeze(0).to(device)


@torch.inference_mode()
def text_embedding_outputs(
    model, tokenizer, config: Dict[str, Any], queries: List[str], device: torch.device
) -> Dict[str, Any]:
    """Encode text without reading an image or running the vision encoder."""
    text_input = tokenizer(
        queries,
        padding="max_length",
        truncation=True,
        max_length=config["max_tokens"],
        return_tensors="pt",
    ).to(device)
    text_embeds = model.get_text_embeds(text_input.input_ids, text_input.attention_mask)
    text_features = model.get_features(text_embeds=text_embeds)
    return {
        "embedding_dimension": text_features.shape[-1],
        "text_embeddings": [
            {"query": query, "embedding": embedding.float().cpu().tolist()}
            for query, embedding in zip(queries, text_features)
        ],
    }


@torch.inference_mode()
def image_embedding_outputs(
    model,
    config: Dict[str, Any],
    image: Image.Image,
    device: torch.device,
    include_global_embedding: bool = True,
    include_patch_embeddings: bool = False,
) -> Dict[str, Any]:
    """Encode an image without tokenizing text or running the text encoder."""
    image_embeds, _ = model.get_vision_embeds(_image_tensor(image, config, device))
    result = {"embedding_dimension": model.vision_proj.out_features}
    if include_global_embedding:
        image_feature = model.get_features(image_embeds=image_embeds)[0]
        result["image_embedding"] = image_feature.float().cpu().tolist()
    if include_patch_embeddings:
        patch_features = torch.nn.functional.normalize(
            model.vision_proj(image_embeds[:, 1:, :]), dim=-1
        )[0]
        patch_count = patch_features.shape[0]
        columns = int(round(patch_count**0.5))
        if columns <= 0 or patch_count % columns:
            columns = patch_count
        rows = patch_count // columns
        patches = []
        for index, embedding in enumerate(patch_features.float().cpu().tolist()):
            row, column = divmod(index, columns)
            patches.append(
                {
                    "patch_index": index,
                    "row": row,
                    "column": column,
                    "source_pixel_xyxy": [
                        image.width * column / columns,
                        image.height * row / rows,
                        image.width * (column + 1) / columns,
                        image.height * (row + 1) / rows,
                    ],
                    "embedding": embedding,
                }
            )
        result["patch_embeddings"] = {
            "embedding_dimension": patch_features.shape[-1],
            "grid": {"rows": rows, "columns": columns},
            "patches": patches,
        }
    return result


@torch.inference_mode()
def model_outputs(
    model,
    tokenizer,
    config: Dict[str, Any],
    image: Image.Image,
    queries: List[str],
    device: torch.device,
    include_embeddings: bool = False,
    include_bboxes: bool = False,
    include_patch_embeddings: bool = False,
) -> Dict[str, Any]:
    image_tensor = _image_tensor(image, config, device)
    text_input = tokenizer(
        queries,
        padding="max_length",
        truncation=True,
        max_length=config["max_tokens"],
        return_tensors="pt",
    ).to(device)

    image_embeds, _ = model.get_vision_embeds(image_tensor)
    text_embeds = model.get_text_embeds(text_input.input_ids, text_input.attention_mask)
    image_features, text_features = model.get_features(image_embeds, text_embeds)
    scores = (image_features @ text_features.T).flatten().float().cpu().tolist()
    result = {
        "ranked_queries": sorted(
            (
                {"query": query, "similarity": score}
                for query, score in zip(queries, scores)
            ),
            key=lambda item: item["similarity"],
            reverse=True,
        )
    }

    if include_embeddings:
        result["image_embedding"] = image_features[0].float().cpu().tolist()
        result["text_embeddings"] = [
            {"query": query, "embedding": embedding.float().cpu().tolist()}
            for query, embedding in zip(queries, text_features)
        ]

    if include_patch_embeddings:
        # GeoText's vision encoder returns one CLS token followed by its spatial
        # patch tokens.  Project every patch with the same retrieval head used
        # for the global CLS embedding, so their 256-D vectors share the text
        # embedding space.
        patch_features = torch.nn.functional.normalize(
            model.vision_proj(image_embeds[:, 1:, :]), dim=-1
        )[0]
        patch_count = patch_features.shape[0]
        columns = int(round(patch_count**0.5))
        if columns <= 0 or patch_count % columns:
            columns = patch_count
        rows = patch_count // columns
        patches = []
        for index, embedding in enumerate(patch_features.float().cpu().tolist()):
            row, column = divmod(index, columns)
            x1 = image.width * column / columns
            y1 = image.height * row / rows
            x2 = image.width * (column + 1) / columns
            y2 = image.height * (row + 1) / rows
            patches.append(
                {
                    "patch_index": index,
                    "row": row,
                    "column": column,
                    "source_pixel_xyxy": [x1, y1, x2, y2],
                    "embedding": embedding,
                }
            )
        result["patch_embeddings"] = {
            "embedding_dimension": patch_features.shape[-1],
            "grid": {"rows": rows, "columns": columns},
            "patches": patches,
        }

    if include_bboxes:
        repeated_image_embeds = image_embeds.expand(len(queries), -1, -1)
        boxes = model.predict_bbox(
            repeated_image_embeds, text_embeds, text_input.attention_mask
        ).float().cpu().tolist()
        result["text_conditioned_boxes"] = [
            {
                "query": query,
                "format": "normalized_cx_cy_width_height",
                "box": box,
            }
            for query, box in zip(queries, boxes)
        ]

    return result


def rank_queries(
    model,
    tokenizer,
    config: Dict[str, Any],
    image: Image.Image,
    queries: List[str],
    device: torch.device,
):
    return model_outputs(model, tokenizer, config, image, queries, device)["ranked_queries"]


def main() -> None:
    args = parse_args()
    device = resolve_device(args.device)
    image, image_metadata = fetch_rgb_geotiff(args.image_url)
    model, tokenizer, config = load_model(args.source_dir, args.model_dir, device)
    ranked_queries = rank_queries(model, tokenizer, config, image, args.query, device)
    result = {
        "model": MODEL_REPOSITORY,
        "image_url": args.image_url,
        "image_metadata": image_metadata,
        "device": str(device),
        "ranked_queries": ranked_queries,
    }
    output = json.dumps(result, indent=2)
    print(output)
    if args.output:
        args.output.write_text(output + "\n")


if __name__ == "__main__":
    main()