"""Geospatial feature engineering for hedonic pricing analysis."""

from __future__ import annotations

import logging
from typing import Any

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.neighbors import BallTree

logger = logging.getLogger(__name__)

EARTH_RADIUS_M = 6_371_000.0

# ---------------------------------------------------------------------------
# Default city centres (major Quebec cities)
# ---------------------------------------------------------------------------

DEFAULT_CITY_CENTERS: dict[str, tuple[float, float]] = {
    "montréal": (45.5017, -73.5673),
    "montreal": (45.5017, -73.5673),
    "québec": (46.8139, -71.2080),
    "quebec": (46.8139, -71.2080),
    "gatineau": (45.4765, -75.7013),
    "sherbrooke": (45.4042, -71.8929),
    "saguenay": (48.4280, -71.0686),
    "laval": (45.6066, -73.7124),
    "trois-rivières": (46.3432, -72.5432),
    "trois-rivieres": (46.3432, -72.5432),
    "lévis": (46.8032, -71.1780),
    "levis": (46.8032, -71.1780),
    "terrebonne": (45.6960, -73.6473),
    "saint-jean-sur-richelieu": (45.3073, -73.2628),
    "repentigny": (45.7402, -73.4598),
    "drummondville": (45.8838, -72.4843),
    "granby": (45.4000, -72.7333),
    "saint-hyacinthe": (45.6307, -72.9571),
    "rimouski": (48.4490, -68.5233),
    "victoriaville": (46.0500, -71.9667),
    "saint-georges": (46.1167, -70.6667),
    "matane": (48.8500, -67.5333),
    "mont-tremblant": (46.2103, -74.5847),
}

# ---------------------------------------------------------------------------
# Haversine helpers
# ---------------------------------------------------------------------------

def _to_radians(lat: np.ndarray, lon: np.ndarray) -> np.ndarray:
    """Convert lat/lon arrays to radians, shape (n, 2)."""
    return np.column_stack([np.radians(lat), np.radians(lon)])


def _haversine_vector(lat1: np.ndarray, lon1: np.ndarray, lat2: float, lon2: float) -> np.ndarray:
    """Vectorised haversine distance in metres."""
    lat1_r, lon1_r = np.radians(lat1), np.radians(lon1)
    lat2_r, lon2_r = np.radians(lat2), np.radians(lon2)
    dlat = lat2_r - lat1_r
    dlon = lon2_r - lon1_r
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1_r) * np.cos(lat2_r) * np.sin(dlon / 2) ** 2
    return 2 * EARTH_RADIUS_M * np.arcsin(np.sqrt(a))


# ---------------------------------------------------------------------------
# BallTree construction
# ---------------------------------------------------------------------------

def _build_balltree(lat: np.ndarray, lon: np.ndarray) -> BallTree:
    """Build a BallTree from latitude/longitude in radians using haversine metric."""
    coords_rad = _to_radians(lat, lon)
    return BallTree(coords_rad, metric="haversine")


# ---------------------------------------------------------------------------
# Distance to city centre
# ---------------------------------------------------------------------------

def compute_distance_to_center(
    df: pd.DataFrame,
    city_centers_dict: dict[str, tuple[float, float]] | None = None,
) -> pd.Series:
    """Compute haversine distance (m) from each listing to its city centre."""
    centers = city_centers_dict or DEFAULT_CITY_CENTERS

    distances = pd.Series(np.nan, index=df.index, dtype=float)
    city_col = df["city"].str.lower().str.strip()

    for city_key, (clat, clon) in centers.items():
        mask = city_col == city_key
        if mask.any():
            d = _haversine_vector(
                df.loc[mask, "lat"].values,
                df.loc[mask, "long"].values,
                clat,
                clon,
            )
            distances.loc[mask] = d

    # For cities not in the dictionary, use the median of their own coordinates as centre
    missing_mask = distances.isna()
    if missing_mask.any():
        for city_name, grp in df.loc[missing_mask].groupby(city_col[missing_mask]):
            clat = grp["lat"].median()
            clon = grp["long"].median()
            d = _haversine_vector(grp["lat"].values, grp["long"].values, clat, clon)
            distances.loc[grp.index] = d

    return distances


# ---------------------------------------------------------------------------
# Local density
# ---------------------------------------------------------------------------

def compute_local_density(df: pd.DataFrame, radius_m: float = 1000.0) -> pd.Series:
    """Count of other listings within radius_m metres using BallTree."""
    valid = df[["lat", "long"]].notna().all(axis=1)
    sub = df.loc[valid]
    lat = sub["lat"].values
    lon = sub["long"].values
    tree = _build_balltree(lat, lon)

    radius_rad = radius_m / EARTH_RADIUS_M
    coords_rad = _to_radians(lat, lon)
    counts = tree.query_radius(coords_rad, r=radius_rad, count_only=True)

    density = np.maximum(counts - 1, 0)
    logger.info("Local density (r=%dm): mean=%.1f, max=%d", radius_m, density.mean(), density.max())
    result = pd.Series(np.nan, index=df.index, name=f"local_density_{int(radius_m)}m")
    result.loc[valid] = density
    return result


# ---------------------------------------------------------------------------
# Local average price
# ---------------------------------------------------------------------------

def compute_local_avg_price(df: pd.DataFrame, radius_m: float = 1000.0) -> pd.Series:
    """Average price of other listings within radius_m metres."""
    coord_valid = df[["lat", "long"]].notna().all(axis=1)
    sub = df.loc[coord_valid]
    lat = sub["lat"].values
    lon = sub["long"].values
    prices = sub["price_numeric"].values.copy() if "price_numeric" in sub.columns else np.full(len(sub), np.nan)
    tree = _build_balltree(lat, lon)

    radius_rad = radius_m / EARTH_RADIUS_M
    coords_rad = _to_radians(lat, lon)
    indices_list = tree.query_radius(coords_rad, r=radius_rad)

    avg_prices = np.full(len(sub), np.nan)
    for i, neighbours in enumerate(indices_list):
        others = neighbours[neighbours != i]
        neighbour_prices = prices[others]
        valid = neighbour_prices[~np.isnan(neighbour_prices)]
        if len(valid) > 0:
            avg_prices[i] = valid.mean()

    result = pd.Series(np.nan, index=df.index, name=f"local_avg_price_{int(radius_m)}m")
    result.loc[coord_valid] = avg_prices
    logger.info("Local avg price (r=%dm): computed for %d listings", radius_m, np.sum(~np.isnan(avg_prices)))
    return result


# ---------------------------------------------------------------------------
# Spatial lag (KNN)
# ---------------------------------------------------------------------------

def compute_spatial_lag(df: pd.DataFrame, col: str, k: int = 5) -> pd.Series:
    """K-nearest-neighbour spatial lag of any numeric column."""
    coord_valid = df[["lat", "long"]].notna().all(axis=1)
    sub = df.loc[coord_valid]
    lat = sub["lat"].values
    lon = sub["long"].values
    values = sub[col].values.astype(float)
    tree = _build_balltree(lat, lon)

    coords_rad = _to_radians(lat, lon)
    dists, indices = tree.query(coords_rad, k=k + 1)

    lag = np.full(len(sub), np.nan)
    for i in range(len(sub)):
        nbr_idx = indices[i, 1:]
        nbr_vals = values[nbr_idx]
        valid = nbr_vals[~np.isnan(nbr_vals)]
        if len(valid) > 0:
            lag[i] = valid.mean()

    result = pd.Series(np.nan, index=df.index, name=f"spatial_lag_{col}_k{k}")
    result.loc[coord_valid] = lag
    logger.info("Spatial lag of '%s' (k=%d): computed for %d listings", col, k, np.sum(~np.isnan(lag)))
    return result


# ---------------------------------------------------------------------------
# Spatial clusters
# ---------------------------------------------------------------------------

def create_spatial_clusters(df: pd.DataFrame, n_clusters: int = 20) -> pd.Series:
    """K-means clustering on lat/long coordinates."""
    coords = df[["lat", "long"]].values
    valid_mask = ~np.isnan(coords).any(axis=1)

    labels = pd.Series(-1, index=df.index, dtype=int, name="spatial_cluster")
    if valid_mask.sum() < n_clusters:
        logger.warning("Too few valid coordinates (%d) for %d clusters", valid_mask.sum(), n_clusters)
        return labels

    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    kmeans.fit(coords[valid_mask])
    labels.iloc[np.where(valid_mask)[0]] = kmeans.labels_

    logger.info("Spatial clusters: %d clusters from %d valid points", n_clusters, valid_mask.sum())
    return labels


# ---------------------------------------------------------------------------
# Combined distance features
# ---------------------------------------------------------------------------

def compute_distance_features(
    df: pd.DataFrame,
    city_centers_dict: dict[str, tuple[float, float]] | None = None,
) -> pd.DataFrame:
    """Compute the standard suite of distance-based features."""
    out = df.copy()

    out["distance_to_center"] = compute_distance_to_center(out, city_centers_dict)
    out["log_distance_to_center"] = np.log1p(out["distance_to_center"])

    out["local_density_500m"] = compute_local_density(out, radius_m=500.0)
    out["local_density_1km"] = compute_local_density(out, radius_m=1000.0)

    out["local_avg_price_500m"] = compute_local_avg_price(out, radius_m=500.0)
    out["local_avg_price_1km"] = compute_local_avg_price(out, radius_m=1000.0)

    out["spatial_cluster"] = create_spatial_clusters(out)

    logger.info("Distance features complete: %d new columns", out.shape[1] - df.shape[1])
    return out


# ---------------------------------------------------------------------------
# Spatial weights matrix
# ---------------------------------------------------------------------------

def build_spatial_weights(df: pd.DataFrame, k: int = 5) -> np.ndarray:
    """Build a row-standardised KNN spatial weights matrix (n x n)."""
    lat = df["lat"].values
    lon = df["long"].values
    tree = _build_balltree(lat, lon)
    coords_rad = _to_radians(lat, lon)

    n = len(df)
    W = np.zeros((n, n), dtype=float)

    _, indices = tree.query(coords_rad, k=k + 1)
    for i in range(n):
        nbrs = indices[i, 1:]  # exclude self
        W[i, nbrs] = 1.0

    # Row-standardise
    row_sums = W.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1.0
    W = W / row_sums

    logger.info("Spatial weights matrix built: %dx%d, k=%d", n, n, k)
    return W


# ---------------------------------------------------------------------------
# Moran's I
# ---------------------------------------------------------------------------

def compute_morans_i(
    df: pd.DataFrame,
    col: str,
    w: np.ndarray | None = None,
    k: int = 5,
) -> dict[str, float]:
    """Compute Moran's I statistic for spatial autocorrelation."""
    y = df[col].values.astype(float)
    valid_mask = ~np.isnan(y)

    if valid_mask.sum() < 10:
        logger.warning("Too few valid observations (%d) for Moran's I", valid_mask.sum())
        return {"morans_i": np.nan, "expected_i": np.nan, "z_score": np.nan}

    if w is None:
        w = build_spatial_weights(df, k=k)

    # Restrict to valid observations
    idx = np.where(valid_mask)[0]
    y_sub = y[idx]
    w_sub = w[np.ix_(idx, idx)]

    # Re-row-standardise the sub-matrix
    row_sums = w_sub.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1.0
    w_sub = w_sub / row_sums

    n = len(y_sub)
    y_mean = y_sub.mean()
    z = y_sub - y_mean

    numerator = float(n * z @ w_sub @ z)
    denominator = float(w_sub.sum() * (z ** 2).sum())

    if denominator == 0:
        return {"morans_i": np.nan, "expected_i": np.nan, "z_score": np.nan}

    I = numerator / denominator
    expected_I = -1.0 / (n - 1)

    # Variance under normality assumption
    S0 = w_sub.sum()
    S1 = 0.5 * ((w_sub + w_sub.T) ** 2).sum()
    S2 = ((w_sub.sum(axis=1) + w_sub.sum(axis=0)) ** 2).sum()
    k2 = (n * (z ** 4).sum()) / ((z ** 2).sum() ** 2)

    A = n * ((n ** 2 - 3 * n + 3) * S1 - n * S2 + 3 * S0 ** 2)
    B = k2 * ((n ** 2 - n) * S1 - 2 * n * S2 + 6 * S0 ** 2)
    C = (n - 1) * (n - 2) * (n - 3) * S0 ** 2

    if C == 0:
        z_score = np.nan
    else:
        var_I = (A - B) / C - expected_I ** 2
        if var_I > 0:
            z_score = float((I - expected_I) / np.sqrt(var_I))
        else:
            z_score = np.nan

    result = {"morans_i": float(I), "expected_i": float(expected_I), "z_score": z_score}
    logger.info("Moran's I for '%s': I=%.4f, E[I]=%.4f, z=%.2f", col, I, expected_I, z_score if not np.isnan(z_score) else 0.0)
    return result