import streamlit as st
import pandas as pd
import numpy as np
import joblib
import holidays
import xgboost as xgb
from datetime import date

# Load the model
@st.cache_resource
def load_artifacts():
    model = joblib.load('src/xgb_model.joblib')
    encoders = joblib.load('src/encoders.joblib')
    return model, encoders

try:
    model, encoders = load_artifacts()
except FileNotFoundError:
    st.error("Model files not found. Please upload 'xgb_model.joblib' and 'encoders.joblib'.")
    st.stop()

# Feature Engineering 
def create_features(df):
    df = df.copy()
    df['date'] = pd.to_datetime(df['date'])
    df['year'] = df['date'].dt.year
    df['month'] = df['date'].dt.month
    df['day'] = df['date'].dt.day
    df['day_of_week'] = df['date'].dt.dayofweek
    df['day_of_year'] = df['date'].dt.dayofyear
    df['week_of_year'] = df['date'].dt.isocalendar().week.astype(int)
    df['is_weekend'] = (df['day_of_week'] >= 5).astype(int)
    
    # Cyclical Encoding
    df['day_sin'] = np.sin(2 * np.pi * df['day_of_year'] / 365.0)
    df['day_cos'] = np.cos(2 * np.pi * df['day_of_year'] / 365.0)
    df['month_sin'] = np.sin(2 * np.pi * df['month'] / 12.0)
    df['month_cos'] = np.cos(2 * np.pi * df['month'] / 12.0)
    
    return df

def add_holiday_features(df, encoders):
    df = df.copy()
    # Inverse transform to get country names for holiday library
    df['country_name'] = encoders['country'].inverse_transform(df['country'])
    years = df['date'].dt.year.unique()
    country_codes = {'Canada': 'CA','Finland': 'FI','Italy': 'IT','Kenya': 'KE','Norway': 'NO','Singapore': 'SG'}
    
    df['is_holiday'] = 0
    
    for country_name, code in country_codes.items():
        try:
            country_holidays = holidays.Country(code, years=years)
            mask = (df['country_name'] == country_name) & (df['date'].isin(country_holidays))
            df.loc[mask, 'is_holiday'] = 1
        except:
            continue
    df = df.drop(columns=['country_name'])
    return df

# UI
st.title("🛒 Sticker Sales Prediction")
st.write("Enter the details below to predict the number of items sold.")

# Input
col1, col2 = st.columns(2)

with col1:
    country_options = list(encoders['country'].classes_)
    store_options = list(encoders['store'].classes_)
    product_options = list(encoders['product'].classes_)
    
    selected_date = st.date_input("Select Date", value=date(2025, 1, 1))
    selected_country = st.selectbox("Select Country", country_options)

with col2:
    selected_store = st.selectbox("Select Store", store_options)
    selected_product = st.selectbox("Select Product", product_options)

# 4. Prediction Button
if st.button("Predict Sales"):
    input_data = pd.DataFrame({
        'date': [pd.to_datetime(selected_date)],
        'country': [selected_country],
        'store': [selected_store],
        'product': [selected_product]})
    
    processed_data = create_features(input_data)
    categorical_cols = ['country', 'store', 'product']
    for col in categorical_cols:
        processed_data[col] = encoders[col].transform(processed_data[col])
    processed_data = add_holiday_features(processed_data, encoders)
    
    features = ['country', 'store', 'product', 'year', 'month', 'day', 'day_of_week', 'day_of_year', 'is_weekend', 
                'day_sin', 'day_cos', 'month_sin', 'month_cos', 'is_holiday']
    
    X_input = processed_data[features]
    
    # Prediction
    try:
        pred_log = model.predict(X_input)
        
        final_prediction = np.expm1(pred_log)[0]
        final_prediction = max(0, final_prediction) # Ensure no negative sales
        
        st.success(f"Predicted Num Sold: **{final_prediction:.2f}**")
        
        with st.expander("See processed features"):
            st.dataframe(X_input)
            
    except Exception as e:
        st.error(f"An error occurred during prediction: {e}")