import streamlit as st
import pandas as pd
from datetime import datetime
import joblib
from sklearn.base import BaseEstimator, TransformerMixin
from datetime import datetime
from transformers import SugarContentReplacer,StoreAgeCalculator  # Import the custom transformer

# Load the trained model
def load_model():
    return joblib.load("src/SuperKart_sales_prediction_model_v1_0.joblib")

model = load_model()

# Streamlit UI for Customer Churn Prediction
st.title("Sales Prediction App")
st.write("This tool predicts customer Sales details. Enter the required information below.")

# Input fields for product and store data based on SuperKart dataset features
product_weight = st.number_input("Product Weight", min_value=0.0, value=12.66)
product_sugar_content = st.selectbox("Product Sugar Content", ["Low Sugar", "Regular", "No Sugar"])
product_allocated_area = st.number_input("Product Allocated Area", min_value=0.0, value=0.027)
product_type = st.selectbox("Product Type", ['Baking Goods', 'Breads', 'Breakfast', 'Canned', 'Dairy', 'Frozen Foods', 'Fruits and Vegetables', 'Hard Drinks', 'Health and Hygiene', 'Household', 'Meat', 'Others', 'Seafood', 'Snack Foods', 'Soft Drinks', 'Starchy Foods'])
product_mrp = st.number_input("Product MRP", min_value=0.0, value=117.08)
store_id = st.selectbox("Store ID", ['OUT001', 'OUT002', 'OUT003', 'OUT004'])
store_establishment_year = st.number_input("Store Establishment Year", min_value=1985, max_value=datetime.now().year, value=2009)
store_size = st.selectbox("Store Size", ["Small", "Medium", "High"])
store_location_city_type = st.selectbox("Store Location City Type", ["Tier 1", "Tier 2", "Tier 3"])
store_type = st.selectbox("Store Type", ["Departmental Store", "Supermarket Type1", "Supermarket Type2", "Food Mart"])

# Convert categorical inputs to match model training
input_data = {
    'Product_Weight': [product_weight],
    'Product_Sugar_Content': [product_sugar_content],
    'Product_Allocated_Area': [product_allocated_area],
    'Product_Type': [product_type],
    'Product_MRP': [product_mrp],
    'Store_Id': [store_id],
    'Store_Establishment_Year': [store_establishment_year],
    'Store_Size': [store_size],
    'Store_Location_City_Type': [store_location_city_type],
    'Store_Type': [store_type],
}

# Convert the input data to a DataFrame
input_df = pd.DataFrame(input_data)

# Custom transformer to replace 'reg' with 'Regular' in Product_Sugar_Content
class SugarContentReplacer(BaseEstimator, TransformerMixin):
    def fit(self, input_df, y=None):
        return self

    def transform(self, input_df):
        input_df = input_df.copy()
        input_df['Product_Sugar_Content'] = input_df['Product_Sugar_Content'].replace('reg', 'Regular')
        return input_df


# Convert categorical columns to category type
input_df['Product_Sugar_Content'] = input_df['Product_Sugar_Content'].astype('category')
input_df['Product_Type'] = input_df['Product_Type'].astype('category')
input_df['Store_Id'] = input_df['Store_Id'].astype('category')
input_df['Store_Size'] = input_df['Store_Size'].astype('category')
input_df['Store_Location_City_Type'] = input_df['Store_Location_City_Type'].astype('category')
input_df['Store_Type'] = input_df['Store_Type'].astype('category')


# Make predictions
if st.button("Predict"):
    predictions = model.predict(input_df)
    st.write(f"Prediction: The customer is likely to **{predictions[0]}**.")