import os
import io
import base64
import numpy as np
import pandas as pd
import rasterio
import pyproj
import matplotlib.pyplot as plt
from matplotlib import cm, colors
from geopy.geocoders import Nominatim
from config import OUTPUT_DIR
import plotly.graph_objects as go
from PIL import Image

CELL_SIZE_M = 100  # meters


def make_map(city, show_grid, show_georef):
    city = city.strip()
    if not city:
        return go.Figure().add_annotation(text="Please enter a city")

    # --- Geocode city ---
    geolocator = Nominatim(
        user_agent="histOSM_gradioAPP (maria.u.kuznetsova@gmail.com)",
        timeout=10
    )
    loc = geolocator.geocode(city)
    if loc is None:
        return go.Figure().add_annotation(text=f"Could not find '{city}'")

    lat_center, lon_center = loc.latitude, loc.longitude
    fig = go.Figure()

    # --- Add city marker ---
    fig.add_trace(go.Scattermapbox(
        lat=[lat_center],
        lon=[lon_center],
        mode="markers+text",
        text=["City center"],
        textposition="top right",
        marker=dict(size=12, color="blue")
    ))

    # --- Load raster ---
    raster_path = os.path.join(OUTPUT_DIR, "georeferenced.tif")
    if not os.path.exists(raster_path):
        return go.Figure().add_annotation(text="Georeferenced raster not found")

    with rasterio.open(raster_path) as src:
        bounds = src.bounds
        crs = src.crs
    xmin, ymin, xmax, ymax = bounds
    transformer = pyproj.Transformer.from_crs(crs, "EPSG:4326", always_xy=True)
    lon0, lat0 = transformer.transform(xmin, ymin)
    lon1, lat1 = transformer.transform(xmax, ymax)
    lat_min, lat_max = sorted([lat0, lat1])
    lon_min, lon_max = sorted([lon0, lon1])

    # --- Optional raster overlay ---
    if show_georef:
        with rasterio.open(raster_path) as src:
            arr = src.read(out_dtype="uint8")
            if arr.shape[0] >= 3:
                img = arr[:3].transpose(1, 2, 0)
            else:
                img = arr[0]

        img = np.clip(img, 0, 255)
        image = Image.fromarray(img)
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        encoded = base64.b64encode(buffer.getvalue()).decode()

        fig.update_layout(mapbox_layers=[
            dict(
                sourcetype="image",
                source="data:image/png;base64," + encoded,
                coordinates=[
                    [lon_min, lat_max],
                    [lon_max, lat_max],
                    [lon_max, lat_min],
                    [lon_min, lat_min]
                ],
                opacity=0.65
            )
        ])

        # Add raster center marker
        cx, cy = (xmin + xmax) / 2, (ymin + ymax) / 2
        clon, clat = transformer.transform(cx, cy)
        fig.add_trace(go.Scattermapbox(
            lat=[clat],
            lon=[clon],
            mode="markers+text",
            text=["Raster center"],
            textposition="bottom right",
            marker=dict(size=10, color="red")
        ))

    # --- Optional grid overlay ---
    if show_grid:
        _add_grid_overlay_plotly(fig, transformer)

    # --- Layout ---
    fig.update_layout(
        mapbox_style="open-street-map",
        mapbox_zoom=12,
        mapbox_center={"lat": lat_center, "lon": lon_center},
        margin={"l": 0, "r": 0, "t": 0, "b": 0},
        showlegend=False
    )
    return fig


def _add_grid_overlay_plotly(fig, transformer):
    """Add grid overlay as filled polygons on the Plotly map."""
    grid_values = []

    for fname in os.listdir(OUTPUT_DIR):
        if fname.startswith("street_matches_tile") and fname.endswith(".csv"):
            df = pd.read_csv(os.path.join(OUTPUT_DIR, fname))
            if df.empty:
                continue

            tile_xmin, tile_xmax = df['x'].min(), df['x'].max()
            tile_ymin, tile_ymax = df['y'].min(), df['y'].max()
            n_cols = int(np.ceil((tile_xmax - tile_xmin) / CELL_SIZE_M))
            n_rows = int(np.ceil((tile_ymax - tile_ymin) / CELL_SIZE_M))

            grid = np.zeros((n_rows, n_cols))
            counts = np.zeros((n_rows, n_cols))

            for _, row in df.iterrows():
                col = int((row['x'] - tile_xmin) // CELL_SIZE_M)
                row_idx = int((tile_ymax - row['y']) // CELL_SIZE_M)
                if 0 <= col < n_cols and 0 <= row_idx < n_rows:
                    grid[row_idx, col] += row['osm_match_score']
                    counts[row_idx, col] += 1

            mask = counts > 0
            grid[mask] /= counts[mask]
            grid_values.append((grid, tile_xmin, tile_ymin, n_rows, n_cols))

    if not grid_values:
        return

    all_scores = np.concatenate([g[0].flatten() for g in grid_values])
    min_val, max_val = all_scores.min(), all_scores.max()
    if min_val == max_val:
        max_val += 1e-6

    cmap = plt.get_cmap("Reds")

    for grid, tile_xmin, tile_ymin, n_rows, n_cols in grid_values:
        for r in range(n_rows):
            for c in range(n_cols):
                val = grid[r, c]
                if val <= 0:
                    continue
                norm_val = (val - min_val) / (max_val - min_val)
                color = colors.to_hex(cmap(norm_val))
                x0 = tile_xmin + c * CELL_SIZE_M
                y0 = tile_ymin + r * CELL_SIZE_M
                x1 = x0 + CELL_SIZE_M
                y1 = y0 + CELL_SIZE_M
                
                lon0, lat0 = transformer.transform(x0, y0)
                lon1, lat1 = transformer.transform(x1, y1)
                lons = [lon0, lon1, lon1, lon0, lon0]
                lats = [lat0, lat0, lat1, lat1, lat0]
                fig.add_trace(go.Scattermapbox(
                    lon=lons,
                    lat=lats,
                    mode="lines",
                    fill="toself",
                    fillcolor=color,
                    line=dict(width=0),
                    hoverinfo="text",
                    text=f"{val:.2f}"
                ))