import streamlit as st
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from io import BytesIO

st.title("KPI Std. Deviation")

# KPI Dropdown with units
kpi_options = {
    'CRF': '%',
    'Feed Water Temp': '°C',
    'S:F': '',
    'SSC': 'kg/kg',
    'SWC': 'kg/kg',
    'Make up water': 'kL'
}
kpi_selected = st.selectbox("Select KPI", list(kpi_options.keys()))
unit = kpi_options[kpi_selected]

# Inputs for Min, Max, and Mean
col1, col2, col3 = st.columns(3)
with col1:
    min_val = st.number_input(f"Enter Min Value ({unit})", value=60.0)
with col2:
    mean_val = st.number_input(f"Enter Mean (Avg) Value ({unit})", value=100.0)
with col3:
    max_val = st.number_input(f"Enter Max Value ({unit})", value=140.0)

# Inputs for Probability Densities
col4, col5, col6 = st.columns(3)
with col4:
    min_pdf = st.number_input("Enter Min Value Probability Density", value=0.0)
with col5:
    mean_pdf = st.number_input("Enter Mean Value Probability Density", value=1.0)
with col6:
    max_pdf = st.number_input("Enter Max Value Probability Density", value=0.0)

# Validation
if max_val > mean_val > min_val:
    # Generate X values (200 points evenly spaced)
    x = np.linspace(min_val, max_val, 200)

    # Custom quadratic bell shape through interpolation of given PDF values
    A = np.array([
        [min_val**2, min_val, 1],
        [mean_val**2, mean_val, 1],
        [max_val**2, max_val, 1]
    ])
    B = np.array([min_pdf, mean_pdf, max_pdf])
    coeffs = np.linalg.solve(A, B)
    a, b, c = coeffs

    # Compute Y values based on custom quadratic
    y = a * x**2 + b * x + c
    y = np.maximum(y, 0)  # Ensure no negative values

    df = pd.DataFrame({
        f"{kpi_selected} Value ({unit})": x,
        "Probability Density": y
    })

    # Plot
    fig, ax = plt.subplots()
    ax.plot(x, y, color='royalblue', linewidth=2, label='Deviation Curve')

    # Markers for Min, Mean, Max
    ax.plot(min_val, min_pdf, 'o', color='blue')
    ax.plot(mean_val, mean_pdf, 'o', color='green')
    ax.plot(max_val, max_pdf, 'o', color='red')

    ax.annotate(f'{min_val} {unit}', (min_val, min_pdf), textcoords="offset points", xytext=(-10,10), ha='center', color='blue')
    ax.annotate(f'{mean_val} {unit}', (mean_val, mean_pdf), textcoords="offset points", xytext=(0,10), ha='center', color='green')
    ax.annotate(f'{max_val} {unit}', (max_val, max_pdf), textcoords="offset points", xytext=(10,10), ha='center', color='red')

    ax.set_title(f"{kpi_selected} - Std. Deviation Curve")
    ax.set_xlabel(f"{kpi_selected} ({unit})")
    ax.set_ylabel("Probability Density")
    ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
    st.pyplot(fig)

    # Download data as CSV
    csv = df.to_csv(index=False).encode('utf-8')
    st.download_button(
        label="Download Data as CSV",
        data=csv,
        file_name=f"{kpi_selected}_custom_bell_curve_data.csv",
        mime='text/csv'
    )

    # Download plot as PNG
    buf = BytesIO()
    fig.savefig(buf, format="png")
    st.download_button(
        label="Download Plot as PNG",
        data=buf.getvalue(),
        file_name=f"{kpi_selected}_custom_bell_curve_plot.png",
        mime="image/png"
    )
else:
    st.warning("Please ensure that: Min < Mean < Max to generate a valid bell curve.")