"""Getting started with GeoPandas - TechnicalGIS tutorial.

Reads Natural Earth countries and populated places, reprojects to an
equal-area CRS, performs a spatial join and a dissolve, and draws
choropleth maps.  Run with:  python geopandas_basics.py
"""
import urllib.request
from pathlib import Path

import geopandas as gpd
import matplotlib.pyplot as plt
import pandas as pd

HERE = Path(__file__).parent
DATA = HERE / "data"
OUT = HERE / "output"
DATA.mkdir(exist_ok=True)
OUT.mkdir(exist_ok=True)

pd.set_option("display.width", 120)
pd.set_option("display.max_columns", 8)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

NE = "https://naciscdn.org/naturalearth/50m/cultural/"


def fetch(name: str) -> Path:
    """Download a Natural Earth zip once and cache it in ./data."""
    dest = DATA / name
    if not dest.exists():
        req = urllib.request.Request(NE + name, headers={"User-Agent": "technicalgis-tutorial"})
        with urllib.request.urlopen(req) as r:
            dest.write_bytes(r.read())
    return dest


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: read the data -------------------------------------------------
countries = gpd.read_file(fetch("ne_50m_admin_0_countries.zip"))
places = gpd.read_file(fetch("ne_50m_populated_places.zip"))

section("countries")
print(type(countries))
print(countries.shape)
print(countries.crs)
print(countries.geom_type.value_counts())

countries = countries[["NAME", "ISO_A3_EH", "CONTINENT", "POP_EST", "GDP_MD", "geometry"]]
countries = countries.rename(columns={"NAME": "name", "ISO_A3_EH": "iso3", "CONTINENT": "continent",
                                      "POP_EST": "pop_est", "GDP_MD": "gdp_musd"})
section("countries.head()")
print(countries.head())

# --- Step 2: reproject and calculate areas ---------------------------------
countries_ea = countries.to_crs("ESRI:54009")  # World Mollweide (equal area)
countries_ea["area_km2"] = countries_ea.area / 1e6
countries_ea["density"] = countries_ea["pop_est"] / countries_ea["area_km2"]

section("largest countries by area")
print(countries_ea.nlargest(5, "area_km2")[["name", "area_km2", "pop_est", "density"]].round(1).to_string(index=False))

fig, axes = plt.subplots(1, 2, figsize=(11, 3.4))
countries.plot(ax=axes[0], color="#9ec5b8", edgecolor="white", linewidth=0.3)
axes[0].set_title("EPSG:4326 (geographic lat/lon)")
countries_ea.plot(ax=axes[1], color="#9ec5b8", edgecolor="white", linewidth=0.3)
axes[1].set_title("ESRI:54009 (World Mollweide, equal area)")
for ax in axes:
    ax.set_axis_off()
fig.savefig(OUT / "01-projections.png")
plt.close(fig)

# --- Step 3: choropleth -------------------------------------------------------
world = countries_ea[countries_ea["continent"] != "Antarctica"]
fig, ax = plt.subplots(figsize=(10, 5))
world.plot(column="density", scheme="Quantiles", k=6, cmap="YlGnBu", legend=True,
           edgecolor="white", linewidth=0.25, ax=ax,
           legend_kwds={"title": "People per km²", "loc": "lower left", "fmt": "{:.0f}", "fontsize": 7})
ax.set_title("Population density by country (quantile classes)")
ax.set_axis_off()
fig.savefig(OUT / "02-density-choropleth.png")
plt.close(fig)

# --- Step 4: spatial join -----------------------------------------------------
places = places[["NAME", "POP_MAX", "geometry"]].rename(columns={"NAME": "city", "POP_MAX": "city_pop"})
places_ea = places.to_crs(countries_ea.crs)
joined = gpd.sjoin(places_ea, countries_ea[["name", "continent", "geometry"]], how="inner", predicate="within")

section("joined.head()")
print(joined[["city", "city_pop", "name", "continent"]].head().to_string())
section("unmatched places")
print(len(places_ea) - len(joined), "of", len(places_ea), "places fall outside every country polygon")

# Coastal cities often sit just offshore of generalised 1:50m coastlines.
unmatched = places_ea[~places_ea.index.isin(joined.index)]
nearest = gpd.sjoin_nearest(unmatched, countries_ea[["name", "continent", "geometry"]],
                            max_distance=25_000, distance_col="dist_m")
section("recovered with sjoin_nearest")
print(len(nearest), "recovered within 25 km")
print(nearest.nlargest(5, "city_pop")[["city", "city_pop", "name", "dist_m"]].round(0).to_string())
joined = pd.concat([joined, nearest.drop(columns="dist_m")])

big = joined[joined["city_pop"] >= 1_000_000]
per_country = big.groupby("name").agg(cities_1m=("city", "count"), city_pop=("city_pop", "sum"))
section("countries with the most 1M+ cities")
print(per_country.sort_values("cities_1m", ascending=False).head(8))

xmin, xmax, ymin, ymax = -1.2e6, 3.5e6, 3.9e6, 7.6e6
europe = countries_ea[countries_ea["continent"] == "Europe"].cx[xmin:xmax, ymin:ymax]
eu_cities = joined[joined["continent"] == "Europe"].cx[xmin:xmax, ymin:ymax]
fig, ax = plt.subplots(figsize=(7, 6))
countries_ea.plot(ax=ax, color="#f3f4f4", edgecolor="#c9cfcd", linewidth=0.3)
europe.plot(ax=ax, color="#e3ebe7", edgecolor="#8a9a94", linewidth=0.4)
eu_cities.plot(ax=ax, markersize=eu_cities["city_pop"] / 120_000, color="#0f766e", alpha=0.7,
               edgecolor="white", linewidth=0.3)
for _, r in eu_cities.nlargest(10, "city_pop").iterrows():
    ax.annotate(r["city"], (r.geometry.x, r.geometry.y), xytext=(4, 3), textcoords="offset points", fontsize=7)
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
ax.set_title("Natural Earth populated places joined to European countries\n(marker size = population)")
ax.set_axis_off()
fig.savefig(OUT / "03-spatial-join-europe.png")
plt.close(fig)

# --- Step 5: dissolve ---------------------------------------------------------
continents = countries_ea[countries_ea["continent"] != "Seven seas (open ocean)"].dissolve(
    by="continent", aggfunc={"pop_est": "sum", "gdp_musd": "sum", "area_km2": "sum"})
continents["density"] = continents["pop_est"] / continents["area_km2"]
section("continents")
with pd.option_context("display.float_format", "{:,.1f}".format):
    print(continents.drop(columns="geometry").sort_values("pop_est", ascending=False))

fig, ax = plt.subplots(figsize=(10, 5))
continents.reset_index().plot(column="continent", cmap="Set2", edgecolor="white", linewidth=0.6, ax=ax, legend=True,
                              legend_kwds={"loc": "lower left", "fontsize": 7, "frameon": False})
ax.set_title("Countries dissolved into continents")
ax.set_axis_off()
fig.savefig(OUT / "04-dissolve-continents.png")
plt.close(fig)

# --- Step 6: save -------------------------------------------------------------
gpkg = OUT / "natural_earth_analysis.gpkg"
countries_ea.to_file(gpkg, layer="countries", driver="GPKG")
continents.to_file(gpkg, layer="continents", driver="GPKG")
joined.drop(columns="index_right").to_file(gpkg, layer="places_joined", driver="GPKG")
section("saved")
print(gpkg.name, round(gpkg.stat().st_size / 1e6, 2), "MB")

# HTML version of the joined table so it can be screenshotted like a notebook cell.
(OUT / "joined_head.html").write_text(
    joined[["city", "city_pop", "name", "continent", "geometry"]].head(8).to_html(), encoding="utf8")