# app.py
import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import requests
import json
import os
from groq import Groq
from dotenv import load_dotenv
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime

# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv('GROQ_API_KEY')

# Page configuration
st.set_page_config(
    page_title="Hydrogen Production Optimizer",
    page_icon="⚡",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main {
        padding: 1rem;
    }
    .stButton>button {
        width: 100%;
        background-color: #4CAF50;
        color: white;
    }
    .stTabs [data-baseweb="tab-list"] {
        gap: 10px;
    }
    .stTabs [data-baseweb="tab"] {
        padding: 10px;
        border-radius: 4px 4px 0px 0px;
    }
    h1, h2, h3 {
        color: #1E88E5;
    }
    .highlight {
        background-color: #f0f8ff;
        padding: 1rem;
        border-radius: 0.5rem;
        border-left: 5px solid #1E88E5;
    }
</style>
""", unsafe_allow_html=True)

# Helper functions for calculations
def calculate_h2_production(method, water_quantity, energy_input, current_density, voltage):
    """Calculate hydrogen production based on input parameters"""
    # Conversion factors and efficiencies for different methods
    method_efficiencies = {
        "Alkaline Electrolysis": 0.65,
        "PEM Electrolysis": 0.75,
        "SOEC": 0.85
    }
    
    # Basic Faraday's law calculation (simplified)
    # Assuming ideal conditions and standard molar volume
    faraday_constant = 96485  # C/mol
    molar_mass_h2 = 2.02  # g/mol
    
    # Adjusting efficiency based on method
    efficiency = method_efficiencies[method]
    
    # Calculate total charge (Q = I * t)
    # Assuming current_density is in A/cm² and we convert water_quantity to surface area
    surface_area = water_quantity * 0.1  # Simplified conversion
    current = current_density * surface_area  # Total current in A
    
    # Assuming energy_input helps determine operation time
    time_hours = energy_input / (voltage * current)  # Time in hours
    
    # Calculate hydrogen production using Faraday's Law
    moles_h2 = (current * time_hours * 3600 * efficiency) / (2 * faraday_constant)
    mass_h2 = moles_h2 * molar_mass_h2  # grams
    volume_h2 = moles_h2 * 22.4  # Standard liters
    
    return {
        "production_rate_g_per_hour": mass_h2 / time_hours,
        "total_production_g": mass_h2,
        "total_production_L": volume_h2,
        "efficiency": efficiency,
        "operation_time_hours": time_hours
    }

def calculate_cost(method, water_cost, water_purification_cost, energy_source, energy_input, h2_production):
    """Calculate the cost of hydrogen production"""
    # Energy costs by source ($/kWh)
    energy_costs = {
        "Grid Electricity": 0.12,
        "Solar": 0.08,
        "Wind": 0.06,
        "Nuclear": 0.10,
        "Hydroelectric": 0.07
    }
    
    # Operational costs by method ($/kg H2)
    operational_costs = {
        "Alkaline Electrolysis": 1.2,
        "PEM Electrolysis": 1.5,
        "SOEC": 1.8
    }
    
    # Calculate water cost
    total_water_cost = water_cost * (h2_production["total_production_g"] / 1000)  # $/kg
    
    # Calculate purification cost
    total_purification_cost = water_purification_cost * (h2_production["total_production_g"] / 1000)  # $/kg
    
    # Calculate energy cost
    energy_cost_rate = energy_costs[energy_source]
    total_energy_cost = energy_cost_rate * energy_input  # $
    
    # Calculate operational cost
    operational_cost_rate = operational_costs[method]
    total_operational_cost = operational_cost_rate * (h2_production["total_production_g"] / 1000)  # $
    
    # Calculate total cost
    total_cost = total_water_cost + total_purification_cost + total_energy_cost + total_operational_cost
    
    # Calculate cost per kg of H2
    cost_per_kg = total_cost / (h2_production["total_production_g"] / 1000) if h2_production["total_production_g"] > 0 else 0
    
    return {
        "water_cost": total_water_cost,
        "purification_cost": total_purification_cost,
        "energy_cost": total_energy_cost,
        "operational_cost": total_operational_cost,
        "total_cost": total_cost,
        "cost_per_kg": cost_per_kg
    }

def call_groq_api(user_inputs, production_data, cost_data):
    """Call Groq API with Llama 3 to analyze production parameters and provide recommendations"""
        
    # Initialize Groq client
    try:
        client = Groq(api_key=os.environ.get("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"))
    except Exception as e:
        return {"error": f"Failed to initialize Groq client: {str(e)}"}
    
    # Prepare the prompt without f-string for the JSON schema part
    prompt_part1 = f"""
    As a hydrogen production expert, analyze the following electrolysis parameters and provide recommendations for optimization:
    
    Input Parameters:
    - Water Source: {user_inputs['water_source']}
    - Production Method: {user_inputs['production_method']}
    - Energy Source: {user_inputs['energy_source']}
    - Current Density: {user_inputs['current_density']} A/cm²
    - Voltage: {user_inputs['voltage']} V
    - Membrane Material: {user_inputs['membrane']}
    - Electrode Materials: {user_inputs['electrodes']}
    
    Production Results:
    - Production Rate: {production_data['production_rate_g_per_hour']:.2f} g/hour
    - Total Production: {production_data['total_production_g']:.2f} g
    - Efficiency: {production_data['efficiency'] * 100:.1f}%
    - Operation Time: {production_data['operation_time_hours']:.2f} hours
    
    Cost Analysis:
    - Water Cost: ${cost_data['water_cost']:.2f}
    - Purification Cost: ${cost_data['purification_cost']:.2f}
    - Energy Cost: ${cost_data['energy_cost']:.2f}
    - Operational Cost: ${cost_data['operational_cost']:.2f}
    - Total Cost: ${cost_data['total_cost']:.2f}
    - Cost per kg H₂: ${cost_data['cost_per_kg']:.2f}
    
    Please provide:
    1. An efficiency assessment of the current setup
    2. Three specific recommendations to improve efficiency
    3. Three specific recommendations to reduce costs
    4. An ideal parameter configuration based on the provided inputs
    
    Format your response as a structured JSON with these fields:
    """
    
    # The JSON schema part without using f-string
    prompt_part2 = """
    {
        "efficiency_assessment": "text analysis",
        "efficiency_recommendations": ["recommendation1", "recommendation2", "recommendation3"],
        "cost_recommendations": ["recommendation1", "recommendation2", "recommendation3"],
        "ideal_parameters": {
            "current_density": value,
            "voltage": value,
            "membrane": "recommendation",
            "electrodes": "recommendation",
            "energy_source": "recommendation"
        },
        "estimated_improvement": {
            "efficiency_increase": "percentage",
            "cost_reduction": "percentage"
        }
    }
    """
    
    # Combine the parts
    prompt = prompt_part1 + prompt_part2
    
    try:
        # Call Groq API using the client
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.5,
            max_tokens=1024,
            response_format={"type": "json_object"}
        )
        
        # Extract the response content
        response_content = chat_completion.choices[0].message.content
        
        # Parse the JSON response
        try:
            recommendations_json = json.loads(response_content)
            return recommendations_json
        except json.JSONDecodeError:
            return {"error": "Failed to parse API response as JSON"}
            
    except Exception as e:
        return {"error": f"Error calling Groq API: {str(e)}"}

# Main application
def main():
    # Sidebar for inputs
    st.sidebar.image("https://upload.wikimedia.org/wikipedia/commons/thumb/1/18/Creative-Tail-Objects-flask.svg/256px-Creative-Tail-Objects-flask.svg.png", width=100)
    st.sidebar.title("H₂ Production Parameters")
    
    # Water parameters
    st.sidebar.subheader("Water Parameters")
    water_source = st.sidebar.selectbox("Water Source", ["Tap Water", "Deionized Water", "Seawater", "Wastewater", "Ultrapure Water"])
    water_cost = st.sidebar.number_input("Water Cost ($/m³)", min_value=0.1, max_value=50.0, value=2.0, step=0.1)
    water_purification_cost = st.sidebar.number_input("Water Purification Cost ($/m³)", min_value=0.0, max_value=100.0, value=5.0, step=0.5)
    water_quantity = st.sidebar.number_input("Water Quantity (L)", min_value=1.0, max_value=10000.0, value=100.0, step=10.0)
    
    # Electrolysis parameters
    st.sidebar.subheader("Electrolysis Parameters")
    production_method = st.sidebar.selectbox("Production Method", ["Alkaline Electrolysis", "PEM Electrolysis", "SOEC"])
    current_density = st.sidebar.slider("Current Density (A/cm²)", min_value=0.1, max_value=2.0, value=0.5, step=0.1)
    voltage = st.sidebar.slider("Voltage (V)", min_value=1.4, max_value=5.0, value=2.0, step=0.1)
    
    # Materials
    membrane = st.sidebar.selectbox("Membrane Material", ["Nafion", "Zirfon", "Ceramic", "PBI", "SPEEK"])
    electrodes = st.sidebar.selectbox("Electrode Materials", ["Platinum", "Nickel", "Iridium Oxide", "Stainless Steel", "Carbon-based"])
    
    # Energy parameters
    st.sidebar.subheader("Energy Parameters")
    energy_source = st.sidebar.selectbox("Energy Source", ["Grid Electricity", "Solar", "Wind", "Nuclear", "Hydroelectric"])
    energy_input = st.sidebar.number_input("Energy Input (kWh)", min_value=1.0, max_value=10000.0, value=100.0, step=10.0)
    
    # Collect user inputs
    user_inputs = {
        "water_source": water_source,
        "water_cost": water_cost,
        "water_purification_cost": water_purification_cost,
        "water_quantity": water_quantity,
        "production_method": production_method,
        "current_density": current_density,
        "voltage": voltage,
        "membrane": membrane,
        "electrodes": electrodes,
        "energy_source": energy_source,
        "energy_input": energy_input
    }
    
    # Main content area
    st.title("Hydrogen Production Analysis & Optimization")
    st.markdown("Analyze and optimize your hydrogen production process with AI-driven recommendations")
    
    # Process button
    analyze_button = st.button("Analyze Production Parameters")
    
    if analyze_button:
        with st.spinner("Calculating production parameters and generating AI recommendations..."):
            # Calculate production
            production_data = calculate_h2_production(
                production_method, 
                water_quantity, 
                energy_input, 
                current_density, 
                voltage
            )
            
            # Calculate cost
            cost_data = calculate_cost(
                production_method,
                water_cost,
                water_purification_cost,
                energy_source,
                energy_input,
                production_data
            )
            
            # Get AI recommendations
            ai_recommendations = call_groq_api(user_inputs, production_data, cost_data)
            
            # Display results in tabs
            tabs = st.tabs(["Production Analysis", "Cost Analysis", "AI Recommendations", "Visualization"])
            
            # Tab 1: Production Analysis
            with tabs[0]:
                st.header("Hydrogen Production Analysis")
                
                # Key metrics in columns
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("Production Rate", f"{production_data['production_rate_g_per_hour']:.2f} g/hour")
                    st.metric("Total H₂ Produced", f"{production_data['total_production_g']:.2f} g")
                
                with col2:
                    st.metric("Volume of H₂", f"{production_data['total_production_L']:.2f} L")
                    st.metric("Operation Time", f"{production_data['operation_time_hours']:.2f} hours")
                
                with col3:
                    st.metric("System Efficiency", f"{production_data['efficiency']*100:.1f}%")
                    energy_consumption = energy_input / (production_data['total_production_g']/1000)
                    st.metric("Energy Consumption", f"{energy_consumption:.2f} kWh/kg H₂")
                
                # Detailed production information
                st.subheader("Process Details")
                process_df = pd.DataFrame({
                    "Parameter": ["Production Method", "Current Density", "Voltage", "Membrane", "Electrodes", 
                                 "Water Source", "Water Quantity", "Energy Source", "Energy Input"],
                    "Value": [production_method, f"{current_density} A/cm²", f"{voltage} V", membrane, electrodes,
                             water_source, f"{water_quantity} L", energy_source, f"{energy_input} kWh"]
                })
                st.table(process_df)
            
            # Tab 2: Cost Analysis
            with tabs[1]:
                st.header("Cost Analysis")
                
                # Key metrics
                col1, col2 = st.columns(2)
                with col1:
                    st.metric("Total Production Cost", f"${cost_data['total_cost']:.2f}")
                    st.metric("Cost per kg H₂", f"${cost_data['cost_per_kg']:.2f}")
                
                # Cost breakdown
                st.subheader("Cost Breakdown")
                cost_data_viz = {
                    "Category": ["Water", "Purification", "Energy", "Operation"],
                    "Cost ($)": [
                        cost_data["water_cost"],
                        cost_data["purification_cost"],
                        cost_data["energy_cost"],
                        cost_data["operational_cost"]
                    ]
                }
                cost_df = pd.DataFrame(cost_data_viz)
                
                # Create cost breakdown chart
                fig = px.pie(
                    cost_df, 
                    values="Cost ($)", 
                    names="Category",
                    title="Cost Distribution",
                    color_discrete_sequence=px.colors.sequential.Blues_r
                )
                st.plotly_chart(fig, use_container_width=True)
                
                # Detailed cost table
                st.subheader("Detailed Cost Breakdown")
                detailed_cost_df = pd.DataFrame({
                    "Cost Component": ["Water Cost", "Water Purification", "Energy Cost", "Operational Cost", "Total Cost"],
                    "Amount ($)": [
                        f"${cost_data['water_cost']:.2f}",
                        f"${cost_data['purification_cost']:.2f}",
                        f"${cost_data['energy_cost']:.2f}",
                        f"${cost_data['operational_cost']:.2f}",
                        f"${cost_data['total_cost']:.2f}"
                    ],
                    "Percentage": [
                        f"{cost_data['water_cost']/cost_data['total_cost']*100:.1f}%",
                        f"{cost_data['purification_cost']/cost_data['total_cost']*100:.1f}%",
                        f"{cost_data['energy_cost']/cost_data['total_cost']*100:.1f}%",
                        f"{cost_data['operational_cost']/cost_data['total_cost']*100:.1f}%",
                        "100%"
                    ]
                })
                st.table(detailed_cost_df)
            
            # Tab 3: AI Recommendations
            with tabs[2]:
                st.header("AI-Driven Recommendations")
                
                if "error" in ai_recommendations:
                    st.error(f"Error getting AI recommendations: {ai_recommendations['error']}")
                else:
                    # Efficiency Assessment
                    st.subheader("Efficiency Assessment")
                    st.markdown(f"<div class='highlight'>{ai_recommendations['efficiency_assessment']}</div>", unsafe_allow_html=True)
                    
                    # Recommendations
                    col1, col2 = st.columns(2)
                    with col1:
                        st.subheader("Efficiency Recommendations")
                        for i, rec in enumerate(ai_recommendations['efficiency_recommendations'], 1):
                            st.markdown(f"**{i}.** {rec}")
                    
                    with col2:
                        st.subheader("Cost Reduction Recommendations")
                        for i, rec in enumerate(ai_recommendations['cost_recommendations'], 1):
                            st.markdown(f"**{i}.** {rec}")
                    
                    # Ideal parameters
                    st.subheader("Ideal Parameter Configuration")
                    ideal_params = ai_recommendations['ideal_parameters']
                    
                    param_comparison = pd.DataFrame({
                        "Parameter": ["Current Density (A/cm²)", "Voltage (V)", "Membrane", "Electrodes", "Energy Source"],
                        "Current Value": [
                            current_density,
                            voltage,
                            membrane,
                            electrodes,
                            energy_source
                        ],
                        "Recommended Value": [
                            ideal_params['current_density'],
                            ideal_params['voltage'],
                            ideal_params['membrane'],
                            ideal_params['electrodes'],
                            ideal_params['energy_source']
                        ]
                    })
                    st.table(param_comparison)
                    
                    # Estimated improvements
                    st.subheader("Estimated Improvements")
                    col1, col2 = st.columns(2)
                    with col1:
                        st.metric("Efficiency Increase", ai_recommendations['estimated_improvement']['efficiency_increase'])
                    with col2:
                        st.metric("Cost Reduction", ai_recommendations['estimated_improvement']['cost_reduction'])
            
            # Tab 4: Visualization
            with tabs[3]:
                st.header("Production Visualization")
                
                # Create comparison data
                methods = ["Alkaline Electrolysis", "PEM Electrolysis", "SOEC"]
                production_rates = []
                costs_per_kg = []
                
                for method in methods:
                    # Calculate for each method
                    prod_data = calculate_h2_production(
                        method, 
                        water_quantity, 
                        energy_input, 
                        current_density, 
                        voltage
                    )
                    
                    cost_result = calculate_cost(
                        method,
                        water_cost,
                        water_purification_cost,
                        energy_source,
                        energy_input,
                        prod_data
                    )
                    
                    production_rates.append(prod_data['production_rate_g_per_hour'])
                    costs_per_kg.append(cost_result['cost_per_kg'])
                
                # Create comparison charts
                col1, col2 = st.columns(2)
                
                with col1:
                    # Production rate comparison
                    fig1 = px.bar(
                        x=methods,
                        y=production_rates,
                        title="Production Rate Comparison",
                        labels={"x": "Production Method", "y": "Production Rate (g/hour)"},
                        color=production_rates,
                        color_continuous_scale="Blues"
                    )
                    st.plotly_chart(fig1, use_container_width=True)
                
                with col2:
                    # Cost comparison
                    fig2 = px.bar(
                        x=methods,
                        y=costs_per_kg,
                        title="Cost per kg Comparison",
                        labels={"x": "Production Method", "y": "Cost ($/kg)"},
                        color=costs_per_kg,
                        color_continuous_scale="Reds_r"
                    )
                    st.plotly_chart(fig2, use_container_width=True)
                
                # Efficiency vs Cost scatter plot
                efficiencies = [0.65, 0.75, 0.85]  # Method efficiencies
                
                fig3 = px.scatter(
                    x=efficiencies,
                    y=costs_per_kg,
                    size=production_rates,
                    text=methods,
                    title="Efficiency vs Cost Trade-off",
                    labels={"x": "Efficiency", "y": "Cost ($/kg)"},
                    color=methods
                )
                
                # Add a vertical line for current efficiency
                current_efficiency = production_data['efficiency']
                fig3.add_shape(
                    type="line",
                    x0=current_efficiency,
                    y0=0,
                    x1=current_efficiency,
                    y1=max(costs_per_kg) * 1.1,
                    line=dict(color="red", width=2, dash="dash")
                )
                
                # Add annotation for current efficiency
                fig3.add_annotation(
                    x=current_efficiency,
                    y=max(costs_per_kg) * 0.9,
                    text="Current Efficiency",
                    showarrow=True,
                    arrowhead=1
                )
                
                st.plotly_chart(fig3, use_container_width=True)
                
                # Current density vs production rate
                current_densities = np.linspace(0.1, 2.0, 10)
                production_results = []
                
                for cd in current_densities:
                    prod_data = calculate_h2_production(
                        production_method, 
                        water_quantity, 
                        energy_input, 
                        cd, 
                        voltage
                    )
                    production_results.append(prod_data['production_rate_g_per_hour'])
                
                fig4 = px.line(
                    x=current_densities,
                    y=production_results,
                    title=f"Impact of Current Density on {production_method} Production Rate",
                    labels={"x": "Current Density (A/cm²)", "y": "Production Rate (g/hour)"},
                    markers=True
                )
                
                # Add a vertical line for current current density
                fig4.add_shape(
                    type="line",
                    x0=current_density,
                    y0=0,
                    x1=current_density,
                    y1=max(production_results) * 1.1,
                    line=dict(color="green", width=2, dash="dash")
                )
                
                st.plotly_chart(fig4, use_container_width=True)
                
                # Add a simulation over time
                st.subheader("Production Simulation Over Time")
                
                # Create time series data
                hours = list(range(0, int(production_data['operation_time_hours']) + 1))
                if len(hours) > 1:
                    production_over_time = [h * production_data['production_rate_g_per_hour'] for h in hours]
                    
                    fig5 = px.line(
                        x=hours,
                        y=production_over_time,
                        title="Cumulative Hydrogen Production Over Time",
                        labels={"x": "Time (hours)", "y": "Cumulative Production (g)"},
                        markers=True
                    )
                    st.plotly_chart(fig5, use_container_width=True)
                else:
                    st.info("Operation time too short for meaningful time series visualization.")
                
                # Export results button
                st.download_button(
                    label="Export Results as CSV",
                    data=pd.DataFrame({
                        "Parameter": ["Production Method", "Water Source", "Current Density (A/cm²)", "Voltage (V)",
                                    "Production Rate (g/h)", "Total Production (g)", "Efficiency (%)", 
                                    "Total Cost ($)", "Cost per kg ($/kg)"],
                        "Value": [production_method, water_source, current_density, voltage,
                                production_data['production_rate_g_per_hour'], production_data['total_production_g'],
                                production_data['efficiency']*100, cost_data['total_cost'], cost_data['cost_per_kg']]
                    }).to_csv(index=False),
                    file_name=f"hydrogen_production_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                    mime="text/csv"
                )

# Run the app
if __name__ == "__main__":
    main()