import streamlit as st
import pandas as pd
import numpy as np
import torch
import matplotlib.pyplot as plt
import pickle
import os
import warnings
from lllm_model_all_token import LLMConcreteModel

# Optimize for deployment
warnings.filterwarnings('ignore')
torch.set_num_threads(2)

# Canva-style colors
CANVA_PURPLE = "#8B5CF6"
CANVA_LIGHT_PURPLE = "#A78BFA"
CANVA_DARK_PURPLE = "#7C3AED"
CANVA_BACKGROUND = "#FAFAFA"
CANVA_WHITE = "#FFFFFF"

# Set page config with Canva-style theme
st.set_page_config(
    page_title="Concrete Creep Prediction",
    page_icon="🏗️",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# Custom CSS for Canva-style design
st.markdown(f"""
<style>
    .main {{
        background-color: {CANVA_BACKGROUND};
    }}
    
    .stApp {{
        background-color: {CANVA_BACKGROUND};
    }}
    
    .css-1d391kg {{
        background-color: {CANVA_WHITE};
        padding: 2rem;
        border-radius: 15px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin: 1rem 0;
    }}
    
    .stButton > button {{
        background-color: {CANVA_PURPLE};
        color: white;
        border: none;
        border-radius: 25px;
        padding: 0.75rem 2rem;
        font-weight: 600;
        font-size: 16px;
        transition: all 0.3s ease;
        width: 100%;
    }}
    
    .stButton > button:hover {{
        background-color: {CANVA_DARK_PURPLE};
        transform: translateY(-2px);
        box-shadow: 0 4px 12px rgba(139, 92, 246, 0.3);
    }}
    
    .stNumberInput > div > div > input {{
        border-radius: 10px;
        border: 2px solid #E5E7EB;
        padding: 0.75rem;
    }}
    
    .stNumberInput > div > div > input:focus {{
        border-color: {CANVA_PURPLE};
        box-shadow: 0 0 0 3px rgba(139, 92, 246, 0.1);
    }}
    
    .metric-card {{
        background: linear-gradient(135deg, {CANVA_PURPLE}, {CANVA_LIGHT_PURPLE});
        color: white;
        padding: 1.5rem;
        border-radius: 15px;
        text-align: center;
        margin: 0.5rem 0;
    }}
    
    .result-card {{
        background-color: {CANVA_WHITE};
        padding: 2rem;
        border-radius: 15px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin: 1rem 0;
    }}
    
    h1 {{
        color: {CANVA_DARK_PURPLE};
        text-align: center;
        font-weight: 700;
        margin-bottom: 2rem;
    }}
    
    h2, h3 {{
        color: {CANVA_DARK_PURPLE};
        font-weight: 600;
    }}
    
    .stSuccess {{
        background-color: #10B981;
        color: white;
        border-radius: 10px;
    }}
</style>
""", unsafe_allow_html=True)

# Simple CreepScaler class
class CreepScaler:
    def __init__(self, factor=1000):
        self.factor = factor
        self.mean_ = 0
        self.scale_ = factor
        self.is_standard_scaler = False
    
    def transform(self, X):
        if self.is_standard_scaler:
            return (X - self.mean_) / self.scale_
        return X / self.factor
    
    def inverse_transform(self, X):
        if self.is_standard_scaler:
            return (X * self.scale_) + self.mean_
        return X * self.factor

@st.cache_resource
def load_model():
    """Load model and scalers"""
    # Find model file
    model_files = ['best_llm_model-17.pt', 'final_llm_model-5.pt']
    model_path = None
    for file in model_files:
        if os.path.exists(file):
            model_path = file
            break
    
    if model_path is None:
        st.error("❌ Model file not found")
        st.stop()
    
    # Load scalers
    try:
        with open('scalers/feature_scaler.pkl', 'rb') as f:
            feature_scaler = pickle.load(f)
        
        try:
            with open('scalers/creep_scaler.pkl', 'rb') as f:
                creep_scaler = pickle.load(f)
        except:
            creep_scaler = CreepScaler(factor=1000)
        
        try:
            with open('scalers/time_values.pkl', 'rb') as f:
                time_values = pickle.load(f)
        except:
            time_values = np.arange(1, 1001)  # Default 1000 time points
            
    except Exception as e:
        st.error(f"❌ Error loading files: {e}")
        st.stop()
    
    # Load model
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = LLMConcreteModel(
        feature_dim=3,
        d_model=192,
        num_layers=4,
        num_heads=4,
        d_ff=768,
        dropout=0.057,
        target_len=1,
        pooling_method='hybrid'
    )
    
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model = model.to(device)
        model.eval()
    except Exception as e:
        st.error(f"❌ Error loading model: {e}")
        st.stop()
    
    return model, feature_scaler, creep_scaler, time_values, device

def predict_creep(model, features, time_values, feature_scaler, creep_scaler, device, max_days=365):
    """Simple prediction function"""
    # Scale features
    scaled_features = feature_scaler.transform(features)
    scaled_features_tensor = torch.FloatTensor(scaled_features).to(device)
    
    # Limit time values
    pred_time_values = time_values[:max_days] if max_days < len(time_values) else time_values
    
    predictions = [0.0]  # Start with 0
    scaled_predictions = [0.0]
    
    with torch.no_grad():
        for i in range(1, len(pred_time_values)):
            history = np.array(scaled_predictions)
            history_tensor = torch.FloatTensor(history).unsqueeze(0).to(device)
            
            time_history = np.log1p(pred_time_values[:i])
            time_tensor = torch.FloatTensor(time_history).unsqueeze(0).to(device)
            
            length = torch.tensor([len(history)], device=device)
            
            next_value = model(
                creep_history=history_tensor,
                features=scaled_features_tensor,
                lengths=length,
                time_history=time_tensor
            ).item()
            
            scaled_predictions.append(next_value)
            next_creep = creep_scaler.inverse_transform(np.array([[next_value]])).flatten()[0]
            predictions.append(next_creep)
    
    return np.array(predictions), pred_time_values

# Load model
model, feature_scaler, creep_scaler, time_values, device = load_model()

def get_base64_of_image(path):
    """Convert image to base64 string"""
    import base64
    try:
        with open(path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode()
    except:
        return ""

# App title with logo
st.markdown("""
<div style='text-align: center; padding: 2rem 0;'>
    <div style='display: flex; justify-content: center; align-items: center; margin-bottom: 1.5rem; flex-wrap: wrap;'>
        <img src='data:image/png;base64,{}' style='width: 120px; height: auto; max-height: 100px; margin-right: 1.5rem; margin-bottom: 1rem; border-radius: 10px; box-shadow: 0 4px 12px rgba(139, 92, 246, 0.2); object-fit: contain;'>
        <div style='text-align: center;'>
            <h1 style='margin: 0; color: {}; font-size: 2.5rem; font-weight: 700;'>🏗️ Concrete Creep Prediction</h1>
            <p style='margin: 0; font-size: 18px; color: #6B7280; font-weight: 500;'>AI-Powered Concrete Analysis</p>
        </div>
    </div>
</div>
""".format(
    get_base64_of_image("AI_logo.png"),
    CANVA_DARK_PURPLE
), unsafe_allow_html=True)

# Input form in a clean card
with st.container():
    st.markdown('<div class="css-1d391kg">', unsafe_allow_html=True)
    
    st.markdown("### 📝 Enter Concrete Properties")
    
    col1, col2 = st.columns(2)
    
    with col1:
        density = st.number_input(
            "Density (kg/m³)", 
            min_value=2000.0, 
            max_value=3000.0, 
            value=2490.0,
            step=10.0
        )
        
        fc = st.number_input(
            "Compressive Strength (ksc)", 
            min_value=10.0, 
            max_value=1000.0, 
            value=670.0,
            step=10.0
        )
    
    with col2:
        e_modulus = st.number_input(
            "Elastic Modulus (ksc)", 
            min_value=10000.0, 
            max_value=1000000.0, 
            value=436000.0,
            step=1000.0
        )
    
    st.markdown('</div>', unsafe_allow_html=True)

# Predict button
if st.button("🚀 Predict Creep Strain"):
    # Set default prediction days
    max_days = 365
    
    # Create features
    features_dict = {
        'Density': density,
        'fc': fc,
        'E': e_modulus
    }
    df_features = pd.DataFrame([features_dict])
    
    # Run prediction
    with st.spinner("🔄 Predicting..."):
        try:
            predictions, pred_time_values = predict_creep(
                model, df_features, time_values, 
                feature_scaler, creep_scaler, device, max_days
            )
            
            # Results
            st.markdown('<div class="result-card">', unsafe_allow_html=True)
            
            # Key metrics
            col1, col2 = st.columns(2)
            with col1:
                st.markdown(f"""
                <div class="metric-card">
                    <h3>{predictions[-1]:.1f}</h3>
                    <p>Final Creep (µε)</p>
                </div>
                """, unsafe_allow_html=True)
            
            with col2:
                st.markdown(f"""
                <div class="metric-card">
                    <h3>{np.max(predictions):.1f}</h3>
                    <p>Maximum Creep (µε)</p>
                </div>
                """, unsafe_allow_html=True)
            
            # Simple plot
            st.markdown("### 📊 Creep Strain Over Time")
            
            # Set plot style to match Canva theme
            plt.style.use('default')
            fig, ax = plt.subplots(figsize=(10, 6))
            fig.patch.set_facecolor('white')
            
            ax.plot(pred_time_values, predictions, 
                   color=CANVA_PURPLE, linewidth=3, alpha=0.8)
            ax.fill_between(pred_time_values, predictions, 
                           alpha=0.2, color=CANVA_LIGHT_PURPLE)
            
            ax.set_xlabel('Time (days)', fontsize=12, color='#374151')
            ax.set_ylabel('Creep Strain (µε)', fontsize=12, color='#374151')
            ax.grid(True, alpha=0.3, color='#E5E7EB')
            ax.set_facecolor('#FAFAFA')
            
            # Remove top and right spines
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['left'].set_color('#E5E7EB')
            ax.spines['bottom'].set_color('#E5E7EB')
            
            plt.tight_layout()
            st.pyplot(fig)
            
            # Download data
            results_df = pd.DataFrame({
                'Time (days)': pred_time_values,
                'Creep Strain (µε)': predictions
            })
            
            csv = results_df.to_csv(index=False)
            st.download_button(
                label="💾 Download Results",
                data=csv,
                file_name="creep_predictions.csv",
                mime="text/csv"
            )
            
            st.markdown('</div>', unsafe_allow_html=True)
            
        except Exception as e:
            st.error(f"❌ Prediction failed: {e}")

# Simple footer
st.markdown("""
<div style='text-align: center; padding: 2rem 0; color: #9CA3AF;'>
    <p>🏗️ Concrete Creep Prediction Tool</p>
    <p style='margin-top: 0.5rem; font-size: 14px;'>Developed by <strong>CIFIR</strong> and <strong>AI Research Group KMUTT</strong></p>
</div>
""", unsafe_allow_html=True)