#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
#
"""Step 5 — Generate all paper figures (fig1-fig11).

Faithful port of paper_figures_v2.py: same layouts, colors and parameters,
but reading the pipeline's analysis dataset instead of re-encoding the
embeddings, and writing into the repository's figures/ directory.

Inputs : data/processed/hedonic_maison_results.csv
Outputs: figures/fig{1..11}_*.pdf and .png
"""

import sys
import warnings
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from scipy.stats import probplot

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src import config
from src.models import fit_all_models
from src.references import ENGLISH_LABELS, REFERENCES, SIM_COLS

warnings.filterwarnings("ignore")

plt.rcParams.update({
    "font.family": "serif",
    "font.serif": ["Times New Roman", "Times", "DejaVu Serif"],
    "font.size": 10,
    "axes.labelsize": 11,
    "axes.titlesize": 12,
    "xtick.labelsize": 9,
    "ytick.labelsize": 9,
    "legend.fontsize": 9,
    "figure.dpi": 300,
    "savefig.dpi": 300,
    "savefig.bbox": "tight",
    "savefig.pad_inches": 0.08,
    "axes.spines.top": False,
    "axes.spines.right": False,
})

LABELS = [ENGLISH_LABELS[slug] for slug in REFERENCES]


def save(fig_name):
    plt.tight_layout()
    for ext in ("pdf", "png"):
        plt.savefig(config.FIGURES_DIR / f"{fig_name}.{ext}")
    plt.close()
    print(f"  {fig_name}")


def fig1_model_comparison(models):
    fig, ax = plt.subplots(figsize=(6.2, 3.8))
    names = ["Model A\nStructural", "Model B\n+ Length", "Model C\n+ Semantic",
             "Model D\nFull", "Model E\nParsim."]
    r2v = [m.rsquared for m in models.values()]
    r2a = [m.rsquared_adj for m in models.values()]
    x = np.arange(len(names))
    b1 = ax.bar(x - 0.15, r2v, 0.28, label=r"$R^2$", color="#1565C0", alpha=0.9)
    ax.bar(x + 0.15, r2a, 0.28, label=r"Adj. $R^2$", color="#FB8C00", alpha=0.9)
    for b, v in zip(b1, r2v):
        ax.text(b.get_x() + b.get_width() / 2, b.get_height() + 0.006,
                f"{v:.3f}", ha="center", fontsize=7.5)
    ax.set_xticks(x)
    ax.set_xticklabels(names, fontsize=8)
    ax.set_ylabel(r"$R^2$")
    ax.set_ylim(0, 0.58)
    ax.legend(loc="upper left", framealpha=0.9)
    ax.axhline(models["A"].rsquared, color="grey", ls="--", lw=0.7, alpha=0.5)
    save("fig1_model_comparison")


def fig2_coefficient_plot(mD):
    fig, ax = plt.subplots(figsize=(6.2, 7))
    conf = mD.conf_int()
    sd = []
    for v, label in zip(SIM_COLS, LABELS):
        sd.append((label, mD.params[v], conf.loc[v, 0], conf.loc[v, 1],
                   mD.pvalues[v], (np.exp(mD.params[v]) - 1) * 100))
    sd.sort(key=lambda t: t[1])
    names = [d[0] for d in sd]
    coefs = [d[1] for d in sd]
    pvs = [d[4] for d in sd]
    yp = np.arange(len(names))
    cols = ["#C62828" if c < 0 and p < 0.05 else
            "#2E7D32" if c > 0 and p < 0.05 else "#BDBDBD"
            for c, p in zip(coefs, pvs)]
    ax.axvline(0, color="black", lw=0.8)
    ax.barh(yp, coefs, color=cols, height=0.55, alpha=0.85,
            edgecolor="white", lw=0.4)
    for i, (name, c, cl, ch, p, imp) in enumerate(sd):
        ax.plot([cl, ch], [yp[i], yp[i]], color="#333", lw=1)
        st = "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else ""
        if st:
            off = 0.004 if c >= 0 else -0.004
            ha = "left" if c >= 0 else "right"
            ax.text(c + off, yp[i], f"{imp:+.1f}%{st}", va="center", ha=ha,
                    fontsize=7.5)
    ax.set_yticks(yp)
    ax.set_yticklabels(names, fontsize=9)
    ax.set_xlabel("Standardized coefficient (log-price)")
    save("fig2_coefficient_plot")


def fig3_similarity_distributions(df):
    fig, ax = plt.subplots(figsize=(6.2, 5.5))
    bd = [df[c].values for c in SIM_COLS]
    order = np.argsort([np.median(d) for d in bd])[::-1]
    bd = [bd[i] for i in order]
    bl = [LABELS[i] for i in order]
    bp = ax.boxplot(bd, vert=False, patch_artist=True, widths=0.55,
                    flierprops=dict(marker=".", markersize=1.5, alpha=0.2),
                    medianprops=dict(color="black", lw=1.2))
    cm = plt.cm.viridis
    for i, patch in enumerate(bp["boxes"]):
        patch.set_facecolor(cm(i / len(bp["boxes"])))
        patch.set_alpha(0.7)
    ax.set_yticklabels(bl, fontsize=8)
    ax.set_xlabel("Cosine similarity")
    save("fig3_similarity_distributions")


def fig4_quintile_heatmap(df):
    df = df.copy()
    df["pq"] = pd.qcut(df["price"], q=5,
                       labels=["Q1 (Low)", "Q2", "Q3 (Med)", "Q4", "Q5 (High)"])
    profile = df.groupby("pq", observed=False)[SIM_COLS].mean()
    fig, ax = plt.subplots(figsize=(6.2, 6.5))
    im = ax.imshow(profile.T.values, aspect="auto", cmap="RdYlGn",
                   interpolation="nearest")
    ax.set_xticks(range(5))
    ax.set_xticklabels(profile.index, fontsize=9)
    ax.set_yticks(range(len(LABELS)))
    ax.set_yticklabels(LABELS, fontsize=8)
    ax.set_xlabel("Price quintile")
    for i in range(len(LABELS)):
        for j in range(5):
            v = profile.T.values[i, j]
            ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=6.5,
                    color="white" if v < 0.25 or v > 0.40 else "black")
    plt.colorbar(im, ax=ax, shrink=0.7, label="Mean cosine similarity", pad=0.02)
    save("fig4_quintile_heatmap")


def fig5_correlation_matrix(df):
    corr = df[SIM_COLS].corr().values
    mask = np.triu(np.ones_like(corr, dtype=bool), k=1)
    corr_show = np.where(mask, np.nan, corr)
    fig, ax = plt.subplots(figsize=(6.5, 6))
    im = ax.imshow(corr_show, cmap="RdBu_r", vmin=-0.1, vmax=1.0,
                   interpolation="nearest")
    ax.set_xticks(range(len(LABELS)))
    ax.set_xticklabels(LABELS, rotation=55, ha="right", fontsize=7)
    ax.set_yticks(range(len(LABELS)))
    ax.set_yticklabels(LABELS, fontsize=7)
    for i in range(len(LABELS)):
        for j in range(i + 1):
            v = corr[i, j]
            ax.text(j, i, f"{v:.2f}", ha="center", va="center", fontsize=5.5,
                    color="white" if abs(v) > 0.65 else "black")
    plt.colorbar(im, ax=ax, shrink=0.7, label="Pearson r", pad=0.02)
    save("fig5_correlation_matrix")


def fig6_scatter_plots(df):
    fig, axes = plt.subplots(2, 2, figsize=(6.2, 5.5))
    pairs = [
        ("sim_luxe", "Luxury", "#7B1FA2"),
        ("sim_a_renover", "Needs Renovation", "#C62828"),
        ("sim_moderne_contemporain", "Modern/Contemporary", "#1565C0"),
        ("sim_urgence_motivation", "Motivated Seller", "#E65100"),
    ]
    samp = df.sample(min(2500, len(df)), random_state=config.SEED)
    for ax, (col, label, color) in zip(axes.flat, pairs):
        ax.scatter(samp[col], samp["log_price"], alpha=0.12, s=5, c=color,
                   edgecolors="none")
        z = np.polyfit(df[col], df["log_price"], 1)
        xr = np.linspace(df[col].min(), df[col].max(), 100)
        ax.plot(xr, np.poly1d(z)(xr), color="black", lw=1.8)
        r = df[col].corr(df["log_price"])
        ax.set_xlabel(f"Sim: {label}", fontsize=8)
        ax.set_ylabel("log(Price)", fontsize=8)
        ax.set_title(f"{label} (r={r:.3f})", fontsize=9)
        ax.tick_params(labelsize=7)
    plt.tight_layout(h_pad=1.2, w_pad=0.8)
    for ext in ("pdf", "png"):
        plt.savefig(config.FIGURES_DIR / f"fig6_scatter_plots.{ext}")
    plt.close()
    print("  fig6_scatter_plots")


def fig7_methodology(n_obs):
    fig, ax = plt.subplots(figsize=(6.2, 2.8))
    ax.set_xlim(0, 12)
    ax.set_ylim(0, 3.5)
    ax.axis("off")
    boxes = [
        (1.2, 1.75, f"Property\nListings\n(n={n_obs:,})", "#E3F2FD"),
        (3.6, 1.75, "Sentence\nEmbeddings\n(384-d)", "#E8F5E9"),
        (6.0, 1.75, "Cosine\nSimilarity\n(20 refs)", "#FFF3E0"),
        (8.4, 1.75, "Hedonic\nOLS Model", "#F3E5F5"),
        (10.8, 1.75, "Implicit\nPrice\nEstimates", "#FFEBEE"),
    ]
    for x, yy, text, color in boxes:
        ax.add_patch(FancyBboxPatch((x - 0.65, yy - 0.65), 1.3, 1.3,
                                    boxstyle="round,pad=0.08",
                                    facecolor=color, edgecolor="#444", lw=1.3))
        ax.text(x, yy, text, ha="center", va="center", fontsize=7.5,
                fontweight="bold")
    for i in range(len(boxes) - 1):
        ax.annotate("", xy=(boxes[i + 1][0] - 0.7, 1.75),
                    xytext=(boxes[i][0] + 0.7, 1.75),
                    arrowprops=dict(arrowstyle="->", color="#444", lw=1.8))
    labs = [
        (2.4, 0.55, "PublicRemarks\nextraction"),
        (4.8, 0.55, "all-MiniLM-L6-v2\ntransformer"),
        (7.2, 0.55, "Reference-based\nfeatures"),
        (9.6, 0.55, "log(P) = βX+γS+ε"),
    ]
    for x, yy, text in labs:
        ax.text(x, yy, text, ha="center", va="center", fontsize=7,
                style="italic", color="#666")
    save("fig7_methodology")


def fig8_r2_decomposition(models):
    fig, ax = plt.subplots(figsize=(4.5, 4))
    mA, mB, mD = models["A"], models["B"], models["D"]
    comps = [
        ("Structural variables", mA.rsquared, "#1565C0"),
        ("Text length", mB.rsquared - mA.rsquared, "#43A047"),
        ("Semantic similarities", mD.rsquared - mB.rsquared, "#FB8C00"),
    ]
    bot = 0
    for lab, val, col in comps:
        ax.bar(0, val, bottom=bot, color=col, width=0.5, edgecolor="white",
               lw=0.5, label=f"{lab}: {val:.4f}")
        if val > 0.008:
            ax.text(0, bot + val / 2, f"{val:.4f}\n({val / mD.rsquared * 100:.1f}%)",
                    ha="center", va="center", fontsize=8, fontweight="bold",
                    color="white")
        bot += val
    ax.set_ylabel(r"$R^2$")
    ax.set_ylim(0, 0.58)
    ax.set_xticks([0])
    ax.set_xticklabels(["Full Model (D)"], fontsize=9)
    ax.legend(loc="upper left", fontsize=8, framealpha=0.9)
    save("fig8_r2_decomposition")


def fig9_price_distribution(df):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.2, 3))
    ax1.hist(df["price"] / 1000, bins=80, color="#1565C0", alpha=0.8,
             edgecolor="white", lw=0.3)
    ax1.set_xlabel("Price (\\$000s)")
    ax1.set_ylabel("Frequency")
    ax1.set_title("(a) Price distribution", fontsize=10)
    ax1.set_xlim(0, 3000)
    ax2.hist(df["log_price"], bins=80, color="#2E7D32", alpha=0.8,
             edgecolor="white", lw=0.3)
    ax2.set_xlabel("log(Price)")
    ax2.set_ylabel("Frequency")
    ax2.set_title("(b) Log-price distribution", fontsize=10)
    plt.tight_layout(w_pad=1.5)
    for ext in ("pdf", "png"):
        plt.savefig(config.FIGURES_DIR / f"fig9_price_distribution.{ext}")
    plt.close()
    print("  fig9_price_distribution")


STRUCT_LABELS = {
    "bedrooms": "Bedrooms", "bathrooms": "Bathrooms",
    "half_baths": "Half-bathrooms", "parking": "Parking",
    "stories": "Stories", "land_size": "Lot size",
    "remarks_length": "Description length",
}


def fig10_structural_coefficients(mD):
    fig, ax = plt.subplots(figsize=(6.2, 3))
    sv = config.STRUCTURAL_VARS + ["remarks_length"]
    conf = mD.conf_int()
    sdata = sorted(
        [(v, mD.params[v], conf.loc[v, 0], conf.loc[v, 1], mD.pvalues[v])
         for v in sv],
        key=lambda t: t[1],
    )
    yy = np.arange(len(sdata))
    cols = ["#1565C0" if d[4] < 0.05 else "#BDBDBD" for d in sdata]
    ax.barh(yy, [d[1] for d in sdata], color=cols, height=0.5, alpha=0.85)
    for i, (v, c, cl, ch, p) in enumerate(sdata):
        ax.plot([cl, ch], [yy[i], yy[i]], color="#333", lw=1)
        st = "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else ""
        if st:
            ax.text(c + 0.005, yy[i], f"{(np.exp(c) - 1) * 100:+.1f}%{st}",
                    va="center", fontsize=7.5)
    ax.axvline(0, color="black", lw=0.7)
    ax.set_yticks(yy)
    ax.set_yticklabels([STRUCT_LABELS[d[0]] for d in sdata], fontsize=9)
    ax.set_xlabel("Standardized coefficient")
    save("fig10_structural_coefficients")


def fig11_residual_diagnostics(mD):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6.2, 3))
    ax1.scatter(mD.fittedvalues, mD.resid, s=2, alpha=0.1, c="#1565C0",
                edgecolors="none")
    ax1.axhline(0, color="red", lw=0.8, ls="--")
    ax1.set_xlabel("Fitted values")
    ax1.set_ylabel("Residuals")
    ax1.set_title("(a) Residuals vs. fitted", fontsize=10)
    probplot(mD.resid, dist="norm", plot=ax2)
    ax2.set_title("(b) Normal Q-Q plot", fontsize=10)
    ax2.get_lines()[0].set(markersize=2, alpha=0.15, color="#1565C0")
    ax2.get_lines()[1].set(color="red", lw=1)
    plt.tight_layout(w_pad=1.5)
    for ext in ("pdf", "png"):
        plt.savefig(config.FIGURES_DIR / f"fig11_residual_diagnostics.{ext}")
    plt.close()
    print("  fig11_residual_diagnostics")


def main():
    df = pd.read_csv(config.ANALYSIS_CSV)
    print(f"{len(df):,} observations — fitting models ...")
    models, _, _ = fit_all_models(df)
    mD = models["D"]

    config.FIGURES_DIR.mkdir(parents=True, exist_ok=True)
    print("Generating figures ...")
    fig1_model_comparison(models)
    fig2_coefficient_plot(mD)
    fig3_similarity_distributions(df)
    fig4_quintile_heatmap(df)
    fig5_correlation_matrix(df)
    fig6_scatter_plots(df)
    fig7_methodology(len(df))
    fig8_r2_decomposition(models)
    fig9_price_distribution(df)
    fig10_structural_coefficients(mD)
    fig11_residual_diagnostics(mD)
    print("All figures saved to figures/")


if __name__ == "__main__":
    main()