"""Core hedonic regression models for Airbnb pricing analysis."""

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

import numpy as np
import pandas as pd
import statsmodels.api as sm
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Feature group definitions for the 9-model nested sequence
# ---------------------------------------------------------------------------

STRUCTURAL_COLS = [
    "bedrooms", "bathrooms", "guests_count", "is_entire_home",
    "is_private_room", "log_accommodates", "log_bedrooms", "log_bathrooms",
    "bathrooms_per_guest", "bedrooms_per_guest", "large_capacity_listing",
    "accommodates_sq", "bedrooms_sq", "bathrooms_sq",
]

AMENITY_COLS = [
    "amenities_count", "amenity_richness_index", "luxury_index",
    "business_travel_index", "family_comfort_index", "leisure_index",
    "remote_work_index", "pets_allowed",
]

HOST_COLS = [
    "is_superhost", "superhost_flag", "host_professional_dummy",
    "multi_listing_host", "superhost_x_entire_home", "guest_favorite_flag",
    "is_guest_favorite",
]

REVIEW_COLS = [
    "rating", "num_reviews", "log_reviews", "review_quality",
    "reputation_score", "cleanliness_premium", "location_rating_premium",
    "has_reviews", "quality_score",
]

LOCATION_COLS = [
    "distance_to_center", "lat", "long",
]

NEIGHBOURHOOD_COLS = [
    "spatial_cluster",
]

NLP_COLS = [
    "title_length", "title_word_count", "luxury_language_score",
    "cozy_language_score", "professional_tone_score", "tourism_language_score",
    "text_embed_0", "text_embed_1", "text_embed_2", "text_embed_3",
]

SPATIAL_CONTROL_COLS = [
    "local_density_500m", "local_density_1km",
    "local_avg_price_500m", "local_avg_price_1km",
]

MODEL_SEQUENCE = {
    "Model 1 - Structural":      STRUCTURAL_COLS,
    "Model 2 - + Amenities":     STRUCTURAL_COLS + AMENITY_COLS,
    "Model 3 - + Host":          STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS,
    "Model 4 - + Reviews":       STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS,
    "Model 5 - + Location":      STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS + LOCATION_COLS,
    "Model 6 - + Neighbourhood": STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS + LOCATION_COLS + NEIGHBOURHOOD_COLS,
    "Model 7 - + NLP":           STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS + LOCATION_COLS + NEIGHBOURHOOD_COLS + NLP_COLS,
    "Model 8 - + Spatial":       STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS + LOCATION_COLS + NEIGHBOURHOOD_COLS + NLP_COLS + SPATIAL_CONTROL_COLS,
    "Model 9 - Full Hedonic":    STRUCTURAL_COLS + AMENITY_COLS + HOST_COLS + REVIEW_COLS + LOCATION_COLS + NEIGHBOURHOOD_COLS + NLP_COLS + SPATIAL_CONTROL_COLS,
}


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


def _prepare_data(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
) -> Tuple[pd.Series, pd.DataFrame]:
    """Drop rows with NaN in y or X; add constant; return (y, X)."""
    cols = _available_cols(df, x_cols)
    if not cols:
        raise ValueError("No valid feature columns found in the DataFrame.")
    subset = df[[y_col] + cols].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))
    logger.info("Prepared data: n=%d, k=%d (excl. constant)", len(y), len(cols))
    return y, X


# ---------------------------------------------------------------------------
# OLS estimators
# ---------------------------------------------------------------------------

def estimate_ols(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    robust: str = "HC3",
) -> sm.regression.linear_model.RegressionResultsWrapper:
    """OLS with heteroskedasticity-robust standard errors."""
    y, X = _prepare_data(df, y_col, x_cols)
    model = sm.OLS(y, X)
    results = model.fit(cov_type=robust)
    logger.info(
        "OLS (%s): R2=%.4f, Adj-R2=%.4f, n=%d",
        robust, results.rsquared, results.rsquared_adj, int(results.nobs),
    )
    return results


def estimate_ols_clustered(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    cluster_col: str,
) -> sm.regression.linear_model.RegressionResultsWrapper:
    """OLS with cluster-robust standard errors."""
    cols = _available_cols(df, x_cols)
    subset = df[[y_col] + cols + [cluster_col]].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))
    groups = subset[cluster_col]
    model = sm.OLS(y, X)
    results = model.fit(cov_type="cluster", cov_kwds={"groups": groups})
    logger.info(
        "OLS (clustered on %s): R2=%.4f, n=%d, clusters=%d",
        cluster_col, results.rsquared, int(results.nobs), groups.nunique(),
    )
    return results


# ---------------------------------------------------------------------------
# Nested model sequence
# ---------------------------------------------------------------------------

def build_model_sequence(
    df: pd.DataFrame,
    y_col: str = "log_price",
    robust: str = "HC3",
) -> Dict[str, sm.regression.linear_model.RegressionResultsWrapper]:
    """Estimate the 9-model nested hedonic sequence."""
    results: Dict[str, sm.regression.linear_model.RegressionResultsWrapper] = {}
    for name, xcols in MODEL_SEQUENCE.items():
        avail = _available_cols(df, xcols)
        if not avail:
            logger.warning("Skipping %s: no available features.", name)
            continue
        try:
            res = estimate_ols(df, y_col, avail, robust=robust)
            results[name] = res
            logger.info("Estimated %s  (k=%d)", name, len(avail))
        except Exception as exc:
            logger.error("Failed to estimate %s: %s", name, exc)
    return results


# ---------------------------------------------------------------------------
# Model comparison & diagnostics
# ---------------------------------------------------------------------------

def compare_models(
    results_dict: Dict[str, sm.regression.linear_model.RegressionResultsWrapper],
) -> pd.DataFrame:
    """Extract R2, adj-R2, RMSE, MAE, AIC, BIC for each model."""
    rows = []
    for name, res in results_dict.items():
        resid = res.resid
        n = int(res.nobs)
        rmse = np.sqrt(np.mean(resid ** 2))
        mae = np.mean(np.abs(resid))
        rows.append({
            "model": name,
            "n_obs": n,
            "n_params": int(res.df_model) + 1,
            "R2": round(res.rsquared, 6),
            "Adj_R2": round(res.rsquared_adj, 6),
            "RMSE": round(rmse, 6),
            "MAE": round(mae, 6),
            "AIC": round(res.aic, 2),
            "BIC": round(res.bic, 2),
        })
    comp = pd.DataFrame(rows)
    logger.info("Model comparison table built (%d models).", len(comp))
    return comp


def compute_price_premiums(
    results: sm.regression.linear_model.RegressionResultsWrapper,
) -> pd.DataFrame:
    """Convert log-linear coefficients to percentage price premiums: 100*(exp(b)-1)."""
    params = results.params.drop("const", errors="ignore")
    se = results.bse.drop("const", errors="ignore")
    pvals = results.pvalues.drop("const", errors="ignore")
    ci = results.conf_int().drop("const", errors="ignore")

    premium = 100.0 * (np.exp(params) - 1)
    premium_lo = 100.0 * (np.exp(ci[0]) - 1)
    premium_hi = 100.0 * (np.exp(ci[1]) - 1)

    df_out = pd.DataFrame({
        "coefficient": params.values,
        "std_error": se.values,
        "p_value": pvals.values,
        "premium_pct": premium.values,
        "premium_lo_95": premium_lo.values,
        "premium_hi_95": premium_hi.values,
    }, index=params.index)
    df_out = df_out.sort_values("premium_pct", ascending=False)
    logger.info("Computed price premiums for %d variables.", len(df_out))
    return df_out


# ---------------------------------------------------------------------------
# Cross-validation
# ---------------------------------------------------------------------------

def cross_validate_ols(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    k: int = 5,
    random_state: int = 42,
) -> Dict[str, float]:
    """K-fold cross-validation returning out-of-sample RMSE and MAE."""
    y, X = _prepare_data(df, y_col, x_cols)
    kf = KFold(n_splits=k, shuffle=True, random_state=random_state)
    rmses, maes = [], []
    for fold, (train_idx, test_idx) in enumerate(kf.split(X), 1):
        X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
        y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
        model = sm.OLS(y_train, X_train).fit()
        y_pred = model.predict(X_test)
        rmses.append(np.sqrt(mean_squared_error(y_test, y_pred)))
        maes.append(mean_absolute_error(y_test, y_pred))
        logger.debug("Fold %d: RMSE=%.4f, MAE=%.4f", fold, rmses[-1], maes[-1])

    result = {
        "mean_rmse": float(np.mean(rmses)),
        "std_rmse": float(np.std(rmses)),
        "mean_mae": float(np.mean(maes)),
        "std_mae": float(np.std(maes)),
        "k": k,
    }
    logger.info(
        "%d-fold CV: RMSE=%.4f (+/-%.4f), MAE=%.4f (+/-%.4f)",
        k, result["mean_rmse"], result["std_rmse"],
        result["mean_mae"], result["std_mae"],
    )
    return result


# ---------------------------------------------------------------------------
# Publication-quality regression table
# ---------------------------------------------------------------------------

def format_regression_table(
    results_dict: Dict[str, sm.regression.linear_model.RegressionResultsWrapper],
    stars: bool = True,
) -> pd.DataFrame:
    """Create a publication-ready regression comparison table (coef + SE rows)."""
    all_vars: List[str] = []
    for res in results_dict.values():
        for v in res.params.index:
            if v not in all_vars:
                all_vars.append(v)

    rows: List[Dict[str, str]] = []
    for var in all_vars:
        coef_row: Dict[str, str] = {"variable": var}
        se_row: Dict[str, str] = {"variable": ""}
        for mname, res in results_dict.items():
            if var in res.params.index:
                b = res.params[var]
                s = res.bse[var]
                p = res.pvalues[var]
                star_str = ""
                if stars:
                    if p < 0.001:
                        star_str = "***"
                    elif p < 0.01:
                        star_str = "**"
                    elif p < 0.05:
                        star_str = "*"
                coef_row[mname] = f"{b:.4f}{star_str}"
                se_row[mname] = f"({s:.4f})"
            else:
                coef_row[mname] = ""
                se_row[mname] = ""
        rows.append(coef_row)
        rows.append(se_row)

    # Footer rows
    for stat_name, extractor in [
        ("N", lambda r: f"{int(r.nobs)}"),
        ("R2", lambda r: f"{r.rsquared:.4f}"),
        ("Adj. R2", lambda r: f"{r.rsquared_adj:.4f}"),
        ("AIC", lambda r: f"{r.aic:.1f}"),
        ("BIC", lambda r: f"{r.bic:.1f}"),
    ]:
        row: Dict[str, str] = {"variable": stat_name}
        for mname, res in results_dict.items():
            row[mname] = extractor(res)
        rows.append(row)

    table = pd.DataFrame(rows)
    logger.info("Formatted regression table: %d rows x %d models.", len(table), len(results_dict))
    return table