# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
05_descriptive_analysis.py
==========================
Generate descriptive statistics tables and figures for the Airbnb-rent
linkage analysis.

Tables  -> results/tables/   (.tex  +  .csv)
Figures -> figures/          (.pdf)

Inputs:
  data/processed/merged_analysis.parquet   (preferred)
  data/processed/merged_spatial.parquet    (fallback)
  data/processed/airbnb_clean.parquet
  data/processed/rent_clean.parquet
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.plotting import use_descriptive_style

use_descriptive_style()

import matplotlib.pyplot as plt  # noqa: E402
import matplotlib.ticker as mticker  # noqa: E402
import pandas as pd  # noqa: E402

from src.config import (  # noqa: E402
    AIRBNB_CLEAN,
    RENT_CLEAN,
    MERGED_ANALYSIS,
    MERGED_SPATIAL,
    FIG_DIR,
    TABLE_DIR,
    require,
)

PROCESSED_HINT = "Run scripts/04_merge_data.py first."


# ===================================================================
# Helper: summary statistics
# ===================================================================

def summary_stats(df: pd.DataFrame, cols: list[str], labels: dict | None = None) -> pd.DataFrame:
    """
    Compute N, mean, sd, min, p25, median, p75, max for *cols*.
    Columns absent from *df* are silently skipped (matches the published
    tables, which omit variables not present in the processed data).
    """
    rows = []
    for c in cols:
        if c not in df.columns:
            continue
        s = df[c].dropna()
        label = (labels or {}).get(c, c)
        rows.append({
            "Variable": label,
            "N": int(len(s)),
            "Mean": s.mean(),
            "SD": s.std(),
            "Min": s.min(),
            "P25": s.quantile(0.25),
            "Median": s.median(),
            "P75": s.quantile(0.75),
            "Max": s.max(),
        })
    return pd.DataFrame(rows)


def save_table(df: pd.DataFrame, name: str):
    """Save as .csv and .tex (a tabular fragment, wrapped by the paper)."""
    csv_path = TABLE_DIR / f"{name}.csv"
    tex_path = TABLE_DIR / f"{name}.tex"
    df.to_csv(csv_path, index=False)
    # LaTeX: format floats to 2 decimals except integers
    float_fmt = "%.2f"
    df.to_latex(tex_path, index=False, float_format=float_fmt)
    print(f"    Saved {csv_path.name}  +  {tex_path.name}")


# ===================================================================
# Tables
# ===================================================================

def make_tables(airbnb: pd.DataFrame, rent: pd.DataFrame, merged: pd.DataFrame):
    print("\n--- Tables ---")

    # 1. Summary statistics: Airbnb
    airbnb_vars = [
        "price_numeric", "rating", "bedrooms", "bathrooms",
        "guests_count", "amenities_count", "num_images", "num_reviews",
        "quality_score",
    ]
    airbnb_labels = {
        "price_numeric": "Nightly price (CAD)",
        "rating": "Rating",
        "bedrooms": "Bedrooms",
        "bathrooms": "Bathrooms",
        "guests_count": "Max guests",
        "amenities_count": "Amenities count",
        "num_images": "Number of images",
        "num_reviews": "Number of reviews",
        "quality_score": "Quality score",
    }
    tbl_airbnb = summary_stats(airbnb, airbnb_vars, airbnb_labels)
    save_table(tbl_airbnb, "summary_stats_airbnb")

    # 2. Summary statistics: Rent
    rent_vars = [
        "monthly_rent", "bedrooms", "bathrooms", "size_sqft",
    ]
    rent_labels = {
        "monthly_rent": "Monthly rent (CAD)",
        "bedrooms": "Bedrooms",
        "bathrooms": "Bathrooms",
        "size_sqft": "Size (sq ft)",
    }
    tbl_rent = summary_stats(rent, rent_vars, rent_labels)
    save_table(tbl_rent, "summary_stats_rent")

    # 3. Correlation matrix
    corr_vars = []
    # Rent outcome
    for v in ["log_rent", "monthly_rent"]:
        if v in merged.columns:
            corr_vars.append(v)
            break  # prefer log_rent

    # Airbnb exposure measures (500m buffer)
    for v in ["airbnb_count_500m", "airbnb_density_500m",
              "mean_airbnb_price_500m", "share_entire_home_500m"]:
        if v in merged.columns:
            corr_vars.append(v)

    # Property characteristics
    for v in ["bedrooms", "bathrooms", "size_sqft"]:
        if v in merged.columns:
            corr_vars.append(v)

    if len(corr_vars) >= 3:
        corr_df = merged[corr_vars].dropna()
        corr_mat = corr_df.corr()
        save_table(corr_mat.round(3).reset_index().rename(columns={"index": ""}),
                   "correlation_matrix")
    else:
        print("    Skipping correlation matrix — not enough overlapping columns.")


# ===================================================================
# Figures
# ===================================================================

def fig_dist_airbnb_price(airbnb: pd.DataFrame):
    """Histogram of Airbnb nightly prices."""
    prices = airbnb["price_numeric"].dropna()
    # Trim extreme outliers (above 99th percentile) for readability
    p99 = prices.quantile(0.99)
    prices_trim = prices[prices <= p99]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(prices_trim, bins=60, color="#4C72B0", edgecolor="white", linewidth=0.5)
    ax.set_xlabel("Nightly price (CAD)")
    ax.set_ylabel("Count")
    ax.set_title("Distribution of Airbnb Nightly Prices")
    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))
    fig.tight_layout()
    fig.savefig(FIG_DIR / "dist_airbnb_price.pdf")
    plt.close(fig)
    print("    dist_airbnb_price.pdf")


def fig_dist_rent(rent: pd.DataFrame):
    """Histogram of monthly rents."""
    rents = rent["monthly_rent"].dropna()
    p99 = rents.quantile(0.99)
    rents_trim = rents[rents <= p99]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(rents_trim, bins=60, color="#DD8452", edgecolor="white", linewidth=0.5)
    ax.set_xlabel("Monthly rent (CAD)")
    ax.set_ylabel("Count")
    ax.set_title("Distribution of Monthly Rents")
    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))
    fig.tight_layout()
    fig.savefig(FIG_DIR / "dist_rent.pdf")
    plt.close(fig)
    print("    dist_rent.pdf")


def fig_airbnb_by_city(airbnb: pd.DataFrame):
    """Bar chart of Airbnb listing count by top 15 cities."""
    city_col = "city_clean" if "city_clean" in airbnb.columns else "city"
    counts = airbnb[city_col].value_counts().head(15)

    fig, ax = plt.subplots(figsize=(8, 5))
    counts.sort_values().plot.barh(ax=ax, color="#4C72B0", edgecolor="white")
    ax.set_xlabel("Number of listings")
    ax.set_ylabel("")
    ax.set_title("Airbnb Listings by City (Top 15)")
    fig.tight_layout()
    fig.savefig(FIG_DIR / "airbnb_by_city.pdf")
    plt.close(fig)
    print("    airbnb_by_city.pdf")


def fig_rent_by_city(rent: pd.DataFrame):
    """Bar chart of mean rent by top 15 cities."""
    city_counts = rent["city"].value_counts()
    top_cities = city_counts.head(15).index
    sub = rent[rent["city"].isin(top_cities)]
    means = sub.groupby("city")["monthly_rent"].mean().sort_values()

    fig, ax = plt.subplots(figsize=(8, 5))
    means.plot.barh(ax=ax, color="#DD8452", edgecolor="white")
    ax.set_xlabel("Mean monthly rent (CAD)")
    ax.set_ylabel("")
    ax.set_title("Mean Monthly Rent by City (Top 15)")
    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))
    fig.tight_layout()
    fig.savefig(FIG_DIR / "rent_by_city.pdf")
    plt.close(fig)
    print("    rent_by_city.pdf")


def fig_scatter_airbnb_rent(merged: pd.DataFrame):
    """Scatter plot of mean Airbnb count (500m) vs log rent
    at a neighbourhood level (postal code or small-area average)."""
    if "airbnb_count_500m" not in merged.columns:
        print("    Skipping scatter_airbnb_rent.pdf — no airbnb_count_500m column.")
        return

    rent_col = "log_rent" if "log_rent" in merged.columns else "monthly_rent"

    # Aggregate at postal-code level if available, else use city
    group_col = None
    for candidate in ["postal_code", "city"]:
        if candidate in merged.columns:
            group_col = candidate
            break

    if group_col is None:
        print("    Skipping scatter_airbnb_rent.pdf — no grouping column found.")
        return

    agg = merged.groupby(group_col).agg(
        mean_airbnb_count=("airbnb_count_500m", "mean"),
        mean_rent=(rent_col, "mean"),
        n=("airbnb_count_500m", "size"),
    ).reset_index()

    # Keep neighbourhoods with at least 5 observations
    agg = agg[agg["n"] >= 5]

    fig, ax = plt.subplots(figsize=(6, 5))
    ax.scatter(agg["mean_airbnb_count"], agg["mean_rent"],
               alpha=0.6, s=20, color="#4C72B0", edgecolors="none")
    ax.set_xlabel("Mean Airbnb count within 500 m")
    y_label = "Log monthly rent" if rent_col == "log_rent" else "Mean monthly rent (CAD)"
    ax.set_ylabel(y_label)
    ax.set_title("Airbnb Density vs. Rent")
    fig.tight_layout()
    fig.savefig(FIG_DIR / "scatter_airbnb_rent.pdf")
    plt.close(fig)
    print("    scatter_airbnb_rent.pdf")


def fig_map_airbnb(airbnb: pd.DataFrame):
    """Scatter map of Airbnb listings coloured by price."""
    lon_col = "lon" if "lon" in airbnb.columns else "long"
    valid = airbnb.dropna(subset=["lat", lon_col, "price_numeric"])

    # Cap at 99th percentile for colour scale
    p99 = valid["price_numeric"].quantile(0.99)

    fig, ax = plt.subplots(figsize=(8, 7))
    sc = ax.scatter(
        valid[lon_col], valid["lat"],
        c=valid["price_numeric"].clip(upper=p99),
        cmap="YlOrRd", s=4, alpha=0.6, edgecolors="none",
    )
    cbar = fig.colorbar(sc, ax=ax, shrink=0.7)
    cbar.set_label("Nightly price (CAD)")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_title("Airbnb Listings — Quebec")
    fig.tight_layout()
    fig.savefig(FIG_DIR / "map_airbnb.pdf")
    plt.close(fig)
    print("    map_airbnb.pdf")


def fig_map_rent(rent: pd.DataFrame):
    """Scatter map of rental listings coloured by rent."""
    valid = rent.dropna(subset=["lat", "lon", "monthly_rent"])
    p99 = valid["monthly_rent"].quantile(0.99)

    fig, ax = plt.subplots(figsize=(8, 7))
    sc = ax.scatter(
        valid["lon"], valid["lat"],
        c=valid["monthly_rent"].clip(upper=p99),
        cmap="YlGnBu", s=4, alpha=0.6, edgecolors="none",
    )
    cbar = fig.colorbar(sc, ax=ax, shrink=0.7)
    cbar.set_label("Monthly rent (CAD)")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_title("Rental Listings — Quebec")
    fig.tight_layout()
    fig.savefig(FIG_DIR / "map_rent.pdf")
    plt.close(fig)
    print("    map_rent.pdf")


# ===================================================================
# Main
# ===================================================================

def main():
    print("=" * 70)
    print("05  DESCRIPTIVE ANALYSIS")
    print("=" * 70)

    # ------------------------------------------------------------------
    # Load data
    # ------------------------------------------------------------------
    print("\n[1] Loading data ...")

    airbnb = pd.read_parquet(require(AIRBNB_CLEAN, PROCESSED_HINT))
    if "long" in airbnb.columns and "lon" not in airbnb.columns:
        airbnb = airbnb.rename(columns={"long": "lon"})

    rent = pd.read_parquet(require(RENT_CLEAN, PROCESSED_HINT))

    merged_path = MERGED_ANALYSIS
    if not merged_path.exists():
        merged_path = MERGED_SPATIAL
    merged = pd.read_parquet(require(merged_path, PROCESSED_HINT))
    if "long" in merged.columns and "lon" not in merged.columns:
        merged = merged.rename(columns={"long": "lon"})

    print(f"    Airbnb : {airbnb.shape}")
    print(f"    Rent   : {rent.shape}")
    print(f"    Merged : {merged.shape}  ({merged_path.name})")

    # ------------------------------------------------------------------
    # Tables
    # ------------------------------------------------------------------
    make_tables(airbnb, rent, merged)

    # ------------------------------------------------------------------
    # Figures
    # ------------------------------------------------------------------
    print("\n--- Figures ---")
    fig_dist_airbnb_price(airbnb)
    fig_dist_rent(rent)
    fig_airbnb_by_city(airbnb)
    fig_rent_by_city(rent)
    fig_scatter_airbnb_rent(merged)
    fig_map_airbnb(airbnb)
    fig_map_rent(rent)

    print("\nDone.\n")


if __name__ == "__main__":
    main()