"""Spatial econometric models for hedonic pricing analysis."""

import logging
import warnings
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy.spatial import cKDTree

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Try importing spatial libraries (optional dependencies)
# ---------------------------------------------------------------------------

_HAS_LIBPYSAL = False
_HAS_SPREG = False
_HAS_MGWR = False

try:
    import libpysal
    from libpysal.weights import KNN as _KNN
    _HAS_LIBPYSAL = True
except ImportError:
    logger.info("libpysal not available; using manual KNN weights.")

try:
    import spreg
    _HAS_SPREG = True
except ImportError:
    logger.info("spreg not available; spatial models will use 2SLS fallback.")

try:
    import mgwr
    from mgwr.gwr import GWR as _GWR
    from mgwr.sel_bw import Sel_BW as _SelBW
    _HAS_MGWR = True
except ImportError:
    logger.info("mgwr not available; GWR will not be estimated.")


# ---------------------------------------------------------------------------
# Helper: prepare data
# ---------------------------------------------------------------------------

def _available_cols(df: pd.DataFrame, requested: List[str]) -> List[str]:
    """Return the subset of requested columns present in df."""
    present = [c for c in requested if c in df.columns]
    missing = set(requested) - set(present)
    if missing:
        logger.warning("Missing columns (skipped): %s", sorted(missing))
    return present


def _prepare_data(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
) -> Tuple[pd.Series, pd.DataFrame, pd.DataFrame]:
    """Drop NaN, add constant, return (y, X_with_const, clean_df)."""
    cols = _available_cols(df, x_cols)
    if not cols:
        raise ValueError("No valid feature columns found in the DataFrame.")
    lat_col = "lat" if "lat" in df.columns else "latitude"
    lon_col = "long" if "long" in df.columns else "longitude"
    need = [y_col, lat_col, lon_col] + cols
    need = list(dict.fromkeys(need))  # deduplicate, preserve order
    subset = df[need].dropna()
    if subset.empty:
        raise ValueError("No observations remaining after dropping NaN rows.")
    y = subset[y_col].astype(float)
    X = sm.add_constant(subset[cols].astype(float))
    return y, X, subset


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

class _ManualKNNWeights:
    """Minimal KNN spatial weights when libpysal is not available."""

    def __init__(self, coords: np.ndarray, k: int = 5):
        self.n = len(coords)
        self.k = k
        tree = cKDTree(coords)
        distances, indices = tree.query(coords, k=k + 1)
        # exclude self-neighbour (index 0 in results)
        self.neighbors: Dict[int, List[int]] = {}
        self.weights: Dict[int, List[float]] = {}
        for i in range(self.n):
            nbrs = indices[i, 1:].tolist()
            self.neighbors[i] = nbrs
            self.weights[i] = [1.0 / k] * k  # row-standardised

    @property
    def sparse(self):
        """Return a scipy sparse CSR row-standardised weights matrix."""
        from scipy import sparse as sp
        rows, cols_, vals = [], [], []
        for i in range(self.n):
            for j, w in zip(self.neighbors[i], self.weights[i]):
                rows.append(i)
                cols_.append(j)
                vals.append(w)
        W = sp.csr_matrix((vals, (rows, cols_)), shape=(self.n, self.n))
        return W

    def full(self) -> np.ndarray:
        """Dense weight matrix."""
        return self.sparse.toarray()


def build_weights_matrix(
    df: pd.DataFrame,
    k: int = 5,
    lat_col: str = "latitude",
    lon_col: str = "longitude",
) -> Union["_ManualKNNWeights", object]:
    """KNN-based spatial weights matrix."""
    if lat_col not in df.columns or lon_col not in df.columns:
        raise ValueError(f"DataFrame must contain '{lat_col}' and '{lon_col}' columns.")

    coords = df[[lat_col, lon_col]].dropna().values
    n = len(coords)
    if n < k + 1:
        raise ValueError(f"Need at least k+1={k+1} observations, got {n}.")

    if _HAS_LIBPYSAL:
        w = _KNN.from_array(coords, k=k)
        w.transform = "r"  # row-standardise
        logger.info("Built libpysal KNN weights: n=%d, k=%d.", n, k)
        return w

    w = _ManualKNNWeights(coords, k=k)
    logger.info("Built manual KNN weights: n=%d, k=%d.", n, k)
    return w


# ---------------------------------------------------------------------------
# Spatial autocorrelation test
# ---------------------------------------------------------------------------

def test_spatial_autocorrelation(
    df: pd.DataFrame,
    y_col: str,
    w: object,
    permutations: int = 999,
) -> Dict[str, float]:
    """Moran's I test for spatial autocorrelation."""
    y = df[y_col].dropna().values.astype(float)

    if _HAS_LIBPYSAL:
        from esda.moran import Moran
        mi = Moran(y, w, permutations=permutations)
        result = {
            "morans_I": float(mi.I),
            "expected_I": float(mi.EI),
            "z_score": float(mi.z_norm),
            "p_value": float(mi.p_norm),
            "p_sim": float(mi.p_sim) if hasattr(mi, "p_sim") else np.nan,
        }
        logger.info("Moran's I=%.4f, z=%.2f, p=%.4f", result["morans_I"], result["z_score"], result["p_value"])
        return result

    # Manual Moran's I
    n = len(y)
    y_bar = y.mean()
    z = y - y_bar

    if hasattr(w, "sparse"):
        W = w.sparse if hasattr(w.sparse, "toarray") else w.sparse
        if hasattr(W, "toarray"):
            W_dense = W.toarray()
        else:
            W_dense = np.array(W)
    elif hasattr(w, "full"):
        W_dense = w.full()
    else:
        raise ValueError("Weights object must have .sparse or .full() method.")

    numerator = float(z @ W_dense @ z)
    denominator = float(z @ z)
    S0 = float(W_dense.sum())

    I = (n / S0) * (numerator / denominator) if denominator != 0 else 0.0
    EI = -1.0 / (n - 1)

    # Variance under normality assumption
    S1 = float(0.5 * ((W_dense + W_dense.T) ** 2).sum())
    S2 = float(np.sum(np.array(W_dense.sum(axis=0)).flatten() + np.array(W_dense.sum(axis=1)).flatten()) ** 2)
    # Simplified z-score
    var_I = (n ** 2 * S1 - n * S2 + 3 * S0 ** 2) / ((n ** 2 - 1) * S0 ** 2) - EI ** 2
    var_I = max(var_I, 1e-12)
    z_score = (I - EI) / np.sqrt(var_I)

    from scipy import stats as sp_stats
    p_value = 2.0 * (1.0 - sp_stats.norm.cdf(abs(z_score)))

    result = {
        "morans_I": float(I),
        "expected_I": float(EI),
        "z_score": float(z_score),
        "p_value": float(p_value),
        "p_sim": np.nan,
    }
    logger.info("Manual Moran's I=%.4f, z=%.2f, p=%.4f", I, z_score, p_value)
    return result


# ---------------------------------------------------------------------------
# Spatial lag variable creation
# ---------------------------------------------------------------------------

def create_spatial_lags(
    df: pd.DataFrame,
    cols: List[str],
    w: object,
) -> pd.DataFrame:
    """Compute spatial lag (Wy) for specified columns."""
    result = df.copy()
    if hasattr(w, "sparse"):
        W_sp = w.sparse
        if hasattr(W_sp, "toarray"):
            W_dense = W_sp.toarray()
        else:
            W_dense = np.array(W_sp)
    elif hasattr(w, "full"):
        W_dense = w.full()
    else:
        raise ValueError("Weights object must have .sparse or .full() method.")

    avail = _available_cols(df, cols)
    for col in avail:
        vals = df[col].values.astype(float)
        lag = W_dense @ vals
        result[f"W_{col}"] = lag
        logger.debug("Created spatial lag W_%s.", col)

    logger.info("Created spatial lags for %d columns.", len(avail))
    return result


# ---------------------------------------------------------------------------
# Spatial econometric models
# ---------------------------------------------------------------------------

def estimate_spatial_lag(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    w: object,
) -> object:
    """Spatial lag model (SAR): y = rho*Wy + Xb + e."""
    y, X, clean = _prepare_data(df, y_col, x_cols)

    if _HAS_SPREG:
        _lat = "lat" if "lat" in clean.columns else "latitude"
        _lon = "long" if "long" in clean.columns else "longitude"
        coords = clean[[_lat, _lon]].values
        y_arr = y.values.reshape(-1, 1)
        X_arr = X.drop(columns="const", errors="ignore").values
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            model = spreg.ML_Lag(y_arr, X_arr, w=w, name_y=y_col, name_x=list(X.columns.drop("const", errors="ignore")))
        logger.info("Spatial Lag (spreg ML): rho=%.4f, pseudo-R2=%.4f", model.rho, model.pr2)
        return model

    # Fallback: 2SLS with Wy instrumented by WX, W2X
    return _spatial_lag_2sls(y, X, w, y_col, x_cols)


def _spatial_lag_2sls(
    y: pd.Series,
    X: pd.DataFrame,
    w: object,
    y_col: str,
    x_cols: List[str],
) -> object:
    """Manual 2SLS spatial lag estimation."""
    if hasattr(w, "sparse"):
        W_sp = w.sparse
        W_dense = W_sp.toarray() if hasattr(W_sp, "toarray") else np.array(W_sp)
    elif hasattr(w, "full"):
        W_dense = w.full()
    else:
        raise ValueError("Weights object must have .sparse or .full().")

    y_arr = y.values
    Wy = W_dense @ y_arr

    X_no_const = X.drop(columns="const", errors="ignore").values
    WX = W_dense @ X_no_const
    W2X = W_dense @ WX

    # Instruments: X, WX, W2X
    Z = np.column_stack([np.ones(len(y_arr)), X_no_const, WX, W2X])

    # Endogenous RHS: [const, X, Wy]
    endog_rhs = np.column_stack([np.ones(len(y_arr)), X_no_const, Wy])

    # 2SLS
    # Stage 1: project Wy onto Z
    Pz = Z @ np.linalg.lstsq(Z, Wy, rcond=None)[0]
    # Actually need: Wy_hat = Z @ inv(Z'Z) @ Z' @ Wy
    ZtZ_inv = np.linalg.pinv(Z.T @ Z)
    Wy_hat = Z @ ZtZ_inv @ Z.T @ Wy

    # Stage 2: regress y on [const, X, Wy_hat]
    RHS = np.column_stack([np.ones(len(y_arr)), X_no_const, Wy_hat])
    beta_2sls = np.linalg.lstsq(RHS, y_arr, rcond=None)[0]

    # Residuals using actual Wy
    RHS_actual = np.column_stack([np.ones(len(y_arr)), X_no_const, Wy])
    resid = y_arr - RHS_actual @ beta_2sls
    n, k = RHS.shape
    sigma2 = float(resid @ resid / (n - k))

    # Standard errors (2SLS)
    bread = np.linalg.pinv(RHS.T @ Z @ ZtZ_inv @ Z.T @ RHS)
    se = np.sqrt(np.diag(bread) * sigma2)

    x_names = ["const"] + list(X.columns.drop("const", errors="ignore")) + ["W_" + y_col]

    result = {
        "method": "2SLS_spatial_lag",
        "coefficients": dict(zip(x_names, beta_2sls)),
        "std_errors": dict(zip(x_names, se)),
        "rho": float(beta_2sls[-1]),
        "sigma2": sigma2,
        "n_obs": n,
        "n_params": k,
        "residuals": resid,
        "pseudo_R2": float(1.0 - np.var(resid) / np.var(y_arr)),
    }
    logger.info(
        "Spatial Lag (2SLS fallback): rho=%.4f, pseudo-R2=%.4f, n=%d",
        result["rho"], result["pseudo_R2"], n,
    )
    return result


def estimate_spatial_error(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    w: object,
) -> object:
    """Spatial error model (SEM): y = Xb + u, u = lambda*Wu + e."""
    y, X, clean = _prepare_data(df, y_col, x_cols)

    if _HAS_SPREG:
        y_arr = y.values.reshape(-1, 1)
        X_arr = X.drop(columns="const", errors="ignore").values
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            model = spreg.ML_Error(y_arr, X_arr, w=w, name_y=y_col, name_x=list(X.columns.drop("const", errors="ignore")))
        logger.info("Spatial Error (spreg ML): lambda=%.4f, pseudo-R2=%.4f", model.lam, model.pr2)
        return model

    # Fallback: Cochrane-Orcutt-style iterative GLS
    return _spatial_error_iterative(y, X, w, y_col)


def _spatial_error_iterative(
    y: pd.Series,
    X: pd.DataFrame,
    w: object,
    y_col: str,
    max_iter: int = 50,
    tol: float = 1e-6,
) -> Dict:
    """Iterative estimation of spatial error model via filtered GLS."""
    if hasattr(w, "sparse"):
        W_sp = w.sparse
        W_dense = W_sp.toarray() if hasattr(W_sp, "toarray") else np.array(W_sp)
    elif hasattr(w, "full"):
        W_dense = w.full()
    else:
        raise ValueError("Weights object must have .sparse or .full().")

    y_arr = y.values
    X_arr = X.values
    n = len(y_arr)
    I_n = np.eye(n)

    # Initial OLS
    beta = np.linalg.lstsq(X_arr, y_arr, rcond=None)[0]
    lam = 0.0

    for iteration in range(max_iter):
        resid = y_arr - X_arr @ beta
        Wu = W_dense @ resid

        # Estimate lambda by regressing residuals on W*residuals
        lam_new = float(np.dot(Wu, resid) / np.dot(Wu, Wu)) if np.dot(Wu, Wu) != 0 else 0.0
        lam_new = max(-0.99, min(0.99, lam_new))

        # Filtered variables
        A = I_n - lam_new * W_dense
        y_star = A @ y_arr
        X_star = A @ X_arr

        beta_new = np.linalg.lstsq(X_star, y_star, rcond=None)[0]

        if np.max(np.abs(beta_new - beta)) < tol and abs(lam_new - lam) < tol:
            beta = beta_new
            lam = lam_new
            logger.debug("Spatial error converged at iteration %d.", iteration + 1)
            break
        beta = beta_new
        lam = lam_new

    resid_final = y_arr - X_arr @ beta
    sigma2 = float(resid_final @ resid_final / (n - len(beta)))

    result = {
        "method": "iterative_spatial_error",
        "coefficients": dict(zip(X.columns, beta)),
        "lambda": lam,
        "sigma2": sigma2,
        "n_obs": n,
        "n_params": len(beta),
        "residuals": resid_final,
        "pseudo_R2": float(1.0 - np.var(resid_final) / np.var(y_arr)),
    }
    logger.info(
        "Spatial Error (iterative fallback): lambda=%.4f, pseudo-R2=%.4f, n=%d",
        lam, result["pseudo_R2"], n,
    )
    return result


def estimate_spatial_durbin(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    w: object,
) -> object:
    """Spatial Durbin model (SDM): y = rho*Wy + Xb + WX*theta + e."""
    y, X, clean = _prepare_data(df, y_col, x_cols)

    # Create spatial lags of X
    X_no_const = X.drop(columns="const", errors="ignore")
    if hasattr(w, "sparse"):
        W_sp = w.sparse
        W_dense = W_sp.toarray() if hasattr(W_sp, "toarray") else np.array(W_sp)
    elif hasattr(w, "full"):
        W_dense = w.full()
    else:
        raise ValueError("Weights object must have .sparse or .full().")

    WX = pd.DataFrame(
        W_dense @ X_no_const.values,
        columns=[f"W_{c}" for c in X_no_const.columns],
        index=X_no_const.index,
    )
    X_augmented = pd.concat([X_no_const, WX], axis=1)
    x_cols_augmented = list(X_augmented.columns)

    if _HAS_SPREG:
        y_arr = y.values.reshape(-1, 1)
        X_arr = X_augmented.values
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            model = spreg.ML_Lag(y_arr, X_arr, w=w, name_y=y_col, name_x=x_cols_augmented)
        logger.info("Spatial Durbin (spreg ML_Lag with WX): rho=%.4f, pseudo-R2=%.4f", model.rho, model.pr2)
        return model

    # Fallback: 2SLS with augmented X
    X_aug_with_const = sm.add_constant(X_augmented)
    return _spatial_lag_2sls(y, X_aug_with_const, w, y_col, x_cols_augmented)


# ---------------------------------------------------------------------------
# GWR
# ---------------------------------------------------------------------------

def run_gwr(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    lat_col: str = "latitude",
    lon_col: str = "longitude",
    kernel: str = "bisquare",
    fixed: bool = False,
) -> Optional[object]:
    """Geographically weighted regression (requires mgwr)."""
    if not _HAS_MGWR:
        logger.warning(
            "mgwr package not installed. Install with: pip install mgwr. "
            "GWR estimation is skipped."
        )
        return None

    cols = _available_cols(df, x_cols)
    need = [y_col, lat_col, lon_col] + cols
    subset = df[need].dropna()
    if len(subset) < 50:
        logger.warning("Too few observations (%d) for reliable GWR.", len(subset))

    coords = subset[[lat_col, lon_col]].values
    y_arr = subset[y_col].values.reshape(-1, 1)
    X_arr = subset[cols].values

    # Bandwidth selection
    selector = _SelBW(coords, y_arr, X_arr, kernel=kernel, fixed=fixed)
    bw = selector.search()
    logger.info("GWR optimal bandwidth: %.2f", bw)

    model = _GWR(coords, y_arr, X_arr, bw=bw, kernel=kernel, fixed=fixed)
    results = model.fit()
    logger.info(
        "GWR estimated: AICc=%.2f, R2=%.4f, n=%d",
        results.aicc, results.R2, len(y_arr),
    )
    return results


# ---------------------------------------------------------------------------
# Model comparison
# ---------------------------------------------------------------------------

def compare_spatial_models(results_dict: Dict[str, object]) -> pd.DataFrame:
    """Model comparison table for spatial models."""
    rows = []
    for name, res in results_dict.items():
        row: Dict[str, object] = {"model": name}

        if isinstance(res, dict):
            # Fallback dict results
            row["n_obs"] = res.get("n_obs", np.nan)
            row["n_params"] = res.get("n_params", np.nan)
            row["pseudo_R2"] = res.get("pseudo_R2", np.nan)
            row["sigma2"] = res.get("sigma2", np.nan)
            row["rho"] = res.get("rho", np.nan)
            row["lambda"] = res.get("lambda", np.nan)
            resid = res.get("residuals")
            if resid is not None:
                row["RMSE"] = float(np.sqrt(np.mean(resid ** 2)))
                row["MAE"] = float(np.mean(np.abs(resid)))
                row["AIC"] = _aic_from_resid(resid, res.get("n_params", 1))
            else:
                row["RMSE"] = np.nan
                row["MAE"] = np.nan
                row["AIC"] = np.nan
        elif _HAS_SPREG and hasattr(res, "pr2"):
            # spreg results
            row["n_obs"] = getattr(res, "n", np.nan)
            row["n_params"] = getattr(res, "k", np.nan)
            row["pseudo_R2"] = getattr(res, "pr2", np.nan)
            row["sigma2"] = getattr(res, "sig2", np.nan)
            row["rho"] = getattr(res, "rho", np.nan)
            row["lambda"] = getattr(res, "lam", np.nan)
            if hasattr(res, "u"):
                u = np.array(res.u).flatten()
                row["RMSE"] = float(np.sqrt(np.mean(u ** 2)))
                row["MAE"] = float(np.mean(np.abs(u)))
                row["AIC"] = _aic_from_resid(u, getattr(res, "k", 1))
            else:
                row["RMSE"] = np.nan
                row["MAE"] = np.nan
                row["AIC"] = np.nan
        else:
            # GWR or unknown
            row["n_obs"] = getattr(res, "n", np.nan)
            row["n_params"] = getattr(res, "k", np.nan)
            row["pseudo_R2"] = getattr(res, "R2", np.nan)
            row["sigma2"] = np.nan
            row["rho"] = np.nan
            row["lambda"] = np.nan
            row["RMSE"] = np.nan
            row["MAE"] = np.nan
            row["AIC"] = getattr(res, "aicc", np.nan)

        rows.append(row)

    comp = pd.DataFrame(rows)
    logger.info("Spatial model comparison: %d models.", len(comp))
    return comp


def _aic_from_resid(resid: np.ndarray, k: int) -> float:
    """Compute AIC from residuals assuming normal errors."""
    n = len(resid)
    if n <= 0:
        return np.nan
    sse = float(np.sum(resid ** 2))
    if sse <= 0:
        return np.nan
    log_lik = -0.5 * n * (np.log(2 * np.pi) + np.log(sse / n) + 1)
    return float(-2 * log_lik + 2 * k)