# streamlit_cbr_obesitas_fixed.py
import streamlit as st
import pandas as pd
import numpy as np
import time
from ucimlrepo import fetch_ucirepo
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics.pairwise import euclidean_distances
from sklearn.neighbors import KDTree

# ------------------- Utilities & Loading -------------------

@st.cache_data
def load_data():
    """
    Fetch the UCI repository dataset (id=544 as in your original code).
    Returns a raw dataframe (with categorical columns as object dtype).
    """
    dataset = fetch_ucirepo(id=544)
    X = dataset.data.features
    y = dataset.data.targets
    df = pd.concat([X, y], axis=1)
    return df

def safe_label_transform(le: LabelEncoder, series: pd.Series) -> pd.Series:
    """
    Transform a pandas Series using LabelEncoder but handle unseen labels by mapping them to -1.
    """
    # Map values to encoder classes index where possible
    mapping = {val: idx for idx, val in enumerate(le.classes_)}
    return series.map(lambda v: mapping.get(v, -1)).astype(int)

# ------------------- Session State Initialization -------------------
if 'case_base_raw' not in st.session_state:
    # Load raw dataset
    df_original = load_data()
    st.session_state.case_base_raw = df_original.copy()

    # Prepare label encoders for categorical columns and create encoded copy
    label_encoders = {}
    df_encoded_init = df_original.copy()
    for col in df_encoded_init.select_dtypes(include='object').columns:
        le = LabelEncoder()
        df_encoded_init[col] = le.fit_transform(df_encoded_init[col])
        label_encoders[col] = le

    st.session_state.label_encoders = label_encoders
    st.session_state.case_base_encoded = df_encoded_init.copy()

    # Keep list of which columns were categorical originally
    st.session_state.original_categorical_cols = list(df_original.select_dtypes(include='object').columns)
    st.session_state.n_initialized = True

# Convenience references
label_encoders = st.session_state.label_encoders
original_categorical_cols = st.session_state.original_categorical_cols

# ------------------- Encoding / Preprocessing Helpers -------------------
def encode_df(df_raw: pd.DataFrame) -> pd.DataFrame:
    """
    Encode a raw dataframe using stored label_encoders.
    Unseen categorical values are mapped to -1.
    """
    df = df_raw.copy()
    for col in df.select_dtypes(include='object').columns:
        if col in label_encoders:
            df[col] = safe_label_transform(label_encoders[col], df[col])
        else:
            # If a column is object but we don't have an encoder, attempt to infer as numeric or treat as unknown.
            try:
                df[col] = pd.to_numeric(df[col])
            except Exception:
                # fallback map unique values to integer codes
                df[col] = df[col].astype('category').cat.codes
    return df

def rebuild_feature_space():
    """
    Build feature matrix (encoded), scaler, scaled features and KDTree.
    Store them into st.session_state to avoid rebuilds unless case base changes.
    """
    case_base_encoded = st.session_state.case_base_encoded
    # Features = all columns except target
    features = case_base_encoded.drop(columns=['NObeyesdad'])
    target = case_base_encoded['NObeyesdad']

    # numerical and categorical columns in the encoded dataframe:
    # We'll consider originally categorical columns (stored) as categorical,
    # but in encoded df they are integers. Numerical columns are the rest.
    numerical_cols = [c for c in features.columns if c not in original_categorical_cols]
    categorical_cols = [c for c in features.columns if c in original_categorical_cols]

    # Scaler fit on encoded numeric data (note: categorical columns remain as integer labels)
    scaler = StandardScaler()
    # For scaling we'll only scale numerical columns and leave categorical intact,
    # but KDTree and Euclidean will be computed on a combined array where categorical columns are NOT scaled.
    # To keep behavior consistent with existing approach we will scale only numerical cols and concatenate.
    if len(numerical_cols) > 0:
        numeric_part = scaler.fit_transform(features[numerical_cols])
    else:
        numeric_part = np.zeros((len(features), 0))

    # categorical part as-is (integers)
    if len(categorical_cols) > 0:
        cat_part = features[categorical_cols].to_numpy(dtype=float)
    else:
        cat_part = np.zeros((len(features), 0))

    combined = np.hstack([numeric_part, cat_part])

    # build KDTree on combined (numerical scaled + categorical integers)
    tree = KDTree(combined)

    # For normalization of euclidean similarity, compute observed max distance within case base
    if combined.shape[0] > 1:
        # compute pairwise distances (this is O(n^2) but dataset is usually small)
        pairwise = euclidean_distances(combined, combined)
        max_observed_euclid = float(np.max(pairwise))
        if max_observed_euclid == 0:
            max_observed_euclid = 1.0
    else:
        max_observed_euclid = 1.0

    # For HEOM maximum possible distance:
    # numerical part: each numeric normalized difference <= 1 => sum of ones over numerical cols
    # categorical part: each mismatch contributes 1, so sum across categorical cols
    max_heom = np.sqrt(len(numerical_cols) + len(categorical_cols)) if (len(numerical_cols)+len(categorical_cols))>0 else 1.0

    # Save to session_state
    st.session_state.features = features
    st.session_state.numerical_cols = numerical_cols
    st.session_state.categorical_cols = categorical_cols
    st.session_state.scaler = scaler
    st.session_state.features_combined = combined
    st.session_state.features_scaled_combined = combined  # name kept for clarity
    st.session_state.kd_tree = tree
    st.session_state.max_observed_euclid = max_observed_euclid
    st.session_state.max_heom = max_heom

# initial build
rebuild_feature_space()

# ------------------- HEOM -------------------
def heom_distance(x1: pd.Series, x2: pd.Series, numerical_cols, categorical_cols, ranges):
    """
    x1, x2: pandas Series (encoded values) representing a single row each
    numerical_cols: list of numerical column names (these are encoded numeric columns)
    categorical_cols: list of originally categorical column names (also encoded as ints)
    ranges: dict of ranges for numerical columns (max-min from encoded data)
    """
    dist_sq = 0.0
    # numeric part: normalized squared difference (range normalization)
    for col in numerical_cols:
        r = ranges.get(col, 0)
        if r > 0:
            diff = (float(x1[col]) - float(x2[col])) / r
            dist_sq += diff * diff
        else:
            # no range (constant column), contribution 0
            pass
    # categorical part: 0 if same encoded value, 1 if different
    for col in categorical_cols:
        dist_sq += 0.0 if int(x1[col]) == int(x2[col]) else 1.0
    return float(np.sqrt(dist_sq))

# ------------------- Adaptability Score -------------------
def calculate_adaptability_score(new_case_df_encoded: pd.DataFrame, case_base_df_encoded: pd.DataFrame):
    """
    Information score based on negative log2 probability of each feature value in the dataset.
    Both inputs must be encoded (numerical & categorical encoded as ints).
    """
    info_score = 0.0
    epsilon = 1e-9
    # use columns of new_case
    for col in new_case_df_encoded.columns:
        freq = case_base_df_encoded[col].value_counts(normalize=True)
        new_val = new_case_df_encoded.iloc[0][col]
        p = freq.get(new_val, epsilon)
        info_score += -np.log2(p)
    return float(info_score)

@st.cache_data
def get_adaptive_threshold(case_base_df_encoded: pd.DataFrame, percentile=50):
    """
    Compute adaptive threshold across entire encoded case base (drop target column before passing).
    Cached to speed up repeated calls.
    """
    all_scores = [
        calculate_adaptability_score(pd.DataFrame([row]), case_base_df_encoded)
        for _, row in case_base_df_encoded.iterrows()
    ]
    return float(np.percentile(all_scores, percentile))

# ------------------- Retain Case -------------------
def retain_case(new_case_raw, distance_threshold=0.5, adaptive_percentile=50, force=False):
    """
    Evaluate whether to retain new_case_raw (raw dict or row) into case base.
    Returns: retain_flag, min_heom_distance, adaptability_score, used_adaptive_threshold
    """
    # Encode incoming case
    new_case_raw_df = pd.DataFrame([new_case_raw])
    new_case_encoded = encode_df(new_case_raw_df)

    case_base_encoded = st.session_state.case_base_encoded
    features = st.session_state.features  # encoded features dataframe (without target)
    numerical_cols = st.session_state.numerical_cols
    categorical_cols = st.session_state.categorical_cols

    # compute ranges for numerical columns using encoded features (max - min)
    feature_ranges = {col: (features[col].max() - features[col].min()) if col in numerical_cols else 0.0
                      for col in features.columns}

    # compute HEOM distances to all existing cases (encoded)
    raw_input = new_case_encoded.iloc[0]
    heom_distances = [
        heom_distance(raw_input, row, numerical_cols, categorical_cols, feature_ranges)
        for _, row in features.iterrows()
    ]
    min_dist = float(np.min(heom_distances)) if len(heom_distances) > 0 else float('inf')

    # adaptability score computed on encoded features only (drop target)
    adaptability_score = calculate_adaptability_score(new_case_encoded, case_base_encoded.drop(columns=['NObeyesdad']))

    adaptive_threshold = get_adaptive_threshold(case_base_encoded.drop(columns=['NObeyesdad']), percentile=adaptive_percentile)

    retain_flag = ((min_dist > distance_threshold) and (adaptability_score < adaptive_threshold)) or force

    return retain_flag, min_dist, adaptability_score, adaptive_threshold

# ------------------- Distance -> Similarity -------------------
def distance_to_similarity(distance, max_distance):
    """
    Normalize distance to [0,1] similarity where similarity = 1 - (distance / max_distance).
    Clipped to [0,1]. Adds tiny eps to denominator to avoid division by zero.
    """
    denom = max_distance + 1e-9
    sim = 1.0 - (distance / denom)
    return float(max(0.0, min(1.0, sim)))

# ------------------- Case-Based Reasoning (search) -------------------
def case_based_reasoning(new_input_raw, k_candidates=50):
    """
    Run CBR similarity searches (Euclidean, KD-Tree, HEOM, HEOM+KD-Tree hybrid) using encoded data.
    Returns dictionary of results and prints similarity & timing to Streamlit.
    """
    # Encode the input
    new_input_df_raw = pd.DataFrame([new_input_raw])
    new_input_encoded = encode_df(new_input_df_raw)

    # prepare encoded feature space from session
    features = st.session_state.features  # encoded features (without target)
    combined = st.session_state.features_combined
    tree = st.session_state.kd_tree
    scaler = st.session_state.scaler
    max_euclid = st.session_state.max_observed_euclid
    max_heom = st.session_state.max_heom
    numerical_cols = st.session_state.numerical_cols
    categorical_cols = st.session_state.categorical_cols

    # Build combined vector for the new input (scale numeric cols, keep categorical ints)
    if len(numerical_cols) > 0:
        numeric_part = scaler.transform(new_input_encoded[numerical_cols])
    else:
        numeric_part = np.zeros((1, 0))
    if len(categorical_cols) > 0:
        cat_part = new_input_encoded[categorical_cols].to_numpy(dtype=float)
    else:
        cat_part = np.zeros((1, 0))
    combined_input = np.hstack([numeric_part, cat_part])

    result = {}

    # ========== Euclidean (brute-force) ==========
    start_time = time.time()
    eucl_dists = euclidean_distances(combined_input, combined)
    eucl_idx = int(np.argmin(eucl_dists[0]))
    eucl_distance = float(eucl_dists[0][eucl_idx])
    eucl_time = time.time() - start_time
    eucl_similarity = distance_to_similarity(eucl_distance, max_euclid)

    st.write(f"🧮 Euclidean Distance: `{eucl_distance:.6f}`")
    st.write(f"🔗 Euclidean Similarity (1 - norm): `{eucl_similarity:.4f}`")
    st.write(f"⏱️ Euclidean Search Time: `{eucl_time:.6f}` seconds")

    result['euclidean'] = {"index": eucl_idx, "distance": eucl_distance, "similarity": eucl_similarity}

    # ========== KD-Tree (fast Euclidean) ==========
    start_time = time.time()
    kd_dist, kd_idx = tree.query(combined_input, k=1)
    kd_idx = int(kd_idx[0][0])
    kd_distance = float(kd_dist[0][0])
    kd_time = time.time() - start_time
    kd_similarity = distance_to_similarity(kd_distance, max_euclid)

    st.write(f"🧮 KD-Tree Distance: `{kd_distance:.6f}`")
    st.write(f"🔗 KD-Tree Similarity: `{kd_similarity:.4f}`")
    st.write(f"⏱️ KD-Tree Search Time: `{kd_time:.6f}` seconds")

    result['kdtree'] = {"index": kd_idx, "distance": kd_distance, "similarity": kd_similarity}

    # ========== HEOM (brute-force) ==========
    # compute ranges for numerical columns
    ranges = {col: (features[col].max() - features[col].min()) if col in numerical_cols else 0.0 for col in features.columns}
    start_time = time.time()
    heom_distances = [
        heom_distance(new_input_encoded.iloc[0], row, numerical_cols, categorical_cols, ranges)
        for _, row in features.iterrows()
    ]
    heom_idx = int(np.argmin(heom_distances))
    heom_distance_value = float(heom_distances[heom_idx])
    heom_time = time.time() - start_time
    heom_similarity = distance_to_similarity(heom_distance_value, max_heom)

    st.write(f"🧮 HEOM Distance: `{heom_distance_value:.6f}`")
    st.write(f"🔗 HEOM Similarity: `{heom_similarity:.4f}`")
    st.write(f"⏱️ HEOM Search Time: `{heom_time:.6f}` seconds")

    result['heom'] = {"index": heom_idx, "distance": heom_distance_value, "similarity": heom_similarity,
                      "case": st.session_state.case_base_raw.iloc[heom_idx].to_dict()}

    # ========== HEOM + KD-Tree Hybrid ==========
    # Query KDTree for k candidates (fast), then compute HEOM only on candidates
    k = min(k_candidates, combined.shape[0])
    kd_candidates_idx = tree.query(combined_input, k=k)[1][0]
    start_time = time.time()
    heom_candidate_dists = []
    for idx in kd_candidates_idx:
        row = features.iloc[idx]
        d = heom_distance(new_input_encoded.iloc[0], row, numerical_cols, categorical_cols, ranges)
        heom_candidate_dists.append((int(idx), float(d)))
    # find min among candidates
    heom_kdtree_idx, heom_kdtree_dist = min(heom_candidate_dists, key=lambda x: x[1])
    heom_kdtree_time = time.time() - start_time
    heom_kdtree_similarity = distance_to_similarity(heom_kdtree_dist, max_heom)

    st.write(f"🧮 HEOM (KD-Tree Hybrid) Distance: `{heom_kdtree_dist:.6f}`")
    st.write(f"🔗 HEOM (KD-Tree Hybrid) Similarity: `{heom_kdtree_similarity:.4f}`")
    st.write(f"⏱️ HEOM (KD-Tree Hybrid) Time: `{heom_kdtree_time:.6f}` seconds")

    result['heom_kdtree'] = {
        "index": int(heom_kdtree_idx),
        "distance": float(heom_kdtree_dist),
        "similarity": float(heom_kdtree_similarity),
        "case": st.session_state.case_base_raw.iloc[int(heom_kdtree_idx)].to_dict()
    }

    return result

# ------------------- Streamlit UI -------------------
st.title("🧠 CBR Obesitas + Retain Adaptif (Fixed)")
st.markdown(f"Jumlah kasus dalam database saat ini: **{len(st.session_state.case_base_raw)} kasus**")

# USER INPUT
user_input = {
    'Gender': st.selectbox("Gender", ['Male', 'Female']),
    'Age': st.number_input("Age", 10, 100, 25),
    'Height': st.number_input("Height (in meters)", 1.0, 2.5, 1.70, format="%.2f"),
    'Weight': st.number_input("Weight (in kg)", 30.0, 200.0, 70.0, format="%.1f"),
    'family_history_with_overweight': st.selectbox("Family history with overweight", ['yes', 'no']),
    'FAVC': st.selectbox("Frequent consumption of high caloric food", ['yes', 'no']),
    'FCVC': st.slider("Vegetable consumption (0–3)", 0.0, 3.0, 2.0, step=0.5),
    'NCP': st.slider("Number of main meals", 1.0, 5.0, 3.0, step=1.0),
    'CAEC': st.selectbox("Food between meals", ['no', 'Sometimes', 'Frequently', 'Always']),
    'SMOKE': st.selectbox("Do you smoke?", ['yes', 'no']),
    'CH2O': st.slider("Daily water intake", 0.0, 3.0, 2.0, step=0.1),
    'SCC': st.selectbox("Calories monitoring", ['yes', 'no']),
    'FAF': st.slider("Physical activity (hrs/week)", 0.0, 5.0, 1.0, step=0.5),
    'TUE': st.slider("Tech usage (hrs/day)", 0.0, 5.0, 1.0, step=0.5),
    'CALC': st.selectbox("Alcohol consumption", ['no', 'Sometimes', 'Frequently', 'Always']),
    'MTRANS': st.selectbox("Transport", ['Automobile', 'Motorbike', 'Bike', 'Public_Transportation', 'Walking'])
}

force_save = st.checkbox("💾 Simpan paksa jika ditolak?", value=False)

if st.button("🔍 Prediksi dan Evaluasi Retain"):
    # run CBR
    tic = time.time()
    result = case_based_reasoning(user_input)
    toc = time.time()

    st.markdown("### 🔢 Kasus Paling Mirip (HEOM KD-Tree Hybrid)")
    st.json(result['heom_kdtree']['case'])

    # predicted label from matched case
    predicted_label = result['heom_kdtree']['case']['NObeyesdad']
    user_input_with_label = user_input.copy()
    user_input_with_label['NObeyesdad'] = predicted_label

    st.markdown("### 🧾 Prediksi Kategori Obesitas")
    st.write(f"🎯 **{predicted_label}** (berdasarkan kasus paling mirip)")

    # Evaluate retain
    retain_flag, min_dist, adaptability_score, threshold_used = retain_case(
        user_input, distance_threshold=0.5, adaptive_percentile=50, force=force_save
    )

    st.markdown("### 📥 Evaluasi Retain")
    st.write(f"HEOM Distance Terdekat: `{min_dist:.3f}`")
    st.write(f"Skor Adaptabilitas: `{adaptability_score:.2f}` (Threshold adaptif: `{threshold_used:.2f}`)")

    if retain_flag:
        # append to RAW
        st.session_state.case_base_raw = pd.concat(
            [st.session_state.case_base_raw, pd.DataFrame([user_input_with_label])],
            ignore_index=True
        )
        # append to ENCODED
        encoded_new = encode_df(pd.DataFrame([user_input_with_label]))
        st.session_state.case_base_encoded = pd.concat(
            [st.session_state.case_base_encoded, encoded_new],
            ignore_index=True
        )

        # rebuild feature space after insertion
        rebuild_feature_space()

        st.success(f"✅ Kasus berhasil disimpan. Jumlah kasus sekarang: {len(st.session_state.case_base_raw)}")
    else:
        st.info("❌ Kasus tidak disimpan otomatis (tidak adaptif & terlalu mirip).")
        if force_save:
            st.session_state.case_base_raw = pd.concat(
                [st.session_state.case_base_raw, pd.DataFrame([user_input_with_label])],
                ignore_index=True
            )
            encoded_new = encode_df(pd.DataFrame([user_input_with_label]))
            st.session_state.case_base_encoded = pd.concat(
                [st.session_state.case_base_encoded, encoded_new],
                ignore_index=True
            )
            rebuild_feature_space()
            st.warning(f"⚠️ Simpan paksa dilakukan. Jumlah kasus sekarang: {len(st.session_state.case_base_raw)}")

st.markdown("---")
st.caption("Implementasi: HEOM uses encoded values only; similarity scaled to [0,1] via 1 - (distance / max).")