import streamlit as st
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import StandardScaler, LabelEncoder

st.set_page_config(layout="wide")

# Add custom CSS for background image and styling
# Add custom CSS for background image and styling
st.markdown("""
    <style>
        .stApp {
            background-image: url("https://cdn.pixabay.com/photo/2020/01/28/11/14/galaxy-4799471_1280.jpg");
            background-size: cover;
            background-position: center;
            background-repeat: no-repeat;
            height: auto;  /* Allows the page to expand for scrolling */
            overflow: auto;  /* Enables scrolling if the page content overflows */
        }

        /* Adjust opacity of overlay to make content more visible */
        .stApp::before {
            content: "";
            position: absolute;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            background-color: rgba(255, 255, 255, 0);  /* Slightly higher opacity */
            z-index: 0;
        }

        /* Ensure content appears above the overlay */
        .stApp > * {
            position: relative;
            z-index: 1;
        }

        /* Ensure the dataframe is visible */
        .dataframe {
            background-color: rgba(255, 255, 255, 0.9) !important;
            z-index: 2;
        }

        /* Style text elements for better visibility */
        h1, h3, span, div {
            text-shadow: 1px 1px 2px rgba(255, 255, 255, 0.2);
        }
    </style>
""", unsafe_allow_html=True)


# Custom title styling functions
def colored_title(text, color):
    st.markdown(f"<h1 style='color: {color};'>{text}</h1>", unsafe_allow_html=True)

def colored_subheader(text, color):
    st.markdown(f"<h3 style='color: {color};'>{text}</h3>", unsafe_allow_html=True)

def colored_text(text, color):
    st.markdown(f"<span style='color: {color};'>{text}</span>", unsafe_allow_html=True)

class ANNModel(nn.Module):
    def __init__(self, input_size):
        super(ANNModel, self).__init__()
        self.fc1 = nn.Linear(input_size, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

@st.cache_resource
def load_model():
    _, X_scaled, _ = load_and_preprocess_data()
    input_size = X_scaled.shape[1]
    
    model = ANNModel(input_size=input_size)
    try:
        state_dict = torch.load('model_weights.pth', map_location=torch.device('cpu'))
        model.load_state_dict(state_dict)
        model.eval()
        return model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

@st.cache_data
def load_and_preprocess_data():
    df = pd.read_csv('Life Expectancy Data.csv')
    df.columns = df.columns.str.strip()

    df_display = df.copy()

    expected_features = [
        'Adult Mortality', 'infant deaths', 'Alcohol', 'percentage expenditure', 'Hepatitis B', 
        'Measles', 'BMI', 'under-five deaths', 'Polio', 'Total expenditure', 
        'Diphtheria', 'HIV/AIDS', 'GDP', 'Population', 'thinness  1-19 years', 
        'thinness 5-9 years', 'Income composition of resources', 'Schooling',
        'Country', 'Status', 'Year'
    ]

    for feature in expected_features:
        if feature not in df.columns:
            st.warning(f"Missing column '{feature}' - Creating with default values")
            df[feature] = 0
            df_display[feature] = 0

    for column in df.columns:
        if df[column].dtype == 'object':
            fill_value = df[column].mode()[0]
            df[column].fillna(fill_value, inplace=True)
            df_display[column].fillna(fill_value, inplace=True)
        else:
            fill_value = df[column].median()
            df[column].fillna(fill_value, inplace=True)
            df_display[column].fillna(fill_value, inplace=True)

    label_encoders = {}
    categorical_cols = ['Country', 'Status']
    for col in categorical_cols:
        le = LabelEncoder()
        df[col] = le.fit_transform(df[col].astype(str))
        label_encoders[col] = le

    X = df[expected_features]
    y = df['Life expectancy']

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    return df_display, X_scaled, y

def main():
    colored_title("Life Expectancy Estimation", "yellow")

    df_display, X_scaled, y = load_and_preprocess_data()

    colored_subheader("Dataset Preview", "yellow")
    st.dataframe(df_display.head())

    colored_subheader("Select a Row for Prediction:", "yellow")
    colored_text("Select a Row Index", "red")
    selected_row_index = st.selectbox("", options=range(len(df_display)), index=0, label_visibility="collapsed")

    predict_button = st.button("Predict Life Expectancy")

    if predict_button:
        model = load_model()
        if model is not None:
            row_to_predict = X_scaled[selected_row_index].reshape(1, -1)
            row_tensor = torch.tensor(row_to_predict, dtype=torch.float32)

            try:
                with torch.no_grad():
                    prediction = model(row_tensor).item()
                colored_subheader("Prediction Results:", "yellow")
                colored_text(f"Predicted Life Expectancy: {prediction:.2f} years", "red")
            except RuntimeError as e:
                st.error(f"Prediction error: {str(e)}")
                st.write("Input shape:", row_tensor.shape)
                st.write("Expected shape: 1x21")

            # Display the selected row with original categorical values
            colored_subheader("Selected Row for Prediction:", "yellow")
            st.write(df_display.iloc[selected_row_index])  
            
if __name__ == "__main__":
    main()