import os
import pandas as pd
import streamlit as st
from fpdf import FPDF
from groq import Groq

# Initialize Groq client
client = Groq(
    api_key="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",
)

def analyze_water_quality(data, standards):
    """Analyze water quality data against standards."""
    results = []
    for parameter, value in data.items():
        standard = standards.get(parameter, None)
        if standard:
            exceeded = ((value - standard) / standard) * 100
            if value < standard:
                remark = "Below standard. Treatment needed."
                treatment = "Primary Treatment: Screening, Sedimentation, Grit Removal."
            else:
                remark = "Above standard. Treatment needed."
                if parameter in ["Bacteria"]:
                    treatment = "Tertiary Treatment: Disinfection using Chlorination or UV."
                elif parameter in ["Lead", "Mercury", "Arsenic"]:
                    treatment = "Advanced Treatment: Chemical Precipitation, Adsorption."
                elif parameter in ["Temperature"]:
                    treatment = "Cooling, Aeration."
                elif parameter in ["Clarity"]:
                    treatment = "Coagulation, Flocculation."
                elif parameter in ["Conductivity"]:
                    treatment = "Ion Exchange, Reverse Osmosis."
                elif parameter in ["pH"]:
                    treatment = "pH Adjustment: Lime or Acid Addition."
                elif parameter in ["Alkalinity"]:
                    treatment = "Alkalinity Control: Lime or Soda Ash."
                elif parameter in ["Chlorine"]:
                    treatment = "Dechlorination: Activated Carbon or Sulfite."
                elif parameter in ["Hardness"]:
                    treatment = "Softening: Ion Exchange or Reverse Osmosis."
                elif parameter in ["Dissolved Oxygen"]:
                    treatment = "Aeration."
                else:
                    treatment = "Preliminary Treatment: Removal of large debris."
            results.append({
                "Parameter": parameter,
                "Value": value,
                "Standard": standard,
                "Exceedance (%)": exceeded,
                "Remark": remark,
                "Treatment": treatment
            })
    return results

def generate_pdf_report(results, filename="Water_Quality_Report.pdf"):
    """Generate a PDF report for water quality analysis."""
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.cell(200, 10, txt="Water Quality Assessment Report", ln=True, align='C')
    pdf.ln(10)

    for result in results:
        pdf.cell(0, 10, txt=f"Parameter: {result['Parameter']}", ln=True)
        pdf.cell(0, 10, txt=f"Value: {result['Value']} | Standard: {result['Standard']} | Exceedance: {result['Exceedance (%)']:.2f}%", ln=True)
        pdf.cell(0, 10, txt=f"Remark: {result['Remark']}", ln=True)
        pdf.cell(0, 10, txt=f"Treatment: {result['Treatment']}", ln=True)
        pdf.ln(5)

    pdf.output(filename)

def main():
    st.title("Water Quality Assessment App")
    st.write("Upload water sample data and analyze it against Punjab, Pakistan water quality standards.")

    uploaded_file = st.file_uploader("Upload your water sample data (Excel format):", type=["xls", "xlsx"])

    if uploaded_file:
        # Load water sample data
        sample_data = pd.read_excel(uploaded_file)
        st.write("Sample Data:")
        st.dataframe(sample_data)

        # Debug: Show column names
        st.write("Columns in uploaded data:", sample_data.columns)

        # Ensure the column names are correct
        if 'Parameter' not in sample_data.columns or 'Value' not in sample_data.columns:
            st.error("The uploaded file must contain 'Parameter' and 'Value' columns.")
            return

        # Load water quality standards (to be replaced with your data)
        standards = {
            "Lead": 0.01,
            "Mercury": 0.001,
            "Bacteria": 0.0,
            "Temperature": 25,  # Example standard for temperature in Celsius
            "Clarity": 5,  # Example standard for clarity (NTU)
            "Conductivity": 1000,  # Example standard for conductivity (µS/cm)
            "pH": 7,  # Example standard for pH
            "Alkalinity": 200,  # Example standard for alkalinity (mg/L)
            "Chlorine": 0.5,  # Example standard for chlorine (mg/L)
            "Hardness": 200,  # Example standard for hardness (mg/L)
            "Dissolved Oxygen": 5  # Example standard for dissolved oxygen (mg/L)
        }

        # Analyze the data
        results = []
        for _, row in sample_data.iterrows():
            parameter = row['Parameter']
            value = row['Value']
            result = analyze_water_quality({parameter: value}, standards)
            results.extend(result)

        # Display results
        st.write("Analysis Results:")
        results_df = pd.DataFrame(results)
        st.dataframe(results_df)

        # Generate and download report
        if st.button("Generate PDF Report"):
            report_filename = "Water_Quality_Report.pdf"
            generate_pdf_report(results, filename=report_filename)
            st.success("Report generated successfully!")
            with open(report_filename, "rb") as file:
                st.download_button(
                    label="Download Report",
                    data=file,
                    file_name=report_filename,
                    mime="application/pdf"
                )

if __name__ == "__main__":
    main()