import streamlit as st
from streamlit_shap import st_shap
import shap
from datasets import load_dataset
from sklearn.model_selection import train_test_split
import lightgbm as lgb
import numpy as np
import pandas as pd


@st.experimental_memo
def load_data():
    dataset = load_dataset("ttd22/house-price", streaming = True)
    df = pd.DataFrame.from_dict(dataset["train"])
    df = df.drop('Id', axis=1)
    drop_columns = (df.isnull().sum().sort_values(ascending=False).loc[lambda x : x > .90*1460]).index.to_list()
    df = df.drop(drop_columns, axis = 'columns', errors = 'ignore')
    cols_with_missing_values = df.columns[df.isnull().sum() > 0]
    # Iterate through each column with missing values
    for col in cols_with_missing_values:
        # Check if the column is numeric
        if df[col].dtype in ['int64', 'float64']:
            # Impute missing values with median
            median = df[col].median()
            df[col].fillna(median, inplace=True)
        else:
            # Impute missing values with mode
            mode = df[col].mode()[0]
            df[col].fillna(mode, inplace=True)
    X, y = df.drop("SalePrice", axis=1), df["SalePrice"]
    # Extract categoricals and their indices
    cat_features = X.select_dtypes(exclude=np.number).columns.to_list()
    cat_idx = [X.columns.get_loc(col) for col in cat_features]
    # Convert cat_features to pd.Categorical dtype
    for col in cat_features:
        X[col] = pd.Categorical(X[col])
    return X,y,cat_idx

@st.experimental_memo
def load_model(X, y, cat_idx):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    params = {'n_estimators': 569, 'num_leaves': 62, 'max_depth': 10, 'learning_rate': 0.010786783375710743, 'colsample_bytree': 0.5065493231651268, 'subsample': 0.7900705177300663, 'lambda_l1': 4.998785478697207, 'lambda_l2': 2.1857959934319657, 'min_child_weight': 11.187719709451862}
    model = lgb.LGBMRegressor(**params)
    model.fit(X_train, y_train, eval_set=[(X_test, y_test)], categorical_feature=cat_idx, verbose = False)
    return model

# train LightGBM model
X,y,cat_idx = load_data()
model = load_model(X, y, cat_idx)