import streamlit as st
import math
import numpy as np
import matplotlib.pyplot as plt

# Set page config with an icon and title
st.set_page_config(page_title="Pipe Sizing Helper", page_icon="🛠️")

# Custom style
st.markdown(
    """
    <style>
    .title {
        font-size: 2.5rem;
        color: #4CAF50;
        font-weight: bold;
    }
    .subtitle {
        font-size: 1.25rem;
        color: #333;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        border-radius: 8px;
        padding: 10px 20px;
        font-size: 1rem;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# App header
st.markdown('<div class="title">Pipe Sizing Helper</div>', unsafe_allow_html=True)
st.markdown(
    '<div class="subtitle">Get recommendations for pipe sizes based on flow rate and permissible velocity.</div>',
    unsafe_allow_html=True,
)

# App layout
st.image(
    "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a4/Pipe_icon.svg/1024px-Pipe_icon.svg.png",
    width=100,
)  # Add an icon image
st.divider()

# Input fields
st.header("Input Parameters")
flow_rate = st.number_input("Flow Rate (m³/s):", min_value=0.0, step=0.01, value=0.0)
velocity = st.number_input("Permissible Velocity (m/s):", min_value=0.1, step=0.01, value=0.1)

# Calculate pipe diameter when button is clicked
if st.button("Generate Recommended Pipe Diameter"):
    if velocity > 0:
        # Pipe diameter formula: D = sqrt(4 * Q / (π * V))
        diameter = math.sqrt((4 * flow_rate) / (math.pi * velocity))
        diameter_mm = diameter * 1000  # Convert meters to mm
        st.success(f"Recommended Pipe Diameter: {diameter_mm:.2f} mm")
    else:
        st.error("Permissible velocity must be greater than zero.")

# Graphical visualization section
st.header("Graphical Visualization")
st.write("Explore how pipe diameter changes with flow rate and velocity.")

# Generate ranges for flow rate and velocity
flow_rates = np.linspace(0.1, 10, 100)  # Flow rates from 0.1 to 10 m³/s
velocities = np.linspace(0.1, 5, 50)  # Velocities from 0.1 to 5 m/s

# Create a grid for flow rates and velocities
flow_rate_mesh, velocity_mesh = np.meshgrid(flow_rates, velocities)
pipe_diameters = np.sqrt(4 * flow_rate_mesh / (np.pi * velocity_mesh)) * 1000  # Diameter in mm

# Plotting
fig, ax = plt.subplots(figsize=(8, 6))
c = ax.contourf(flow_rate_mesh, velocity_mesh, pipe_diameters, cmap="viridis", levels=20)
plt.colorbar(c, label="Pipe Diameter (mm)")
ax.set_title("Pipe Diameter Heatmap")
ax.set_xlabel("Flow Rate (m³/s)")
ax.set_ylabel("Velocity (m/s)")

# Display plot in Streamlit
st.pyplot(fig)