import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px

# Page configuration
st.set_page_config(
    page_title="Effluent Treatment Load Calculator",
    page_icon="💧",
    layout="wide"
)

# Title and description
st.title("💧 Effluent Treatment Load Calculator")
st.markdown("""
This app calculates the **chemical dosage** required for the **secondary treatment** of wastewater 
based on the wastewater source, flow rate, pollutant loads, and turbidity.
""")

# Sidebar inputs
st.sidebar.header("📊 Input Parameters")

# Wastewater source selection
source_options = [
    "Textile Wastewater",
    "Pharmaceutical Wastewater",
    "Leather Industry Wastewater",
    "Pesticides Industry Wastewater",
    "Paints Wastewater"
]

source = st.sidebar.selectbox("Wastewater Source", source_options)

# Flow rate input
flow_rate = st.sidebar.number_input(
    "Flow Rate",
    min_value=0.0,
    value=100.0,
    step=10.0,
    help="In m³/day or kL/day"
)

# COD input
cod = st.sidebar.number_input(
    "COD (Chemical Oxygen Demand)",
    min_value=0.0,
    value=500.0,
    step=50.0,
    help="In mg/L"
)

# BOD input
bod = st.sidebar.number_input(
    "BOD (Biochemical Oxygen Demand)",
    min_value=0.0,
    value=250.0,
    step=25.0,
    help="In mg/L"
)

# Turbidity input
turbidity = st.sidebar.number_input(
    "Turbidity",
    min_value=0.0,
    value=100.0,
    step=10.0,
    help="In NTU (Nephelometric Turbidity Units)"
)

# Unit selection for output
unit = st.sidebar.radio("Output Unit", ["kg", "lb"])

# Chemical dosage coefficients (based on literature and typical treatment needs)
# Format: (COD_coefficient, BOD_coefficient, turbidity_coefficient, base_dosage)
# These coefficients represent kg of chemical per unit load per day
dosage_factors = {
    "Textile Wastewater": (0.25, 0.30, 0.02, 50),
    "Pharmaceutical Wastewater": (0.35, 0.40, 0.03, 80),
    "Leather Industry Wastewater": (0.20, 0.25, 0.04, 40),
    "Pesticides Industry Wastewater": (0.40, 0.45, 0.05, 100),
    "Paints Wastewater": (0.30, 0.35, 0.03, 70)
}

# Calculation function
def calculate_chemical_dosage(source, flow, cod, bod, turb):
    cod_coef, bod_coef, turb_coef, base = dosage_factors[source]
    
    # Calculate loads (kg/day)
    cod_load = (cod * flow) / 1000  # Convert mg/L * m³/day to kg/day
    bod_load = (bod * flow) / 1000
    turb_load = (turb * flow) / 1000
    
    # Chemical needed (kg/day)
    chemical_kg = (
        (cod_load * cod_coef) +
        (bod_load * bod_coef) +
        (turb_load * turb_coef) +
        base
    )
    
    return chemical_kg, cod_load, bod_load, turb_load

# Perform calculation
if flow_rate > 0:
    chemical_kg, cod_load, bod_load, turb_load = calculate_chemical_dosage(
        source, flow_rate, cod, bod, turbidity
    )
    
    # Convert to lb if needed
    chemical_lb = chemical_kg * 2.20462
    final_value = chemical_lb if unit == "lb" else chemical_kg
    unit_display = "lb" if unit == "lb" else "kg"
    
    # Display result prominently
    st.markdown("---")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.metric(
            label=f"🧪 Chemical Required (per day)",
            value=f"{final_value:.2f} {unit_display}"
        )
    st.markdown("---")
    
    # Show detailed breakdown
    st.subheader("📋 Detailed Load Breakdown")
    
    col_a, col_b, col_c, col_d = st.columns(4)
    col_a.metric("COD Load", f"{cod_load:.2f} kg/day")
    col_b.metric("BOD Load", f"{bod_load:.2f} kg/day")
    col_c.metric("Turbidity Load", f"{turb_load:.2f} kg/day")
    col_d.metric("Base Dosage", f"{dosage_factors[source][3]:.2f} kg/day")
    
    # Explanation of calculation
    with st.expander("📖 How is this calculated?"):
        st.markdown(f"""
        **Formula used:**  
        Chemical Required = (COD Load × {dosage_factors[source][0]}) + 
                          (BOD Load × {dosage_factors[source][1]}) + 
                          (Turbidity Load × {dosage_factors[source][2]}) + 
                          Base Dosage ({dosage_factors[source][3]} kg/day)
        
        **For {source}:**  
        - COD Load = {cod} mg/L × {flow_rate} m³/day ÷ 1000 = {cod_load:.2f} kg/day  
        - BOD Load = {bod} mg/L × {flow_rate} m³/day ÷ 1000 = {bod_load:.2f} kg/day  
        - Turbidity Load = {turbidity} NTU × {flow_rate} m³/day ÷ 1000 = {turb_load:.2f} kg/day  
        
        **Final Calculation:**  
        Chemical = ({cod_load:.2f} × {dosage_factors[source][0]}) + ({bod_load:.2f} × {dosage_factors[source][1]}) + 
                  ({turb_load:.2f} × {dosage_factors[source][2]}) + {dosage_factors[source][3]}  
        Chemical = **{chemical_kg:.2f} kg/day** ≈ **{chemical_lb:.2f} lb/day**
        """)
    
    # Visualization
    st.subheader("📊 Contribution to Chemical Dosage")
    contribution_data = {
        "Component": ["COD Contribution", "BOD Contribution", "Turbidity Contribution", "Base Dosage"],
        "Amount (kg/day)": [
            cod_load * dosage_factors[source][0],
            bod_load * dosage_factors[source][1],
            turb_load * dosage_factors[source][2],
            dosage_factors[source][3]
        ]
    }
    df_contrib = pd.DataFrame(contribution_data)
    
    fig = px.pie(df_contrib, values="Amount (kg/day)", names="Component", 
                 title="What makes up the total chemical dosage?",
                 color_discrete_sequence=px.colors.sequential.Blues_r)
    st.plotly_chart(fig, use_container_width=True)
    
    # Recommendation note
    st.info("💡 **Note:** These calculations are based on standard treatment guidelines for secondary treatment processes (e.g., coagulation, flocculation, or biological nutrient removal). Actual chemical dosage may vary based on site-specific conditions and treatability studies.")
    
else:
    st.warning("⚠️ Please enter a flow rate > 0 to calculate the chemical dosage.")

# Footer
st.markdown("---")
st.caption("Effluent Treatment Load Calculator | For secondary treatment chemical estimation | Deployment on Hugging Face")