# inference_core.py
import pandas as pd
import numpy as np
import joblib

from fe_avm_bdo26 import (
    map_kecamatan,
    transform_area_group
)

import os

# ============================
# OPTIONAL GEMINI CLIENT
# ============================
gemini_client = None
GEMINI_ENABLED = False

try:
    from google import genai

    api_key = os.getenv("GEMINI_API_KEY")

    if api_key:
        gemini_client = genai.Client(api_key=api_key)
        GEMINI_ENABLED = True
        print("Gemini inference summary enabled ✅")
    else:
        print("No GEMINI_API_KEY found — summary fallback mode.")

except Exception as e:
    print("Gemini init error:", e)
    gemini_client = None
    GEMINI_ENABLED = False
    
# =========================
# LOAD BUNDLES

avm_bundle = joblib.load("avm_bdo26_bundle_fixed.pkl")
recc_bundle = joblib.load("recc_bdo26_bundle.pkl")
listing_df = pd.read_csv("bandung_listing_features.csv")

model = avm_bundle["model"]


# =========================
# PREDICTION
# =========================
def prepare_input(user_info: dict):
    location = user_info.get("location")

    building_default = 100
    land_default = 120
    default_source = "Global Default"

    kec_name = None

    if location:
        # map location -> kecamatan
        temp_df = pd.DataFrame([{
            "location": location
        }])

        temp_df = map_kecamatan(temp_df)

        kec_name = temp_df.iloc[0]["kecamatan_mapped"]

        # -------------------------
        # 1) exact location median
        # -------------------------
        area_df = listing_df[
            listing_df["location"] == location
        ]

        if len(area_df) >= 10:
            building_default = area_df[
                "building_area_m2"
            ].median()

            land_default = area_df[
                "land_area_m2"
            ].median()

            default_source = (
                f"Location Median ({location})"
            )

        # -------------------------
        # 2) kec median fallback
        # -------------------------
        elif kec_name is not None:
            kec_df = listing_df[
                listing_df["kecamatan_mapped"] == kec_name
            ]

            if len(kec_df) > 0:
                building_default = kec_df[
                    "building_area_m2"
                ].median()

                land_default = kec_df[
                    "land_area_m2"
                ].median()

                default_source = (
                    f"Kecamatan Median ({kec_name})"
                )

    defaults = {
        "building_type": "Rumah",
        "bedroom_total": 3,
        "bathroom_total": 2,
        "carport_total": 1,
        "land_area_m2": land_default,
        "building_area_m2": building_default,
        "size_default_source": default_source,
        "edu_count_3km": np.nan,
        "health_count_3km": np.nan,
        "ent_count_3km": np.nan,
        "transport_count_3km": np.nan,
        "nearest_transport_1_km": np.nan,
        "nearest_edu_1_km": np.nan,
        "nearest_health_1_km": np.nan,
        "nearest_ent_1_km": np.nan,
    }

    row = defaults.copy()

    for k, v in user_info.items():
        if v is not None and v != "":
            row[k] = v

    df = pd.DataFrame([row])

    df = map_kecamatan(df)
    df = transform_area_group(df)

    return df

def predict_price(user_info: dict):
    avm_input = user_info.copy()
    avm_input.pop("budget_million_idr", None)
    avm_input.pop("notes", None)

    df = prepare_input(avm_input)

    pred = model.predict(df)[0]

    # if model outputs log price
    if pred < 20:
        pred = np.expm1(pred)

    return float(pred)

def calibrate_prediction(pred_price, pct=0.12):
    lower = pred_price * (1 - pct)
    upper = pred_price * (1 + pct)

    return {
        "Predicted Price (IDR Mn)": round(pred_price, 1),
        "Lower Range": round(lower, 1),
        "Upper Range": round(upper, 1),
    }


# =========================
# TARGET AREA OVERVIEW
# =========================
def get_target_area_overview(location):
    area_df = listing_df[
        listing_df["location"] == location
    ].copy()

    if area_df.empty:
        return {
            "Target Area": location,
            "Typical Price": None,
            "Typical Size": None
        }

    typical_price = area_df[
        "price_million_idr"
    ].median()

    size_col = (
        "building_area_m2"
        if "building_area_m2" in area_df.columns
        else None
    )

    typical_size = (
        area_df[size_col].median()
        if size_col is not None
        else None
    )

    return {
        "Target Area": location,
        "Typical Price": round(typical_price, 1),
        "Typical Size": round(typical_size, 1)
        if typical_size is not None else None
    }


# =========================
# RECOMMENDER
# =========================
def recommend_locations(user_info, top_n=5):
    location_profile = recc_bundle["location_profile"].copy()
    similarity_cols = recc_bundle["similarity_cols"]
    artifacts = recc_bundle["artifacts"]

    knn_model = artifacts["knn_location"]
    scaler = artifacts["scaler_location"]
    weight_vector = artifacts["weight_vector"]

    target_area = user_info.get("location")

    target_row = location_profile[
        location_profile["location"] == target_area
    ].copy()

    if target_row.empty:
        budget = user_info.get("budget_million_idr")

        if budget is not None:
            location_profile["budget_gap"] = abs(
                location_profile["price_weighted"] - budget
            )
            target_row = (
                location_profile
                .sort_values("budget_gap")
                .head(1)
                .copy()
            )
        else:
            target_row = location_profile.head(1).copy()

    target_features = target_row[similarity_cols]

    # IMPORTANT: scale first, then apply weights
    target_scaled = scaler.transform(target_features)
    target_weighted = target_scaled * weight_vector

    distances, indices = knn_model.kneighbors(
        target_weighted,
        n_neighbors=top_n + 1
    )

    rec_rows = location_profile.iloc[indices[0]].copy()

    # attach raw distances so we can filter self safely
    rec_rows["__distance"] = distances[0]

    # remove self if exact target area is present
    rec_rows = rec_rows[
        rec_rows["location"] != target_area
    ].head(top_n).copy()

    # similarity score
    rec_rows["Similarity Score"] = (
        1 / (1 + rec_rows["__distance"])
    ).round(3)

    rec_rows = rec_rows.rename(columns={
        "location": "Alternative Area",
        "price_weighted": "Typical Price",
        "building_area_m2": "Typical Size",
        "access_quality_score": "Access Score",
        "location_premium_score": "Area Desirability Score"
    })

    round_cols = [
        "Typical Price",
        "Typical Size",
        "Access Score",
        "Area Desirability Score"
    ]

    for col in round_cols:
        if col in rec_rows.columns:
            rec_rows[col] = rec_rows[col].round(2)

    return rec_rows[[
        "Alternative Area",
        "Similarity Score",
        "Typical Price",
        "Typical Size",
        "Access Score",
        "Area Desirability Score"
    ]]

# =========================
# LISTINGS
# =========================
def get_target_listings(
    location,
    budget=None,
    predicted_price=None,
    top_n=3
):
    """
    Show listings in target area.

    Priority:
    1. use budget if user gives one
    2. fallback to predicted price
    3. fallback to median listing
    """

    area_df = listing_df[
        listing_df["location"] == location
    ].copy()

    if area_df.empty:
        return pd.DataFrame()

    # -------------------------
    # anchor price logic
    # -------------------------
    anchor_price = budget

    if anchor_price is None:
        anchor_price = predicted_price

    if anchor_price is None:
        anchor_price = area_df[
            "price_million_idr"
        ].median()

    # -------------------------
    # choose closest fit
    # -------------------------
    area_df["fit_gap"] = abs(
        area_df["price_million_idr"] - anchor_price
    )

    area_df = area_df.sort_values(
        ["fit_gap", "price_million_idr"]
    )

    keep_cols = [
        c for c in [
            "building_type",
            "location",
            "price_million_idr",
            "listing_summary",
            "url"
        ] if c in area_df.columns
    ]

    return area_df[keep_cols].head(top_n)

def get_recommended_area_examples(
    alt_areas,
    target_price=None,
    budget=None,
    per_area=2
):
    anchor_price = budget if budget is not None else target_price

    all_rows = []

    for area in alt_areas:
        area_df = listing_df[
            listing_df["location"] == area
        ].copy()

        if area_df.empty:
            continue

        if anchor_price is not None:
            area_df["fit_gap"] = abs(
                area_df["price_million_idr"] - anchor_price
            )

            # prioritize closest fit, not cheapest
            area_df = area_df.sort_values(
                ["fit_gap", "price_million_idr"]
            )
        else:
            median_price = area_df[
                "price_million_idr"
            ].median()

            area_df["fit_gap"] = abs(
                area_df["price_million_idr"] - median_price
            )

            area_df = area_df.sort_values("fit_gap")

        keep_cols = [
            c for c in [
                "building_type",
                "location",
                "price_million_idr",
                "listing_summary",
                "url"
            ] if c in area_df.columns
        ]

        all_rows.append(area_df[keep_cols].head(per_area))

    if not all_rows:
        return pd.DataFrame()

    return pd.concat(all_rows).reset_index(drop=True)

def build_comparison_table(target_overview, alternatives):
    target_row = pd.DataFrame([{
        "Area": target_overview["Target Area"],
        "Similarity Score": 1.000,
        "Typical Price": target_overview["Typical Price"],
        "Typical Size": target_overview["Typical Size"],
        "Access Score": np.nan,
        "Area Desirability Score": np.nan,
        "Category": "Target"
    }])

    alt_df = alternatives.copy()
    alt_df["Category"] = "Alternative"

    alt_df = alt_df.rename(columns={
        "Alternative Area": "Area"
    })

    return pd.concat(
        [target_row, alt_df],
        ignore_index=True
    )

def get_target_area_scores(location):
    """
    Pull access + desirability score from
    recommender profile
    """
    location_profile = recc_bundle["location_profile"].copy()

    row = location_profile[
        location_profile["location"] == location
    ]

    if row.empty:
        return {
            "Access Score": np.nan,
            "Desirability Score": np.nan
        }

    row = row.iloc[0]

    return {
        "Access Score": round(
            row.get("access_quality_score", np.nan), 2
        ),
        "Desirability Score": round(
            row.get("location_premium_score", np.nan), 2
        )
    }

# =========================
# SUMMARY
# =========================
def generate_summary(
    user_info,
    pred_result,
    rec_table
):
    target_area = user_info.get("location")
    target_price = pred_result[
        "Predicted Price (IDR Mn)"
    ]

    if rec_table is None or rec_table.empty:
        return (
            f"The estimated market value for "
            f"{target_area} is around "
            f"IDR {target_price:.0f} Mn."
        )

    top_rec = rec_table.iloc[0]

    alt_area = top_rec["Alternative Area"]
    alt_price = top_rec["Typical Price"]
    similarity = top_rec["Similarity Score"]

    pct_gap = 0

    if target_price and target_price != 0:
        pct_gap = (
            (target_price - alt_price)
            / target_price
        ) * 100

    if pct_gap > 0:
        price_text = (
            f"approximately {pct_gap:.0f}% "
            f"more affordable"
        )
    else:
        price_text = (
            f"approximately {abs(pct_gap):.0f}% "
            f"more premium"
        )

    return (
        f"If you like {target_area}, "
        f"you may want to consider "
        f"{alt_area}, which has a similar "
        f"profile (similarity {similarity:.2f}) "
        f"while being {price_text}. "
        f"The estimated market value in "
        f"{target_area} is around "
        f"IDR {target_price:.0f} Mn."
    )

def gemini_summary(user_info, pred_result, rec_table, client=None):
    fallback = generate_summary(
        user_info,
        pred_result,
        rec_table
    )

    active_client = client if client is not None else gemini_client

    if active_client is None:
        return fallback

    try:
        target_area = user_info.get("location")
        target_price = pred_result[
            "Predicted Price (IDR Mn)"
        ]
        budget = user_info.get("budget_million_idr")

        top_recs = rec_table.head(3).to_dict(
            orient="records"
        )

        prompt = f"""
Write short housing advice. consulting tone, client-facing. Focus on tradeoffs, affordability, and next best action.

Area: {target_area}
Predicted price: {target_price} Mn
Budget: {budget}
Alternatives: {top_recs}

Rules:
- 3 sentences max
- under 90 words
- start with best alternative IF target area is out of budget.
- say cheaper or premium
- mention 1 tradeoff
- end with action advice
- plain text only
"""

        model_candidates = [
            "gemini-2.5-flash-lite"
        ]
        
        for model_name in model_candidates:
            try:
                response = active_client.models.generate_content(
                    model=model_name,
                    contents=prompt,
                    config={
                        "max_output_tokens": 120,
                        "temperature": 0.3,
                        "top_p": 0.95
                    }
                )

                if hasattr(response, "text") and response.text:
                    text = response.text.strip()
                else:
                    text = str(response).strip()

                clean_text = text.strip()

                if (
                    clean_text
                    and len(clean_text.split()) >= 8
                    and clean_text.count(".") >= 1
                ):
                    print(f"GEMINI SUCCESS: {model_name}")
                    return clean_text

            except Exception as e:
                print(
                    f"GEMINI FAILED {model_name}: {e}"
                )

        return fallback

    except Exception as e:
        print("Gemini summary fallback:", e)
        return fallback
    
# =========================
# MASTER WRAPPER
# =========================
def run_full_inference(user_info, client=None):
    """
    Full decision-support inference pipeline
    """

    # =========================
    # 1. prediction
    # =========================
    pred_price = predict_price(user_info)

    pred_result = calibrate_prediction(pred_price)

    # =========================
    # 2. prepare input context
    # =========================
    prepared_df = prepare_input(user_info)

    params_used = {
        "Target Area": user_info.get("location"),
        "Building Type": prepared_df.iloc[0]["building_type"],
        "Bedrooms": prepared_df.iloc[0]["bedroom_total"],
        "Bathrooms": prepared_df.iloc[0]["bathroom_total"],
        "Carport": prepared_df.iloc[0]["carport_total"],
        "Building Area (m²)": round(
            prepared_df.iloc[0]["building_area_m2"], 1
        ),
        "Land Area (m²)": round(
            prepared_df.iloc[0]["land_area_m2"], 1
        ),
        "Budget (IDR Mn)": user_info.get(
            "budget_million_idr"
        ),
        "Size Default Source": prepared_df.iloc[0].get(
            "size_default_source",
            "User Input"
        )
    }

    # =========================
    # 3. target overview
    # =========================
    target_overview = get_target_area_overview(
        user_info.get("location")
    )

    # =========================
    # 4. alternatives
    # =========================
    alternatives = recommend_locations(
        user_info
    )

    # =========================
    # 5. comparison table
    # =========================
    target_scores = get_target_area_scores(
        user_info.get("location")
    )

    target_row = pd.DataFrame([{
        "Area": target_overview["Target Area"],
        "Similarity Score": 1.000,
        "Typical Price": target_overview["Typical Price"],
        "Typical Size": target_overview["Typical Size"],
        "Access Score": target_scores["Access Score"],
        "Desirability Score": target_scores["Desirability Score"]
    }])

    alt_df = alternatives.copy()

    alt_df = alt_df.rename(columns={
        "Alternative Area": "Area",
        "Area Desirability Score": "Desirability Score"
    })

    comparison_table = pd.concat(
        [target_row, alt_df],
        ignore_index=True
    ).round(2)
    
    # =========================
    # 6. target listings
    # =========================
    target_listings = get_target_listings(
        location=user_info.get("location"),
        budget=user_info.get(
            "budget_million_idr"
        ),
        predicted_price=pred_result[
            "Predicted Price (IDR Mn)"
        ]
    )

    # =========================
    # 7. recommended listings
    # =========================
    recommended_listings = get_recommended_area_examples(
        alt_areas=alternatives["Alternative Area"]
        .tolist(),
        target_price=pred_result[
            "Predicted Price (IDR Mn)"
        ],
        budget=user_info.get(
            "budget_million_idr"
        ),
        per_area=2
    )

    # =========================
    # 8. narrative market summary
    # =========================
    market_summary = f"""
### 📍 Market Context — {target_overview.get('Target Area', 'Selected Area')}

Typical houses matching your criteria in **{target_overview.get('Target Area', 'this area')}**
are usually priced around **IDR {pred_result['Predicted Price (IDR Mn)']:.0f} Mn**

**Conservative valuation range:**  
IDR {pred_result['Lower Range']:.0f} Mn – {pred_result['Upper Range']:.0f} Mn

For reference, the typical area baseline is approximately:

- **Typical area price:** IDR {target_overview.get('Typical Price', 0):.0f} Mn
- **Typical area size:** {target_overview.get('Typical Size', 0):.0f} m²
"""

    # =========================
    # 9. advisory summary
    # =========================
    summary = gemini_summary(
        user_info,
        pred_result,
        alternatives,
        client=client
    )

    # =========================
    # return
    # =========================
    return {
        "params_used": params_used,
        "prediction": pred_result,
        "target_overview": target_overview,
        "alternatives": alternatives,
        "comparison_table": comparison_table,
        "target_listings": target_listings,
        "recommended_listings": recommended_listings,
        "market_summary": market_summary,
        "summary": summary
    }