import os
# Disable OpenMP
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'

import streamlit as st
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import shap
from sklearn.preprocessing import MinMaxScaler
import plotly.graph_objects as go
import io
from matplotlib.figure import Figure

# Set page config
st.set_page_config(
    page_title="Friction Angle Predictor",
    page_icon="🔄",
    layout="wide"
)

# Custom CSS to improve the app's appearance
st.markdown("""
    <style>
    .stApp {
        max-width: 1200px;
        margin: 0 auto;
    }
    .main {
        padding: 2rem;
    }
    .stButton>button {
        width: 100%;
    }
    </style>
    """, unsafe_allow_html=True)

# Load the trained model and recreate the architecture
class Net(torch.nn.Module):
    def __init__(self, input_size):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(input_size, 64)
        self.fc2 = torch.nn.Linear(64, 1000)
        self.fc3 = torch.nn.Linear(1000, 200)
        self.fc4 = torch.nn.Linear(200, 8)
        self.fc5 = torch.nn.Linear(8, 1)
        self.dropout = torch.nn.Dropout(0.2)
        
        # Initialize weights
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        if isinstance(module, torch.nn.Linear):
            torch.nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                module.bias.data.zero_()
    
    def forward(self, x):
        x = torch.nn.functional.relu(self.fc1(x))
        x = self.dropout(x)
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.dropout(x)
        x = torch.nn.functional.relu(self.fc3(x))
        x = self.dropout(x)
        x = torch.nn.functional.relu(self.fc4(x))
        x = self.dropout(x)
        x = self.fc5(x)
        return x

@st.cache_resource
def load_model_and_data():
    # Set device and random seeds
    np.random.seed(32)
    torch.manual_seed(42)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Load data
    data = pd.read_excel("Data_syw.xlsx")
    X = data.iloc[:, list(range(1, 17)) + list(range(21, 23))]
    y = data.iloc[:, 28].values
    
    # Calculate correlation and select features
    correlation_with_target = abs(X.corrwith(pd.Series(y)))
    selected_features = correlation_with_target[correlation_with_target > 0.1].index
    X = X[selected_features]
    
    # Initialize and fit scalers
    scaler_X = MinMaxScaler()
    scaler_y = MinMaxScaler()
    scaler_X.fit(X)
    scaler_y.fit(y.reshape(-1, 1))
    
    # Load model
    model = Net(input_size=len(selected_features)).to(device)
    model.load_state_dict(torch.load('friction_model.pt'))
    model.eval()
    
    return model, X.columns, scaler_X, scaler_y, device, X

def predict_friction(input_values, model, scaler_X, scaler_y, device):
    # Scale input values
    input_scaled = scaler_X.transform(input_values)
    input_tensor = torch.FloatTensor(input_scaled).to(device)
    
    # Make prediction
    with torch.no_grad():
        prediction_scaled = model(input_tensor)
        prediction = scaler_y.inverse_transform(prediction_scaled.cpu().numpy().reshape(-1, 1))
    
    return prediction[0][0]

def calculate_shap_values(input_values, model, X, scaler_X, scaler_y, device):
    def model_predict(X):
        X_scaled = scaler_X.transform(X)
        X_tensor = torch.FloatTensor(X_scaled).to(device)
        with torch.no_grad():
            scaled_pred = model(X_tensor).cpu().numpy()
        return scaler_y.inverse_transform(scaled_pred.reshape(-1, 1)).flatten()
    
    try:
        # Use a smaller background dataset and fewer samples for stability
        background = shap.kmeans(X.values, k=5)  # Reduced from 10 to 5
        explainer = shap.KernelExplainer(model_predict, background)
        shap_values = explainer.shap_values(input_values.values, nsamples=100)  # Added nsamples parameter
        
        if isinstance(shap_values, list):
            shap_values = np.array(shap_values[0])
        
        return shap_values[0], explainer.expected_value
    except Exception as e:
        st.error(f"Error calculating SHAP values: {str(e)}")
        # Return dummy values in case of error
        return np.zeros(len(input_values.columns)), 0.0

def create_waterfall_plot(shap_values, feature_names, base_value, input_data):
    # Create SHAP explanation object
    explanation = shap.Explanation(
        values=shap_values,
        base_values=base_value,
        data=input_data,
        feature_names=list(feature_names)
    )
    
    # Create figure
    fig = plt.figure(figsize=(12, 8))
    shap.plots.waterfall(explanation, show=False)
    plt.title('Local SHAP Value Contributions')
    plt.tight_layout()
    
    # Save plot to a buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=300)
    plt.close(fig)
    buf.seek(0)
    return buf

def main():
    st.title("🔄 Friction Angle Predictor")
    st.write("This app predicts the friction angle based on waste composition and characteristics.")
    
    try:
        # Load model and data
        model, feature_names, scaler_X, scaler_y, device, X = load_model_and_data()
        
        # Create two columns for input
        col1, col2 = st.columns(2)
        
        # Dictionary to store input values
        input_values = {}
        
        # Create input fields for each feature
        for i, feature in enumerate(feature_names):
            with col1 if i < len(feature_names)//2 else col2:
                min_val = float(X[feature].min())
                max_val = float(X[feature].max())
                mean_val = float(X[feature].mean())
                
                input_values[feature] = st.number_input(
                    f"{feature}",
                    min_value=min_val,
                    max_value=max_val,
                    value=mean_val,
                    help=f"Range: {min_val:.2f} to {max_val:.2f}"
                )
        
        # Create DataFrame from input values
        input_df = pd.DataFrame([input_values])
        
        if st.button("Predict Friction Angle"):
            with st.spinner("Calculating prediction and SHAP values..."):
                # Make prediction
                prediction = predict_friction(input_df, model, scaler_X, scaler_y, device)
                
                # Calculate SHAP values
                shap_values, base_value = calculate_shap_values(input_df, model, X, scaler_X, scaler_y, device)
                
                # Display results
                st.header("Results")
                col1, col2 = st.columns(2)
                
                with col1:
                    st.metric("Predicted Friction Angle", f"{prediction:.2f}°")
                with col2:
                    st.metric("Base Value", f"{base_value:.2f}°")
                
                # Create and display waterfall plot
                st.header("SHAP Waterfall Plot")
                waterfall_plot = create_waterfall_plot(
                    shap_values=shap_values,
                    feature_names=feature_names,
                    base_value=base_value,
                    input_data=input_df.values[0]
                )
                st.image(waterfall_plot)
    
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        st.info("Please try refreshing the page. If the error persists, contact support.")

if __name__ == "__main__":
    main()