import gradio as gr
import torch
# Explicitly trust the repository before the image-matching library tries to load it
torch.hub.set_dir('/root/.cache/torch/hub') # Optional: ensure consistent cache dir
torch.hub.load("verlab/accelerated_features", "XFeat", pretrained=False, trust_repo=True)
import faiss
import numpy as np
import pandas as pd
import folium
from PIL import Image
from pathlib import Path
import torchvision.transforms as tfm
from torchvision.transforms import functional as F
import logging
import sys
import io
import base64
import random
import ast
import webdataset as wds
import os
import pickle
from functools import lru_cache
import tarfile
from huggingface_hub import hf_hub_download, snapshot_download

from models.apl_model_dinov2 import DINOv2FeatureExtractor

sys.path.append(str(Path("image-matching-models")))
sys.path.append(str(Path("image-matching-models/matching/third_party")))

import util_matching
from matching import get_matcher

logging.basicConfig(level=logging.INFO)


HF_TOKEN = os.getenv("HF_TOKEN")  # Set this as a secret in your Space


def ensure_files_exist():

    Path("./faiss_index").mkdir(exist_ok=True)
    Path("./data/webdataset_shards").mkdir(parents=True, exist_ok=True)
    """Check for required files and download if missing"""
    # 1. Check FAISS index
    if not Path("faiss_index/faiss_index_2021.bin").exists():
        print("Downloading FAISS index...")

        hf_hub_download(
            repo_id='pawlo2013/EarthLoc2_FAISS',
            filename="faiss_index.bin",
            local_dir="./faiss_index",
            token=HF_TOKEN,
            repo_type="dataset"
        )
    
    # 2. Check WebDataset shards
    shard_dir = Path("data/webdataset_shards")
    required_shards = [f"shard-{i:06d}.tar" for i in range(11)]  # Adjust range as needed
    required_indices = [f"{s}.index" for s in required_shards]
    
    missing_files = [
        f for f in required_shards + required_indices 
        if not (shard_dir / f).exists()
    ]
    
    if missing_files:
        print(f"Downloading {len(missing_files)} missing shard files...")
        snapshot_download(
            repo_id="pawlo2013/EarthLoc_2021_Database",
            local_dir=shard_dir,
            allow_patterns="*.tar*",  # Gets both .tar and .tar.index
            token=HF_TOKEN,
            repo_type="dataset"
        )

# --- Integration Point ---
# Call this BEFORE loading any models or datasets
ensure_files_exist()



# --- Paths and device ---
MODEL_CHECKPOINT_PATH = Path("weights/best_model_95.6.torch")
FAISS_INDEX_PATH = Path("faiss_index/faiss_index.bin")
CSV_MAPPING_PATH = Path("faiss_index/faiss_index_webdataset.csv")  # Updated CSV with shards

DEVICE = "cpu"
MATCHING_IMG_SIZE = 512
logging.info(f"Using device: {DEVICE}")

for path, desc in [
    (MODEL_CHECKPOINT_PATH, "Model checkpoint"),
    (FAISS_INDEX_PATH, "FAISS index"),
    (CSV_MAPPING_PATH, "Path mapping CSV"),
]:
    if not path.exists():
        raise FileNotFoundError(f"{desc} not found at: {path}")

MODEL_NAME = "xfeat_steerers"
matcher = get_matcher(MODEL_NAME, device=DEVICE, max_num_keypoints=2048,trust_repo=True)
if MODEL_NAME == "xfeat_steerers":
    matcher.model.dev = DEVICE

# Load mapping CSV with keys and shard paths
mapping_df = pd.read_csv(CSV_MAPPING_PATH, index_col="faiss_index")
parsed = mapping_df["key"].str.extract(r"@(?P<z>\d{1,2})_(?P<r>\d{1,5})_(?P<c>\d{1,5})@").astype("int32")
mapping_df = mapping_df.join(parsed)



logging.info(f"Loaded mapping CSV with {len(mapping_df)} entries.")

# Cache for opened shards: {shard_path: WebDataset object}
shard_cache = {}

def get_shard_dataset(shard_path):
    """Load or get cached WebDataset for a shard path."""
    if shard_path not in shard_cache:
        shard_cache[shard_path] = wds.WebDataset(shard_path, handler=wds.warn_and_continue)
    return shard_cache[shard_path]

@lru_cache(maxsize=100)
def load_index(index_path):
    with open(index_path, "rb") as f:
        return pickle.load(f)

def load_image_from_shard(key):
    row = mapping_df[mapping_df["key"] == key]
    if row.empty:
        return None
    
    shard_path = row.iloc[0]["shard_path"]
    index_path = shard_path + ".index"
    
    if not os.path.exists(index_path):
        return _load_via_linear_scan(shard_path, key)  # Fallback
    
    try:
        index = load_index(index_path)
        offset = index.get(key)
        if offset is None:
            return None
            
        with open(shard_path, "rb") as f:
            f.seek(offset)
            with tarfile.open(fileobj=f) as tar:
                member = tar.next()
                if member and member.name.startswith(key):
                    jpg_file = tar.extractfile(member)
                    return Image.open(io.BytesIO(jpg_file.read())).convert("RGB")
        return None
    except Exception as e:
        logging.error(f"Error loading {key}: {str(e)}")
        return _load_via_linear_scan(shard_path, key)  # Fallback on error

# Fallback linear scan (original method)
def _load_via_linear_scan(shard_path, key):
    dataset = get_shard_dataset(shard_path)
    for sample in dataset:
        if sample["__key__"] == key:
            if img_bytes := sample.get("jpg"):
                return Image.open(io.BytesIO(img_bytes)).convert("RGB")
    return None

def pil_to_base64(image):
    """Convert a PIL image to a base64-encoded string for HTML embedding."""
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{img_str}"

def create_map(
    final_footprint, candidate_num, filename, inliers, query_footprint=None
):
    """
    Create and return a Folium map (as HTML string) showing the final footprint (blue)
    and optionally the query's ground truth footprint (orange).
    """
    if final_footprint:
        center = final_footprint[0]
        zoom = 10
    elif query_footprint:
        center = query_footprint[0]
        zoom = 10
    else:
        center = [0, 0]
        zoom = 2

    m = folium.Map(location=center, zoom_start=zoom)

    if query_footprint:
        folium.Polygon(
            locations=query_footprint,
            popup="Ground Truth Query Footprint",
            color="orange",
            fill=True,
            fill_color="orange",
            fill_opacity=0.4,
        ).add_to(m)

    if final_footprint:
        footprint_text = "\n".join(
            [f"({lat:.4f}, {lon:.4f})" for lat, lon in final_footprint]
        )
        popup_text = (
            f"Predicted Footprint:<br>{footprint_text}<br><br>"
            f"Candidate: {candidate_num}<br>Inliers: {inliers}"
        )

        folium.Polygon(
            locations=final_footprint,
            popup=popup_text,
            color="blue",
            fill=True,
            fill_color="cyan",
            fill_opacity=0.5,
        ).add_to(m)

        folium.Marker(
            location=[final_footprint[0][0], final_footprint[0][1]],
            popup=f"Footprint Coordinates:<br>{footprint_text}",
            icon=folium.Icon(color="blue"),
        ).add_to(m)
    elif not query_footprint:
        folium.Marker(
            location=[0, 0],
            popup="No valid location found.",
            icon=folium.Icon(color="red"),
        ).add_to(m)

    return m._repr_html_()

def parse_zoom_row_col_from_key(key: str):
    """
    Extract zoom, row, col from the complex key string.
    The key format is like:
    2021_10_30_90_@34,30714@92,81250@35,46067@92,81250@35,46067@94,21875@34,30714@94,21875@10_0404_0776@2021@34,88391@93,51562@16489@0@
    
    The zoom, row, col are expected to be in the last underscore-separated fields,
    specifically the last three fields before the final '@' or at the end.
    
    This function tries to extract zoom, row, col as integers.
    """
    try:
        # Split by underscore
        parts = util_matching.get_image_metadata_from_path(key)
        

        image_id_str = parts[9]  # 9th field: image_id

        parts = image_id_str.split("_")
        zoom = int(parts[0])
        row = int(parts[1])
        col = int(parts[2])
        

        return zoom, row, col
    except Exception as e:
        raise ValueError(f"Failed to parse zoom,row,col from key: {key}") from e


def get_surrounding_tiles_sharded(candidate_key, zoom):
    """
    Given a candidate key, find all keys in mapping_df with the same zoom,
    and row/col within ±4 offsets, then load images from shards.
    Return list of (img, (row, col), key) sorted by row, col.
    """
    try:
        zoom, row, col = parse_zoom_row_col_from_key(candidate_key)
    except Exception as e:
        logging.warning(f"Failed to parse candidate key {candidate_key}: {e}")
        return []

    row_offsets = [-4, 0, 4]
    col_offsets = [-4, 0, 4]

    desired_rows = {row + r for r in row_offsets}
    desired_cols = {col + c for c in col_offsets}



    # ── 2.  Vectorised filter ──────────────────────────────────────────────────────
    mask = (
    (mapping_df["z"] == zoom) &
    (mapping_df["r"].isin(desired_rows)) &
    (mapping_df["c"].isin(desired_cols))
    )

    matched_rows = mapping_df[mask]

    tiles = []
    seen_positions = set()  # Track (row, col) to avoid duplicates

    for _, row_data in matched_rows.iterrows():
        k = row_data["key"]
        try:
            _, r, c = parse_zoom_row_col_from_key(k)
        except Exception:
            continue

        if (r, c) in seen_positions:
            continue  # Skip duplicate position

        img = load_image_from_shard(k)
        if img is not None:
            tiles.append((img, (r, c), k))
            seen_positions.add((r, c))

    tiles.sort(key=lambda t: (t[1][0], t[1][1]))
    return tiles

def compose_tiles_ordered_sharded(tiles, tile_size, candidate_indices):
    """
    Compose a 3x3 grid image from tiles loaded from shards.
    Missing tiles replaced with blank.
    """
    candidate_row, candidate_col = candidate_indices
    grid_img = Image.new("RGB", (tile_size[0] * 3, tile_size[1] * 3))
    blank = Image.new("RGB", tile_size, color=(0, 0, 0))

    tile_dict = {(rc[0], rc[1]): img for img, rc, key in tiles if img is not None}

    for i, row_offset in enumerate([-4, 0, 4]):
        for j, col_offset in enumerate([-4, 0, 4]):
            desired_row = candidate_row + row_offset
            desired_col = candidate_col + col_offset
            img = tile_dict.get((desired_row, desired_col), blank)

            if img.mode != "RGB":
                img = img.convert("RGB")

            img_resized = tfm.Resize(tile_size, antialias=True)(img).copy()

  

            grid_img.paste(img_resized, (j * tile_size[0], i * tile_size[1]))
    return grid_img

def run_matching(query_image, candidate_image, base_footprint):
    local_fm = None
    viz_params = None

    for iteration in range(4):
        (
            num_inliers,
            local_fm,
            predicted_footprint,
            pretty_footprint,
        ) = util_matching.estimate_footprint(
            local_fm,
            query_image,
            candidate_image,
            matcher,
            base_footprint,
            HW=MATCHING_IMG_SIZE,
            viz_params=viz_params,
        )

        if num_inliers == -1 or num_inliers is None:
            return -1, []

    if hasattr(predicted_footprint, "tolist"):
        best_footprint = predicted_footprint.tolist()
    else:
        best_footprint = predicted_footprint

    return num_inliers, best_footprint

# --- Load assets ---

logging.info("Loading assets. This may take a moment...")

try:
    model = DINOv2FeatureExtractor(
        model_type="vit_base_patch14_reg4_dinov2.lvd142m",
        num_of_layers_to_unfreeze=0,
        desc_dim=768,
        aggregator_type="SALAD",
    )
    logging.info(f"Loading model checkpoint from {MODEL_CHECKPOINT_PATH}...")
    model_state_dict = torch.load(MODEL_CHECKPOINT_PATH, map_location=DEVICE)
    model.load_state_dict(model_state_dict)
    model = model.to(DEVICE)
    model.eval()
    logging.info("DINOv2 model and checkpoint loaded successfully.")
except Exception as e:
    logging.error(f"Failed to load the model: {e}")
    raise

faiss_index = faiss.read_index(str(FAISS_INDEX_PATH))
num_db_images = faiss_index.ntotal // 4
logging.info(
    f"FAISS index loaded. Contains {faiss_index.ntotal} vectors for {num_db_images} unique images."
)

image_transform = tfm.Compose(
    [
        tfm.Resize((model.image_size, model.image_size), antialias=True),
        tfm.ToTensor(),
        tfm.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        ),
    ]
)
logging.info("Assets loaded. Gradio app is ready.")

# --- Core app logic ---

def search_and_retrieve(
    query_image: Image.Image, query_footprint_str: str, num_results: int
):
    progress = gr.Progress()
    query_footprint = None
    if query_footprint_str:
        try:
            print(query_footprint_str)
            query_footprint = ast.literal_eval(query_footprint_str)
            
            query_footprint = [list(coord) for coord in query_footprint]
        except (ValueError, SyntaxError):
            logging.warning("Could not parse query footprint string.")
            query_footprint = None

    if query_image is None:
        yield create_map(None, None, None, None), None
        return

    progress(0.1, desc="Preprocessing query")
    if query_image.mode == "RGBA":
        query_image = query_image.convert("RGB")

    image_tensor = image_transform(query_image).to(DEVICE)
    with torch.no_grad():
        descriptor = model(image_tensor.unsqueeze(0))
    descriptor_np = descriptor.cpu().numpy()


    progress(0.2, desc=f"Searching database for {num_results} neighbors")
    distances, indices = faiss_index.search(descriptor_np, num_results)
    flat_indices = indices.flatten()

    global_best_inliers = -1
    global_best_footprint = None
    global_candidate_num = None
    global_filename = None
    global_best_display_image = None

    candidate_infos = []
    processed_image_indices = set()

    query_tensor = tfm.ToTensor()(
        tfm.Resize((MATCHING_IMG_SIZE, MATCHING_IMG_SIZE), antialias=True)(query_image)
    )

    


    progress(0.4, desc="Processing candidates")
    for faiss_idx in flat_indices:

        image_index = faiss_idx % num_db_images

    
        best_rotation_index = faiss_idx // num_db_images
        query_tensor = F.rotate(query_tensor, [0, -90, -180, -270][best_rotation_index] ) 

        if image_index in processed_image_indices:
            continue

        processed_image_indices.add(image_index)
        candidate_num = len(candidate_infos) + 1

        try:
            candidate_row = mapping_df.loc[int(image_index)]
        except Exception as e:
            logging.warning(f"Failed to get candidate info for index {image_index}: {e}")
            continue

        candidate_key = candidate_row["key"]

        candidate_path_str = candidate_row["local_path"]

        shard_path = candidate_row['shard_path']

        base_footprint = util_matching.path_to_footprint(Path(candidate_path_str))

        
        try:
        

            parts = util_matching.get_image_metadata_from_path(candidate_key)
            
            image_id_str = parts[9]  # 9th field: image_id

            parts = image_id_str.split("_")
            zoom = int(parts[0])
            candidate_row_idx = int(parts[1])
            candidate_col_idx = int(parts[2])
            
    
        except Exception as e:
            logging.warning(f"Failed to parse candidate key {candidate_key}: {e}")
            continue
        debug_dir = Path("debug_tiles")
        debug_dir.mkdir(exist_ok=True)

        tiles = get_surrounding_tiles_sharded(candidate_key, zoom)
        composite_img = compose_tiles_ordered_sharded(
            tiles, (1024, 1024), (candidate_row_idx, candidate_col_idx)
        )

        display_img = F.rotate(
            composite_img, [0, 90, 180, 270][best_rotation_index]
        )




        candidate_img_tensor = tfm.ToTensor()(composite_img)
        candidate_img_tensor = tfm.Resize(
            (MATCHING_IMG_SIZE * 3, MATCHING_IMG_SIZE * 3), antialias=True
        )(candidate_img_tensor)
    

        candidate_img_tensor = candidate_img_tensor.to(DEVICE)

      

        progress(
            0.5 + len(candidate_infos) / num_results * 0.4,
            desc=f"Running matching for candidate {candidate_num}",
        )

        best_inliers, best_footprint = run_matching(
            query_tensor, candidate_img_tensor, base_footprint
        )


        if best_inliers > -1:
            candidate_infos.append(
                {
                    "candidate_num": candidate_num,
                    "filename": Path(candidate_path_str).name,
                    "inliers": best_inliers,
                    "display_image": display_img,
                    "footprint": best_footprint,
                }
            )

        if best_inliers > global_best_inliers:
            global_best_inliers = best_inliers
            global_best_footprint = best_footprint
            global_candidate_num = candidate_num
            global_filename = Path(candidate_path_str).name
            global_best_display_image = display_img

    progress(0.9, desc="Finalizing results")

    folium_map_html = create_map(
        global_best_footprint,
        global_candidate_num,
        global_filename,
        global_best_inliers,
        query_footprint=query_footprint,
    )
    progress(1, desc="Done")

    yield folium_map_html, None
    yield folium_map_html, global_best_display_image

# --- Gradio app setup ---

if __name__ == "__main__":
    example_list = []
    google_examples = []
    queries_folder = Path("./data/queries")
    if queries_folder.exists() and queries_folder.is_dir():
        image_extensions = ["*.jpg", "*.jpeg", "*.png"]
        image_files = []
        for ext in image_extensions:
            image_files.extend(queries_folder.glob(ext))

        if image_files:
            num_examples = min(10, len(image_files))
            random_examples = random.sample(image_files, num_examples)
            example_list = [
                [str(p), str(util_matching.get_footprint_from_path(p))]
                for p in random_examples
            ]
            logging.info(
                f"Loaded {len(example_list)} examples for Gradio with footprints."
            )
        else:
            logging.warning(
                f"No images found in the examples folder: {queries_folder}"
            )
    else:
        logging.warning(f"Examples folder not found: {queries_folder}")

    google_folder = Path("./data/google_maps_queries")
    if google_folder.exists() and google_folder.is_dir():
        image_extensions = ["*.jpg", "*.jpeg", "*.png"]
        google_files = []
        for ext in image_extensions:
            google_files.extend(google_folder.glob(ext))
        if google_files:
            num_google = min(10, len(google_files))
            google_examples = [
                [str(p), str(p.stem).split("_")[0]]  # Empty footprint for Google Maps
                for p in random.sample(google_files, num_google)
            ]

    model_description = """
    ## Model Details
    This is a public API for inference of the EarthLoc2 model, which implements the amazing works of:
    - EarthLoc (https://earthloc-and-earthmatch.github.io/)
    - EarthMatch (https://earthloc-and-earthmatch.github.io/)
    - AstroLoc (https://astro-loc.github.io/)

    ### Architecture
    - DINOv2 base with SALAD aggregator out dim = 3072
    - FAISS index ~ 8gb, indexes 161496 * 4 images (4 rotated versions) from 2021

    ### Training
    - Trained on the original EarthLoc dataset (zooms 9,10,11), in range -60,60 latitude, polar regions not supported
    - Training included additional queries which were not part of the test/val sets 
    - 5000 iterations with a batch size of 96

    ### Performance
    - Achieves R@10 = 90.6 on the original EarthLoc test and val sets (when retrieving against whole db as is)
    - Overall performance is around 10% worse than AstroLoc (https://9d214e4bc329a5c3f9.gradio.live/)
    - Works well on satelite images between 1000 sq.km and 50000 sq.km, smaller or higher areas will not produce good results. 
    ### Matching
    - Uses the Xfeat_steerers matcher with 2048 maximal number of keypoints, we recommend Master with 2048 if you have access to GPU. 
    """

    theme = gr.themes.Soft(
        primary_hue=gr.themes.colors.blue,
        font=gr.themes.GoogleFont("Inter"),
    ).set(
        button_primary_background_fill="*primary_900",
        button_primary_background_fill_hover="*primary_700",
    )

    with gr.Blocks(theme=theme) as demo:
        gr.Markdown("# Aerial Photography Locator ")

        with gr.Row():
            with gr.Column(scale=2):
                image_input = gr.Image(
                    type="pil",
                    label="Aerial Photos of Earth",
                    height=400,
                )
                hidden_footprint_text = gr.Textbox(
                    visible=False, label="Query Footprint"
                )

                slider = gr.Slider(
                    minimum=1,
                    maximum=20,
                    step=1,
                    value=10,
                    label="Number of Candidates to Process",
                    info=(
                        "The higher this number to more likely the model is to find a match, "
                        "however it takes longer to find it. Expect around 5 second more compute per candidate."
                    ),
                )

                submit_btn = gr.Button("Localize Image", variant="primary")

                with gr.Row():
                    with gr.Column():
                        if example_list:
                            gr.Markdown("### ISS Example Queries")
                            gr.Examples(
                                examples=example_list,
                                inputs=[image_input, hidden_footprint_text],
                                examples_per_page=5,
                                cache_examples=False,
                            )
                    
                    with gr.Column():
                        if google_examples:
                            gr.Markdown("### Google Maps Example Queries")
                            gr.Examples(
                                examples=google_examples,
                                inputs=[image_input, hidden_footprint_text],
                                examples_per_page=5,
                                cache_examples=False,
                            )

            with gr.Column(scale=2):
                map_output = gr.HTML(label="Final Footprint Map")
                image_output = gr.Image(
                    type="pil",
                    label="Best Matching Candidate",
                    height=400,
                    show_download_button=True,
                )
                gr.Markdown(model_description)

        submit_btn.click(
            fn=search_and_retrieve,
            inputs=[image_input, hidden_footprint_text, slider],
            outputs=[map_output, image_output],
        )

    demo.launch(share=True)