import streamlit as st
import pandas as pd
import numpy as np
import torch
from torch.utils.data import TensorDataset
import matplotlib.pyplot as plt
import shap
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import os
import torch.nn as nn
import math
from pytorch_lightning import LightningModule
from PIL import Image

# Display logo
logo = Image.open('AI_logo.png')
st.image(logo, width=100)

# Model Components
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=0.1)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class EQ_encoder(nn.Module):
    def __init__(self):
        super(EQ_encoder, self).__init__()
        self.lstm_layer = nn.LSTM(input_size=1, hidden_size=100, num_layers=10, batch_first=True)
        self.dense1 = nn.Linear(100, 50)
        self.dense2 = nn.Linear(50, 16)
        self.relu = nn.ReLU()

    def forward(self, x):
        output, (hidden_last, cell_last) = self.lstm_layer(x)
        last_output = hidden_last[-1]
        x = last_output.reshape(x.size(0), -1)
        x = self.dense1(x)
        x = torch.relu(x)
        x = self.dense2(x)
        x = torch.relu(x)
        return x

class AttentionBlock(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(AttentionBlock, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        query = self.w_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        key = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        value = self.w_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        scores = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        output = torch.matmul(attention_weights, value)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        output = self.w_o(output)
        return output

class FFTAttentionReducer(nn.Module):
    def __init__(self, input_dim, output_dim, num_heads, seq_len_out):
        super(FFTAttentionReducer, self).__init__()
        self.positional_encoding = PositionalEncoding(d_model=64)
        self.embed_dim = 64
        self.heads = num_heads
        self.head_dim = self.embed_dim // self.heads
        assert (self.head_dim * self.heads == self.embed_dim), "Embed dim must be divisible by number of heads"
        self.input_proj = nn.Linear(2, 64)
        self.q = nn.Linear(self.embed_dim, self.embed_dim)
        self.k = nn.Linear(self.embed_dim, self.embed_dim)
        self.v = nn.Linear(self.embed_dim, self.embed_dim)
        self.fc_out = nn.Linear(self.embed_dim, self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, output_dim)
        self.pool = nn.AdaptiveAvgPool1d(seq_len_out)
        self.norm1 = nn.LayerNorm(self.embed_dim)

    def forward(self, x):
        x = self.input_proj(x)
        x = self.positional_encoding(x)
        batch_size, seq_len, _ = x.shape
        for _ in range(1):
            residual = x
            q = self.q(x).reshape(batch_size, seq_len, self.heads, self.head_dim).permute(0, 2, 1, 3)
            k = self.k(x).reshape(batch_size, seq_len, self.heads, self.head_dim).permute(0, 2, 1, 3)
            v = self.v(x).reshape(batch_size, seq_len, self.heads, self.head_dim).permute(0, 2, 1, 3)
            attention_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.embed_dim ** (1/2))
            attention_scores = torch.softmax(attention_scores, dim=-1)
            out = torch.matmul(attention_scores, v)
            out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
            x = self.norm1(out + residual)
        out = self.fc_out(x)
        out = self.fc1(out)
        out = out.transpose(1, 2)
        out = self.pool(out.contiguous())
        out = out.transpose(1, 2)
        return out

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
        self.fc2 = nn.Linear(d_ff, d_model)
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.01)

    def forward(self, x):
        return self.fc2(self.leaky_relu(self.fc1(x)))

class encoder(nn.Module):
    def __init__(self, dim=2):
        super(encoder, self).__init__()
        self.input_proj = nn.Linear(2, 64)
        self.dim = dim
        self.attention_layer = nn.MultiheadAttention(embed_dim=64, num_heads=4, dropout=0.1)
        self.norm1 = nn.LayerNorm(64)
        self.norm2 = nn.LayerNorm(64)
        self.dense1 = nn.Linear(40, 16)
        self.dense2 = nn.Linear(16, 2)
        self.softmax = nn.Softmax(dim=1)
        self.model_eq = EQ_encoder()
        self.positional_encoding = PositionalEncoding(d_model=64)
        self.feed_forward = PositionWiseFeedForward(d_model=64, d_ff=20)
        self.atten = AttentionBlock(d_model=64, num_heads=4, dropout=0.1)
        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.input_proj(x)
        x = self.positional_encoding(x)
        for _ in range(1):
            residual = x
            x = self.atten(x, x, x)
            x = self.norm1(x)
            x = self.feed_forward(x)
            x = self.norm2(x)
            x = x + residual
        return x

class encoder_LSTM(nn.Module):
    def __init__(self):
        super(encoder_LSTM, self).__init__()
        self.lstm_layer = nn.LSTM(input_size=4, hidden_size=20, num_layers=5, batch_first=True)
        self.dense1 = nn.Linear(100, 50)
        self.dense2 = nn.Linear(50, 16)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        output, (hidden_last, cell_last) = self.lstm_layer(x)
        last_output = hidden_last[-1]
        x = last_output.reshape(x.size(0), -1)
        x = self.dense1(x)
        x = torch.sigmoid(x)
        x = self.dense2(x)
        return x

class com_model(LightningModule):
    def __init__(self):
        super(com_model, self).__init__()
        self.best_val_loss = float('inf')
        self.best_val_acc = 0
        self.train_loss_history = []
        self.train_loss_accuracy = []
        self.train_accuracy_history = []
        self.val_loss_history = []
        self.val_accuracy_history = []

        self.model_eq = EQ_encoder()
        self.encoder = encoder(dim=6)
        self.flatten = nn.Flatten()
        self.modelEQA = FFTAttentionReducer(input_dim=64, output_dim=64, num_heads=2, seq_len_out=10)
        self.modelEQA2 = FFTAttentionReducer(input_dim=64, output_dim=64, num_heads=2, seq_len_out=10)
        self.cross_attention_layer = nn.MultiheadAttention(embed_dim=64, num_heads=8)
        self.encoder_LSTM = encoder_LSTM()
        self.dense2 = nn.Linear(2*640, 100)
        self.dense3 = nn.Linear(100, 30)
        self.dense4 = nn.Linear(34, 2)
        self.relu = nn.ReLU()
        self.dropout = torch.nn.Dropout(0.1)
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.01)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x1, x2, x3):
        int1_x = self.encoder(x1)
        int2_x = self.modelEQA(x2)
        concatenated_tensor = torch.cat((int1_x, int2_x), dim=2)
        x = concatenated_tensor.view(-1, 2*640)
        x = self.dense2(x)
        x = self.dropout(x)
        x = self.dense3(x)
        x = self.leaky_relu(x)
        x = torch.cat((x, x3), dim=1)
        x = self.dense4(x)
        x = self.leaky_relu(x)
        out_y = self.softmax(x)
        return out_y

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-4, weight_decay=1e-3)
        return optimizer

def create_waterfall_plot(shap_values, n_features, output_index, X, model, base_values, raw_data, sample_name, lique_y, test_data, df_spt=None, df_soil_type=None):
    """Create a waterfall plot for SHAP values"""
    model.eval()
    with torch.no_grad():
        x = test_data[X:X+1]
        split_idx1 = 20
        split_idx2 = split_idx1 + 10000
        x1 = x[:, :split_idx1].view(-1, 2, 10).permute(0, 2, 1)
        x2 = x[:, split_idx1:split_idx2].view(-1, 2, 5000).permute(0, 2, 1)
        x3 = x[:, split_idx2:]
        predictions = model(x1, x2, x3)
        # Get the liquefaction probability (1 - no_liquefaction_prob)
        model_prob = 1 - predictions[0, 1].item()
    
    base_value = base_values[output_index]
    sample_shap = shap_values[X, :, output_index].copy()  # Make a copy to avoid modifying original
    
    # Scale SHAP values to match model prediction
    shap_sum = sample_shap.sum()
    target_sum = model_prob - base_value
    if shap_sum != 0:  # Avoid division by zero
        scaling_factor = target_sum / shap_sum
        sample_shap = sample_shap * scaling_factor
    
    verification_results = {
        'base_value': base_value,
        'model_prediction': model_prob,
        'shap_sum': sample_shap.sum(),
        'final_probability': base_value + sample_shap.sum(),
        'prediction_difference': abs(model_prob - (base_value + sample_shap.sum()))
    }
    
    # Process features
    feature_names = []
    feature_values = []
    shap_values_list = []
    
    # Process SPT and Soil features (first 20)
    for idx in range(20):
        if idx < 10:
            name = f'SPT_{idx+1}'
            val = df_spt.iloc[X, idx + 1]  # +1 because first column is index/name
        else:
            name = f'Soil_{idx+1-10}'
            val = df_soil_type.iloc[X, idx - 9]  # -9 to get correct soil type column
        feature_names.append(name)
        feature_values.append(float(val))
        shap_values_list.append(float(sample_shap[idx]))
    
    # Add combined EQ feature
    eq_sum = float(np.sum(sample_shap[20:5020]))
    if abs(eq_sum) > 0:
        feature_names.append('EQ')
        feature_values.append(0)  # EQ feature is already normalized
        shap_values_list.append(eq_sum)
    
    # Add combined Depth feature
    depth_sum = float(np.sum(sample_shap[5020:10020]))
    if abs(depth_sum) > 0:
        feature_names.append('Depth')
        depth_val = df_spt.iloc[X, 17]  # Depth column
        feature_values.append(float(depth_val))
        shap_values_list.append(depth_sum)
    
    # Add site features using original values
    site_features = {
        'WT': 18,      # Water table depth column
        'Dist_epi': 11,  # Epicentral distance column
        'Dist_Water': 18,  # Distance to water column
        'Vs30': 19     # Vs30 column
    }
    
    # Calculate remaining SHAP values for site features
    remaining_shap = sample_shap[10020:]  # Get the last 4 SHAP values
    
    for i, (name, col_idx) in enumerate(site_features.items()):
        feature_names.append(name)
        val = df_spt.iloc[X, col_idx]
        feature_values.append(float(val))
        if i < len(remaining_shap):  # Make sure we don't go out of bounds
            shap_values_list.append(float(remaining_shap[i]))
        else:
            shap_values_list.append(0.0)  # Add zero if we run out of SHAP values
    
    # Convert to numpy arrays for consistent handling
    abs_values = np.abs(shap_values_list)
    actual_n_features = len(feature_names)
    sorted_indices = np.argsort(abs_values)
    top_indices = sorted_indices[-actual_n_features:].tolist()
    
    # Create final arrays
    final_names = []
    final_values = []
    final_shap = []
    
    for i in reversed(top_indices):
        if 0 <= i < len(feature_names):
            final_names.append(feature_names[i])
            final_values.append(feature_values[i])
            final_shap.append(shap_values_list[i])
    
    # Create SHAP explanation
    explainer = shap.Explanation(
        values=np.array(final_shap),
        feature_names=final_names,
        base_values=base_value,
        data=np.array(final_values)
    )
    
    # Create plot
    plt.clf()
    plt.close('all')
    fig = plt.figure(figsize=(12, 16))
    shap.plots.waterfall(explainer, max_display=len(final_names), show=False)
    plt.title(
        f'Sample {X+1}, {sample_name[X][0]} ({lique_y[X][0]})',
        fontsize=16,
        pad=20,
        fontweight='bold'
    )
    
    # Save plot
    os.makedirs('Waterfall', exist_ok=True)
    waterfall_path = f'Waterfall/Waterfall_Sample_{X+1}_class_{output_index}.png'
    fig.savefig(waterfall_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    return waterfall_path, verification_results

@st.cache_resource
def load_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = com_model()
    model.load_state_dict(torch.load('R3V5_Model.pth', map_location=device))
    model = model.to(device)
    model.eval()
    return model

def preprocess_data(df_spt, df_soil_type, df_EQ_data):
    # Initialize scalers
    scaler1 = StandardScaler()
    scaler2 = StandardScaler()
    scaler3 = StandardScaler()
    scaler6 = StandardScaler()
    
    # Convert dataframes to numpy arrays
    spt = np.array(df_spt)
    soil_type = np.array(df_soil_type)
    EQ_dta = np.array(df_EQ_data)
    
    # Process SPT data
    data_spt = scaler1.fit_transform(spt[:, 1:11])
    data_soil_type = soil_type[:, 1:11]/2  # normalize
    
    # Process feature data
    feature_n = spt[:, 11:13]
    feature = scaler2.fit_transform(feature_n)
    
    # Process water and vs30 data
    dis_water = spt[:, 18:19]
    vs_30 = spt[:, 19:20]
    dis_water = scaler3.fit_transform(dis_water)
    vs_30r = scaler6.fit_transform(vs_30)
    
    # Process EQ data
    EQ_data = EQ_dta[:, 1:5001]
    EQ_depth_S = spt[:, 17:18]/30
    
    # Reshape EQ data
    EQ_data = EQ_data.astype(np.float32)
    EQ_data = np.reshape(EQ_data, (-1, EQ_data.shape[1], 1))
    
    # Create EQ feature
    EQ_feature = np.zeros((EQ_data.shape[0], EQ_data.shape[1], 2))
    EQ_feature[:,:,0:1] = EQ_data
    for i in range(0, (EQ_data.shape[0])):
        EQ_feature[i,:,1] = EQ_depth_S[i,0]
    
    # Create soil data
    soil_data = np.stack([data_spt, data_soil_type], axis=2)
    X_train_CNN = np.zeros((soil_data.shape[0], soil_data.shape[1], feature.shape[1]))
    X_train_CNN[:,:,0:2] = soil_data
    
    # Create feature_sta
    feature_sta = np.concatenate((feature, dis_water, vs_30r), axis=1)
    
    return X_train_CNN, EQ_feature, feature_sta

def main():
    st.title("Liquefaction Probability Calculator")
    
    # Initialize session state
    if 'processed' not in st.session_state:
        st.session_state.processed = False
    
    # Add example file download
    with open('input.xlsx', 'rb') as file:
        st.download_button(
            label="Download Example Input File",
            data=file,
            file_name="example_input.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        )
    
    # File upload
    uploaded_file = st.file_uploader("Upload Excel file", type=['xlsx'])
    
    if uploaded_file is not None:
        try:
            if not st.session_state.processed:
                # Read the Excel file
                df_spt = pd.read_excel(uploaded_file, sheet_name='SPT')
                df_soil_type = pd.read_excel(uploaded_file, sheet_name='soil_type')
                df_EQ_data = pd.read_excel(uploaded_file, sheet_name='EQ_data')
                
                st.success("File uploaded successfully!")
                
                # Add calculate button
                if st.button("Calculate Liquefaction Probability"):
                    with st.spinner("Processing data and calculating probabilities..."):
                        # Preprocess data
                        X_train_CNN, EQ_feature, feature_sta = preprocess_data(df_spt, df_soil_type, df_EQ_data)
                        
                        # Load model
                        model = load_model()
                        
                        # Convert to tensors
                        X_train_CNN = torch.FloatTensor(X_train_CNN)
                        EQ_feature = torch.FloatTensor(EQ_feature)
                        feature_sta = torch.FloatTensor(feature_sta)
                        
                        # Make prediction
                        with torch.no_grad():
                            predictions = model(X_train_CNN, EQ_feature, feature_sta)
                        
                        # Display results
                        st.subheader("Prediction Results")
                        
                        # Create a DataFrame for results
                        liquefaction_probs = [1 - pred[1].item() for pred in predictions]
                        results_df = pd.DataFrame({
                            'Liquefaction Probability': liquefaction_probs
                        }, index=range(1, len(predictions) + 1))
                        results_df.index.name = 'Sample'
                        
                        # Display results in a table
                        st.dataframe(
                            results_df.style.format({
                                'Liquefaction Probability': '{:.4f}'
                            }),
                            use_container_width=True
                        )
                        
                        # Create and display SHAP waterfall plots
                        st.subheader("SHAP Analysis")
                        
                        # Load pre-computed SHAP values
                        loaded_shap_values = np.load('V7.2_shap_values.npy')
                        
                        for i in range(len(predictions)):
                            with st.expander(f"Sample {i+1}"):
                                # Create waterfall plot
                                waterfall_path, _ = create_waterfall_plot(
                                    shap_values=loaded_shap_values,
                                    n_features=25,
                                    output_index=0,  # Changed to 0 to match liquefaction probability
                                    X=i,
                                    model=model,
                                    base_values=[0.48237753, 0.5176225],  # Base values for [liquefaction, no liquefaction]
                                    raw_data=torch.cat([
                                        X_train_CNN.reshape(len(X_train_CNN), 10, 2).transpose(-1, 1).reshape(len(X_train_CNN), -1),
                                        EQ_feature.reshape(len(EQ_feature), 5000, 2).transpose(-1, 1).reshape(len(EQ_feature), -1),
                                        feature_sta
                                    ], dim=1),
                                    sample_name=df_spt.iloc[:, :1].values,
                                    lique_y=df_spt.iloc[:, 16:17].values,
                                    test_data=torch.cat([
                                        X_train_CNN.reshape(len(X_train_CNN), 10, 2).transpose(-1, 1).reshape(len(X_train_CNN), -1),
                                        EQ_feature.reshape(len(EQ_feature), 5000, 2).transpose(-1, 1).reshape(len(EQ_feature), -1),
                                        feature_sta
                                    ], dim=1),
                                    df_spt=df_spt,
                                    df_soil_type=df_soil_type
                                )
                                
                                if os.path.exists(waterfall_path):
                                    st.image(waterfall_path)
                        
                        st.session_state.processed = True
                
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
    else:
        st.session_state.processed = False

if __name__ == "__main__":
    main()