"""
Airbnb Hedonic Pricing Lab
==========================
Interactive Streamlit dashboard for hedonic pricing analysis of Quebec Airbnb listings.
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import warnings
import os
import joblib

warnings.filterwarnings("ignore")

# ---------------------------------------------------------------------------
# Page config
# ---------------------------------------------------------------------------
st.set_page_config(
    page_title="Airbnb Hedonic Pricing Lab",
    page_icon="🏠",
    layout="wide",
    initial_sidebar_state="expanded",
)

DATA_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "raw", "cleaned_dataset_sample.csv")
MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results", "models")

# ---------------------------------------------------------------------------
# Custom CSS
# ---------------------------------------------------------------------------
st.markdown(
    """
    <style>
    .block-container {padding-top: 1.5rem; padding-bottom: 1rem;}
    h1 {color: #FF5A5F;}
    .stTabs [data-baseweb="tab-list"] {gap: 8px;}
    .stTabs [data-baseweb="tab"] {
        padding: 8px 20px;
        border-radius: 6px 6px 0 0;
        font-weight: 600;
    }
    .metric-card {
        background: #f8f9fa; border-radius: 10px; padding: 16px;
        text-align: center; border: 1px solid #e9ecef;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
@st.cache_data
def load_data():
    df = pd.read_csv(DATA_PATH)
    # Ensure boolean columns are truly boolean
    for col in ["is_superhost", "is_guest_favorite", "pets_allowed"]:
        df[col] = df[col].astype(bool)
    # Log price (for regression)
    df["log_price"] = np.log1p(df["price_numeric"])
    return df


raw_df = load_data()

# ---------------------------------------------------------------------------
# Sidebar: global filters
# ---------------------------------------------------------------------------
st.sidebar.title("Airbnb Hedonic Pricing Lab")
st.sidebar.markdown("---")
st.sidebar.header("Global Filters")

all_cities = sorted(raw_df["city"].dropna().unique())
top_cities = (
    raw_df["city"]
    .value_counts()
    .head(20)
    .index.tolist()
)

selected_cities = st.sidebar.multiselect(
    "Cities",
    options=all_cities,
    default=top_cities[:5],
    help="Leave empty to include all cities.",
)

all_property_types = sorted(raw_df["property_type"].dropna().unique())
selected_property_types = st.sidebar.multiselect(
    "Property Types",
    options=all_property_types,
    default=all_property_types,
)

price_valid = raw_df["price_numeric"].dropna()
price_min_data = float(price_valid.min())
price_max_data = float(min(price_valid.max(), 10000))  # cap slider for UX
price_range = st.sidebar.slider(
    "Price Range ($)",
    min_value=0.0,
    max_value=price_max_data,
    value=(0.0, price_max_data),
    step=10.0,
)

superhost_only = st.sidebar.checkbox("Superhost only", value=False)

st.sidebar.markdown("---")
st.sidebar.caption(
    f"Dataset: {len(raw_df):,} listings across {raw_df['city'].nunique()} cities in Quebec."
)

# ---------------------------------------------------------------------------
# Apply global filters
# ---------------------------------------------------------------------------
def apply_global_filters(df: pd.DataFrame) -> pd.DataFrame:
    mask = pd.Series(True, index=df.index)
    if selected_cities:
        mask &= df["city"].isin(selected_cities)
    if selected_property_types:
        mask &= df["property_type"].isin(selected_property_types)
    mask &= (df["price_numeric"] >= price_range[0]) | df["price_numeric"].isna()
    mask &= (df["price_numeric"] <= price_range[1]) | df["price_numeric"].isna()
    if superhost_only:
        mask &= df["is_superhost"]
    return df.loc[mask].copy()


df = apply_global_filters(raw_df)

# ---------------------------------------------------------------------------
# Hedonic model helpers
# ---------------------------------------------------------------------------
FEATURE_COLS = [
    "bedrooms",
    "bathrooms",
    "guests_count",
    "amenities_count",
    "num_images",
    "rating",
    "is_superhost",
    "is_guest_favorite",
    "pets_allowed",
]

NUMERIC_FEATURES = [
    "bedrooms",
    "bathrooms",
    "guests_count",
    "amenities_count",
    "num_images",
    "rating",
]

BOOL_FEATURES = ["is_superhost", "is_guest_favorite", "pets_allowed"]


def _prepare_model_df(data: pd.DataFrame) -> pd.DataFrame:
    """Return a clean DataFrame ready for modelling (no NaN in target/features)."""
    mdf = data.dropna(subset=["price_numeric"]).copy()
    mdf = mdf[mdf["price_numeric"] > 0].copy()
    mdf["log_price"] = np.log(mdf["price_numeric"])
    for col in BOOL_FEATURES:
        mdf[col] = mdf[col].astype(int)
    # City dummies (top cities)
    city_counts = mdf["city"].value_counts()
    top = city_counts[city_counts >= 20].index.tolist()
    for c in top:
        mdf[f"city_{c}"] = (mdf["city"] == c).astype(int)
    # Property type dummies
    for pt in mdf["property_type"].unique():
        mdf[f"ptype_{pt}"] = (mdf["property_type"] == pt).astype(int)
    # Fill remaining NaN in numeric features with median
    for col in NUMERIC_FEATURES:
        mdf[col] = mdf[col].fillna(mdf[col].median())
    return mdf


@st.cache_resource
def fit_ols_model():
    """Fit an OLS hedonic model on the full raw dataset."""
    from sklearn.linear_model import LinearRegression

    mdf = _prepare_model_df(raw_df)
    feature_names = list(FEATURE_COLS)
    city_cols = [c for c in mdf.columns if c.startswith("city_")]
    ptype_cols = [c for c in mdf.columns if c.startswith("ptype_")]
    feature_names = FEATURE_COLS + city_cols + ptype_cols
    X = mdf[feature_names].values
    y = mdf["log_price"].values
    model = LinearRegression()
    model.fit(X, y)
    return model, feature_names, mdf


@st.cache_resource
def fit_gbm_model():
    """Fit a gradient-boosted model for prediction & SHAP."""
    try:
        from sklearn.ensemble import GradientBoostingRegressor
    except ImportError:
        return None, None, None

    mdf = _prepare_model_df(raw_df)
    feature_names = list(FEATURE_COLS)
    city_cols = sorted([c for c in mdf.columns if c.startswith("city_")])
    ptype_cols = sorted([c for c in mdf.columns if c.startswith("ptype_")])
    feature_names = FEATURE_COLS + city_cols + ptype_cols
    X = mdf[feature_names].values
    y = mdf["log_price"].values
    model = GradientBoostingRegressor(
        n_estimators=200,
        max_depth=5,
        learning_rate=0.1,
        subsample=0.8,
        random_state=42,
    )
    model.fit(X, y)
    return model, feature_names, mdf


# ---------------------------------------------------------------------------
# TABS
# ---------------------------------------------------------------------------
tab1, tab2, tab3, tab4, tab5, tab6, tab7 = st.tabs(
    [
        "Data Explorer",
        "Price Map",
        "Neighbourhood Comparison",
        "Hedonic Premiums",
        "Price Predictor",
        "SHAP Explorer",
        "Price Decomposition",
    ]
)

# ===================================================================
# TAB 1 - Data Explorer
# ===================================================================
with tab1:
    st.header("Data Explorer")

    # Summary metrics
    col1, col2, col3, col4, col5 = st.columns(5)
    with col1:
        st.metric("Listings", f"{len(df):,}")
    with col2:
        st.metric("Cities", f"{df['city'].nunique()}")
    with col3:
        median_price = df["price_numeric"].median()
        st.metric("Median Price", f"${median_price:,.0f}" if pd.notna(median_price) else "N/A")
    with col4:
        avg_rating = df["rating"].mean()
        st.metric("Avg Rating", f"{avg_rating:.2f}" if pd.notna(avg_rating) else "N/A")
    with col5:
        superhost_pct = df["is_superhost"].mean() * 100
        st.metric("Superhosts", f"{superhost_pct:.1f}%")

    st.markdown("---")

    # Summary statistics table
    st.subheader("Summary Statistics")
    numeric_cols = [
        "price_numeric", "rating", "bedrooms", "bathrooms",
        "guests_count", "amenities_count", "num_images", "num_reviews", "quality_score",
    ]
    summary = df[numeric_cols].describe().T
    summary.columns = ["Count", "Mean", "Std", "Min", "25%", "50%", "75%", "Max"]
    st.dataframe(summary.style.format("{:.2f}"), use_container_width=True)

    st.markdown("---")

    # Price distribution
    st.subheader("Price Distribution")
    price_data = df["price_numeric"].dropna()
    if len(price_data) > 0:
        cap = st.slider(
            "Upper price cap for histogram",
            min_value=100,
            max_value=5000,
            value=1500,
            step=50,
            key="price_hist_cap",
        )
        fig_hist = px.histogram(
            price_data[price_data <= cap],
            nbins=60,
            title="Price Distribution (capped)",
            labels={"value": "Price ($)", "count": "Frequency"},
            color_discrete_sequence=["#FF5A5F"],
        )
        fig_hist.update_layout(
            showlegend=False,
            xaxis_title="Price ($)",
            yaxis_title="Count",
            template="plotly_white",
        )
        st.plotly_chart(fig_hist, use_container_width=True)
    else:
        st.info("No price data available for the current filter selection.")

    st.markdown("---")

    # Filtered data table
    st.subheader("Filtered Data")
    display_cols = [
        "name", "city", "property_type", "price_numeric", "rating",
        "bedrooms", "bathrooms", "guests_count", "amenities_count",
        "is_superhost", "pets_allowed",
    ]
    st.dataframe(
        df[display_cols].reset_index(drop=True),
        use_container_width=True,
        height=400,
    )

# ===================================================================
# TAB 2 - Price Map
# ===================================================================
with tab2:
    st.header("Price Map")

    map_df = df.dropna(subset=["price_numeric", "lat", "long"]).copy()
    map_df = map_df[(map_df["lat"] != 0) & (map_df["long"] != 0)]

    if len(map_df) == 0:
        st.warning("No geolocated listings with price data in current filter.")
    else:
        map_col1, map_col2 = st.columns([1, 3])
        with map_col1:
            color_by = st.radio(
                "Color by", ["price_numeric", "log_price", "rating"], index=0, key="map_color"
            )
            map_price_cap = st.slider(
                "Price cap for map color scale",
                100, 5000, 2000, 100,
                key="map_cap",
            )

        map_plot = map_df.copy()
        if color_by == "price_numeric":
            map_plot["color_val"] = map_plot["price_numeric"].clip(upper=map_price_cap)
        elif color_by == "log_price":
            map_plot["color_val"] = map_plot["log_price"]
        else:
            map_plot["color_val"] = map_plot["rating"]

        map_plot["size_val"] = map_plot["guests_count"].clip(lower=1, upper=16)

        with map_col2:
            fig_map = px.scatter_map(
                map_plot,
                lat="lat",
                lon="long",
                color="color_val",
                size="size_val",
                color_continuous_scale="RdYlGn_r" if "price" in color_by else "Viridis",
                size_max=14,
                zoom=5,
                hover_name="name",
                hover_data={
                    "city": True,
                    "property_type": True,
                    "price_numeric": ":.0f",
                    "rating": ":.2f",
                    "guests_count": True,
                    "color_val": False,
                    "size_val": False,
                    "lat": False,
                    "long": False,
                },
                map_style="carto-positron",
                title="Listings Map",
                height=650,
            )
            fig_map.update_layout(
                coloraxis_colorbar_title=color_by.replace("_", " ").title(),
                margin=dict(l=0, r=0, t=40, b=0),
            )
            st.plotly_chart(fig_map, use_container_width=True)

# ===================================================================
# TAB 3 - Neighbourhood Comparison
# ===================================================================
with tab3:
    st.header("Neighbourhood Comparison")

    compare_cities = st.multiselect(
        "Select cities to compare",
        options=all_cities,
        default=top_cities[:8],
        key="compare_cities",
    )

    if not compare_cities:
        st.info("Please select at least one city to compare.")
    else:
        cdf = raw_df[raw_df["city"].isin(compare_cities)].dropna(subset=["price_numeric"])
        cdf = cdf[cdf["price_numeric"] > 0]

        if len(cdf) == 0:
            st.warning("No price data for the selected cities.")
        else:
            c1, c2 = st.columns(2)

            with c1:
                st.subheader("Price Distribution by City")
                fig_box = px.box(
                    cdf,
                    x="city",
                    y="price_numeric",
                    color="city",
                    title="Price Box Plot by City",
                    labels={"price_numeric": "Price ($)", "city": "City"},
                    color_discrete_sequence=px.colors.qualitative.Set2,
                )
                fig_box.update_layout(
                    showlegend=False,
                    template="plotly_white",
                    yaxis=dict(range=[0, cdf["price_numeric"].quantile(0.95)]),
                )
                st.plotly_chart(fig_box, use_container_width=True)

            with c2:
                st.subheader("Median Price by City")
                city_med = (
                    cdf.groupby("city")["price_numeric"]
                    .median()
                    .sort_values(ascending=False)
                    .reset_index()
                )
                city_med.columns = ["City", "Median Price"]
                fig_bar = px.bar(
                    city_med,
                    x="Median Price",
                    y="City",
                    orientation="h",
                    title="Median Nightly Price by City",
                    color="Median Price",
                    color_continuous_scale="Sunset",
                )
                fig_bar.update_layout(template="plotly_white", yaxis=dict(autorange="reversed"))
                st.plotly_chart(fig_bar, use_container_width=True)

            # Summary table
            st.subheader("City Summary Statistics")
            city_stats = (
                cdf.groupby("city")
                .agg(
                    Listings=("price_numeric", "count"),
                    Mean_Price=("price_numeric", "mean"),
                    Median_Price=("price_numeric", "median"),
                    Std_Price=("price_numeric", "std"),
                    Avg_Rating=("rating", "mean"),
                    Avg_Bedrooms=("bedrooms", "mean"),
                    Superhost_Pct=("is_superhost", "mean"),
                )
                .round(2)
            )
            city_stats["Superhost_Pct"] = (city_stats["Superhost_Pct"] * 100).round(1)
            st.dataframe(city_stats.sort_values("Median_Price", ascending=False), use_container_width=True)

# ===================================================================
# TAB 4 - Hedonic Premiums (OLS)
# ===================================================================
with tab4:
    st.header("Hedonic Premiums (OLS Regression)")
    st.markdown(
        """
        A semi-log hedonic model is estimated: **ln(Price) = X * beta + epsilon**.
        Coefficients are converted to percentage premiums via **(exp(beta) - 1) x 100**.
        """
    )

    with st.spinner("Fitting OLS hedonic model..."):
        ols_model, ols_features, ols_mdf = fit_ols_model()

    coefs = pd.DataFrame(
        {"Feature": ols_features, "Coefficient": ols_model.coef_}
    )
    coefs["Premium (%)"] = (np.exp(coefs["Coefficient"]) - 1) * 100
    coefs = coefs.sort_values("Premium (%)", ascending=False).reset_index(drop=True)

    c1, c2 = st.columns([1, 2])

    with c1:
        st.subheader("Coefficient Table")
        st.dataframe(
            coefs.style.format({"Coefficient": "{:.4f}", "Premium (%)": "{:.2f}"}),
            use_container_width=True,
            height=500,
        )

        r2 = ols_model.score(
            ols_mdf[ols_features].values, ols_mdf["log_price"].values
        )
        st.metric("R-squared", f"{r2:.4f}")
        st.metric("N observations", f"{len(ols_mdf):,}")

    with c2:
        st.subheader("Hedonic Premiums: Key Features")
        key_features = FEATURE_COLS + [
            f for f in ols_features if f.startswith("ptype_")
        ]
        key_coefs = coefs[coefs["Feature"].isin(key_features)].sort_values("Premium (%)")

        fig_prem = px.bar(
            key_coefs,
            x="Premium (%)",
            y="Feature",
            orientation="h",
            color="Premium (%)",
            color_continuous_scale="RdYlGn",
            title="Estimated Price Premiums (%)",
            height=500,
        )
        fig_prem.update_layout(template="plotly_white")
        st.plotly_chart(fig_prem, use_container_width=True)

    st.markdown("---")
    st.subheader("Interpretation")
    st.markdown(
        """
        - **Positive premiums** indicate features that command a higher nightly price.
        - **Negative premiums** suggest a discount relative to the baseline.
        - For continuous variables (e.g., bedrooms), the premium is *per unit increase*.
        - For binary variables (e.g., is_superhost), the premium is relative to a non-superhost listing.
        - City and property-type dummies capture location and structural fixed effects.
        """
    )

# ===================================================================
# TAB 5 - Price Predictor
# ===================================================================
with tab5:
    st.header("Price Predictor")
    st.markdown("Adjust the inputs below to predict the nightly Airbnb price.")

    with st.spinner("Loading prediction model..."):
        gbm_model, gbm_features, gbm_mdf = fit_gbm_model()

    if gbm_model is None:
        st.error("Could not fit prediction model. Check dependencies.")
    else:
        pred_c1, pred_c2 = st.columns([1, 2])

        with pred_c1:
            st.subheader("Listing Features")
            pred_city = st.selectbox("City", options=top_cities[:20], index=0, key="pred_city")
            pred_ptype = st.selectbox("Property Type", options=all_property_types, index=0, key="pred_ptype")
            pred_bedrooms = st.slider("Bedrooms", 1, 10, 2, key="pred_bed")
            pred_bathrooms = st.slider("Bathrooms", 0.0, 5.0, 1.0, 0.5, key="pred_bath")
            pred_guests = st.slider("Guests", 1, 16, 4, key="pred_guests")
            pred_amenities = st.slider("Amenities count", 0, 100, 30, key="pred_amen")
            pred_images = st.slider("Number of images", 1, 50, 15, key="pred_img")
            pred_rating = st.slider("Rating", 0.0, 5.0, 4.8, 0.01, key="pred_rat")
            pred_superhost = st.checkbox("Superhost", value=False, key="pred_sh")
            pred_fav = st.checkbox("Guest Favorite", value=False, key="pred_fav")
            pred_pets = st.checkbox("Pets Allowed", value=False, key="pred_pets")

        # Build feature vector
        input_dict = {
            "bedrooms": pred_bedrooms,
            "bathrooms": pred_bathrooms,
            "guests_count": pred_guests,
            "amenities_count": pred_amenities,
            "num_images": pred_images,
            "rating": pred_rating,
            "is_superhost": int(pred_superhost),
            "is_guest_favorite": int(pred_fav),
            "pets_allowed": int(pred_pets),
        }
        # City dummies
        city_cols = [f for f in gbm_features if f.startswith("city_")]
        for c in city_cols:
            input_dict[c] = 1 if c == f"city_{pred_city}" else 0
        # Property type dummies
        ptype_cols = [f for f in gbm_features if f.startswith("ptype_")]
        for p in ptype_cols:
            input_dict[p] = 1 if p == f"ptype_{pred_ptype}" else 0

        X_pred = np.array([[input_dict.get(f, 0) for f in gbm_features]])
        log_pred = gbm_model.predict(X_pred)[0]
        pred_price = np.exp(log_pred)

        # Approximate confidence interval via training residuals
        y_train = gbm_mdf["log_price"].values
        y_hat_train = gbm_model.predict(gbm_mdf[gbm_features].values)
        residual_std = np.std(y_train - y_hat_train)
        ci_low = np.exp(log_pred - 1.96 * residual_std)
        ci_high = np.exp(log_pred + 1.96 * residual_std)

        with pred_c2:
            st.subheader("Predicted Price")

            m1, m2, m3 = st.columns(3)
            with m1:
                st.metric("Lower 95% CI", f"${ci_low:,.0f}")
            with m2:
                st.metric("Predicted Price", f"${pred_price:,.0f}")
            with m3:
                st.metric("Upper 95% CI", f"${ci_high:,.0f}")

            # Feature contribution breakdown
            st.subheader("Feature Contribution Breakdown")
            st.markdown(
                "Approximate marginal contribution of each feature group "
                "(setting each to its baseline and measuring the price change)."
            )

            # Compute contributions by zeroing out each feature
            baseline_dict = {}
            for f in gbm_features:
                baseline_dict[f] = gbm_mdf[f].median() if f in gbm_mdf.columns else 0
            X_base = np.array([[baseline_dict.get(f, 0) for f in gbm_features]])
            base_log = gbm_model.predict(X_base)[0]

            contributions = {}
            feature_groups = {
                "Bedrooms": ["bedrooms"],
                "Bathrooms": ["bathrooms"],
                "Guests": ["guests_count"],
                "Amenities": ["amenities_count"],
                "Images": ["num_images"],
                "Rating": ["rating"],
                "Superhost": ["is_superhost"],
                "Guest Favorite": ["is_guest_favorite"],
                "Pets Allowed": ["pets_allowed"],
                "City": city_cols,
                "Property Type": ptype_cols,
            }

            for group_name, cols in feature_groups.items():
                X_modified = X_pred.copy()
                for col_name in cols:
                    idx = gbm_features.index(col_name)
                    X_modified[0, idx] = baseline_dict.get(col_name, 0)
                modified_log = gbm_model.predict(X_modified)[0]
                # Contribution = what the feature adds relative to baseline
                contributions[group_name] = (log_pred - modified_log) * pred_price

            contrib_df = pd.DataFrame(
                {"Feature": list(contributions.keys()), "Contribution ($)": list(contributions.values())}
            ).sort_values("Contribution ($)")

            fig_contrib = px.bar(
                contrib_df,
                x="Contribution ($)",
                y="Feature",
                orientation="h",
                color="Contribution ($)",
                color_continuous_scale="RdYlGn",
                title="Feature Contributions to Predicted Price",
            )
            fig_contrib.update_layout(template="plotly_white", height=420)
            st.plotly_chart(fig_contrib, use_container_width=True)

# ===================================================================
# TAB 6 - SHAP Explorer
# ===================================================================
with tab6:
    st.header("SHAP Explorer")

    shap_available = True
    try:
        import shap
    except ImportError:
        shap_available = False

    if not shap_available:
        st.error(
            "The `shap` library is not installed. "
            "Install it with `pip install shap` to enable this tab."
        )
    else:
        with st.spinner("Computing SHAP values (this may take a moment)..."):
            gbm_model_s, gbm_features_s, gbm_mdf_s = fit_gbm_model()

        if gbm_model_s is None:
            st.error("Could not load model for SHAP analysis.")
        else:
            @st.cache_data
            def compute_shap_values():
                _model, _features, _mdf = fit_gbm_model()
                X_shap = _mdf[_features].values
                # Use a sample if dataset is large
                sample_size = min(500, len(X_shap))
                rng = np.random.RandomState(42)
                idx = rng.choice(len(X_shap), sample_size, replace=False)
                X_sample = X_shap[idx]
                explainer = shap.TreeExplainer(_model)
                sv = explainer.shap_values(X_sample)
                return sv, X_sample, _features

            shap_vals, X_shap_sample, shap_feature_names = compute_shap_values()

            # SHAP Summary Plot
            st.subheader("SHAP Summary Plot")
            fig_shap_summary, ax_shap = plt.subplots(figsize=(10, 7))
            shap.summary_plot(
                shap_vals,
                X_shap_sample,
                feature_names=shap_feature_names,
                show=False,
                max_display=20,
            )
            plt.tight_layout()
            st.pyplot(fig_shap_summary)
            plt.close(fig_shap_summary)

            # Feature importance ranking
            st.subheader("Feature Importance Ranking (Mean |SHAP|)")
            mean_abs_shap = np.abs(shap_vals).mean(axis=0)
            imp_df = pd.DataFrame(
                {"Feature": shap_feature_names, "Mean |SHAP|": mean_abs_shap}
            ).sort_values("Mean |SHAP|", ascending=False).head(20)

            fig_imp = px.bar(
                imp_df,
                x="Mean |SHAP|",
                y="Feature",
                orientation="h",
                color="Mean |SHAP|",
                color_continuous_scale="Viridis",
                title="Top 20 Features by SHAP Importance",
            )
            fig_imp.update_layout(template="plotly_white", yaxis=dict(autorange="reversed"))
            st.plotly_chart(fig_imp, use_container_width=True)

            # SHAP Dependence Plot
            st.subheader("SHAP Dependence Plot")
            dep_feature = st.selectbox(
                "Select feature for dependence plot",
                options=shap_feature_names,
                index=shap_feature_names.index("bedrooms") if "bedrooms" in shap_feature_names else 0,
                key="shap_dep_feature",
            )
            dep_idx = shap_feature_names.index(dep_feature)
            fig_dep, ax_dep = plt.subplots(figsize=(9, 5))
            shap.dependence_plot(
                dep_idx,
                shap_vals,
                X_shap_sample,
                feature_names=shap_feature_names,
                ax=ax_dep,
                show=False,
            )
            plt.tight_layout()
            st.pyplot(fig_dep)
            plt.close(fig_dep)

# ===================================================================
# TAB 7 - Price Decomposition
# ===================================================================
with tab7:
    st.header("Price Decomposition")
    st.markdown("Select a listing to see its predicted price breakdown into component contributions.")

    gbm_model_d, gbm_features_d, gbm_mdf_d = fit_gbm_model()

    if gbm_model_d is None:
        st.error("Prediction model not available.")
    else:
        # Build listing selector from rows with valid prices
        decomp_df = raw_df.dropna(subset=["price_numeric"]).copy()
        decomp_df = decomp_df[decomp_df["price_numeric"] > 0]
        decomp_df["label"] = (
            decomp_df["name"].str[:60]
            + " | "
            + decomp_df["city"]
            + " | $"
            + decomp_df["price_numeric"].round(0).astype(int).astype(str)
        )

        selected_label = st.selectbox(
            "Select a listing",
            options=decomp_df["label"].tolist()[:500],
            index=0,
            key="decomp_listing",
        )

        sel_row = decomp_df[decomp_df["label"] == selected_label].iloc[0]

        # Prepare features for this listing
        mdf_full = _prepare_model_df(raw_df)
        city_cols_d = sorted([c for c in gbm_features_d if c.startswith("city_")])
        ptype_cols_d = sorted([c for c in gbm_features_d if c.startswith("ptype_")])

        listing_features = {}
        for f in FEATURE_COLS:
            val = sel_row[f]
            if f in BOOL_FEATURES:
                listing_features[f] = int(val)
            else:
                listing_features[f] = float(val) if pd.notna(val) else mdf_full[f].median()

        for c in city_cols_d:
            city_name = c.replace("city_", "")
            listing_features[c] = 1 if sel_row["city"] == city_name else 0
        for p in ptype_cols_d:
            ptype_name = p.replace("ptype_", "")
            listing_features[p] = 1 if sel_row["property_type"] == ptype_name else 0

        X_listing = np.array([[listing_features.get(f, 0) for f in gbm_features_d]])
        log_pred_d = gbm_model_d.predict(X_listing)[0]
        pred_price_d = np.exp(log_pred_d)
        actual_price_d = sel_row["price_numeric"]

        # Metrics
        dc1, dc2, dc3 = st.columns(3)
        with dc1:
            st.metric("Actual Price", f"${actual_price_d:,.0f}")
        with dc2:
            st.metric("Predicted Price", f"${pred_price_d:,.0f}")
        with dc3:
            diff_pct = ((pred_price_d - actual_price_d) / actual_price_d) * 100
            st.metric("Prediction Error", f"{diff_pct:+.1f}%")

        st.markdown("---")

        # Decomposition: compute contribution of each feature group
        baseline_d = {}
        for f in gbm_features_d:
            baseline_d[f] = mdf_full[f].median() if f in mdf_full.columns else 0
        X_base_d = np.array([[baseline_d.get(f, 0) for f in gbm_features_d]])
        base_log_d = gbm_model_d.predict(X_base_d)[0]
        base_price_d = np.exp(base_log_d)

        feature_groups_d = {
            "Bedrooms": ["bedrooms"],
            "Bathrooms": ["bathrooms"],
            "Guests": ["guests_count"],
            "Amenities": ["amenities_count"],
            "Images": ["num_images"],
            "Rating": ["rating"],
            "Superhost": ["is_superhost"],
            "Guest Favorite": ["is_guest_favorite"],
            "Pets Allowed": ["pets_allowed"],
            "City": city_cols_d,
            "Property Type": ptype_cols_d,
        }

        contribs_d = {}
        for group_name, cols in feature_groups_d.items():
            X_mod = X_listing.copy()
            for col_name in cols:
                if col_name in gbm_features_d:
                    idx = gbm_features_d.index(col_name)
                    X_mod[0, idx] = baseline_d.get(col_name, 0)
            mod_log = gbm_model_d.predict(X_mod)[0]
            contribs_d[group_name] = (log_pred_d - mod_log)

        contribs_d["Baseline"] = base_log_d

        # Convert to dollar contributions (approximate)
        total_log = sum(contribs_d.values())
        scale = pred_price_d / np.exp(total_log) if total_log != 0 else 1

        dollar_contribs = {}
        for k, v in contribs_d.items():
            if k == "Baseline":
                dollar_contribs[k] = base_price_d
            else:
                dollar_contribs[k] = v * pred_price_d

        dc_df = pd.DataFrame(
            {"Component": list(dollar_contribs.keys()), "Contribution ($)": list(dollar_contribs.values())}
        )

        col_decomp1, col_decomp2 = st.columns(2)

        with col_decomp1:
            st.subheader("Price Component Breakdown")
            non_baseline = dc_df[dc_df["Component"] != "Baseline"].sort_values(
                "Contribution ($)", ascending=True
            )
            fig_stack = px.bar(
                non_baseline,
                x="Contribution ($)",
                y="Component",
                orientation="h",
                color="Contribution ($)",
                color_continuous_scale="RdYlGn",
                title="Feature Contributions to Predicted Price",
            )
            fig_stack.update_layout(template="plotly_white", height=450)
            st.plotly_chart(fig_stack, use_container_width=True)

        with col_decomp2:
            st.subheader("Top Drivers")

            positive = non_baseline[non_baseline["Contribution ($)"] > 0].sort_values(
                "Contribution ($)", ascending=False
            )
            negative = non_baseline[non_baseline["Contribution ($)"] < 0].sort_values(
                "Contribution ($)", ascending=True
            )

            st.markdown("**Top Positive Drivers**")
            if len(positive) > 0:
                for _, row in positive.head(5).iterrows():
                    st.write(f"  + **{row['Component']}**: +${row['Contribution ($)']:,.0f}")
            else:
                st.write("  None")

            st.markdown("**Top Negative Drivers**")
            if len(negative) > 0:
                for _, row in negative.head(5).iterrows():
                    st.write(f"  - **{row['Component']}**: -${abs(row['Contribution ($)']):,.0f}")
            else:
                st.write("  None")

            st.markdown("---")
            st.markdown(f"**Baseline price (median listing):** ${base_price_d:,.0f}")
            st.markdown(f"**Predicted price for this listing:** ${pred_price_d:,.0f}")