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

# Set page configuration
st.set_page_config(
    page_title="Pipe Size Calculator",
    page_icon="💧",
    layout="centered",
    initial_sidebar_state="expanded",
)

# Sidebar for input and app information
st.sidebar.title("Pipe Size Calculator 💧")
st.sidebar.markdown(
    """
    Calculate the **recommended pipe diameter** based on:
    - Flow Rate (\( Q \)) in cubic meters per second.
    - Velocity (\( V \)) in meters per second.
    
    Also visualize pipe diameters over a range of flow rates and velocities. 🌊
    """
)

# Main application title
st.title("💧 Pipe Size Calculator")

# Add layout with two columns
col1, col2 = st.columns(2)

# Input section
with col1:
    st.header("Input Parameters")
    flow_rate = st.number_input(
        "Enter Flow Rate (m³/s):", 
        min_value=0.0, 
        step=0.01, 
        format="%.2f"
    )
    velocity = st.number_input(
        "Enter Velocity (m/s):", 
        min_value=0.0, 
        step=0.01, 
        format="%.2f"
    )

# Add a button to calculate the diameter
with col2:
    st.header("Action")
    if st.button("Generate Pipe Diameter"):
        if flow_rate > 0 and velocity > 0:
            # Calculate pipe diameter
            diameter = math.sqrt((4 * flow_rate) / (math.pi * velocity))
            diameter_cm = diameter * 100  # Convert to cm for better readability
            
            # Display results in success box
            st.success(
                f"Recommended Pipe Diameter: \n"
                f"- **{diameter:.4f} meters** \n"
                f"- **{diameter_cm:.2f} centimeters**"
            )
        else:
            # Warning for invalid input
            st.error("Please enter positive values for both flow rate and velocity.")
    else:
        st.info("Click the button to calculate the pipe diameter.")

# Graphical visualization
st.header("Graphical Visualization 📊")
st.markdown(
    """
    Visualize how pipe diameter varies with different **flow rates** and **velocities**.
    """
)

# User input for visualization range
col3, col4 = st.columns(2)
with col3:
    flow_rate_range = st.slider(
        "Select Flow Rate Range (m³/s):",
        min_value=0.1, 
        max_value=5.0, 
        value=(0.5, 2.5),
        step=0.1
    )
with col4:
    velocity_range = st.slider(
        "Select Velocity Range (m/s):",
        min_value=0.1, 
        max_value=10.0, 
        value=(1.0, 5.0),
        step=0.5
    )

# Generate and plot graph
flow_rates = np.linspace(flow_rate_range[0], flow_rate_range[1], 50)
velocities = np.linspace(velocity_range[0], velocity_range[1], 50)

flow_grid, velocity_grid = np.meshgrid(flow_rates, velocities)
diameter_grid = np.sqrt((4 * flow_grid) / (np.pi * velocity_grid)) * 100  # Diameter in cm

fig, ax = plt.subplots(figsize=(8, 6))
contour = ax.contourf(flow_grid, velocity_grid, diameter_grid, cmap="coolwarm")
cbar = fig.colorbar(contour, ax=ax)
cbar.set_label("Pipe Diameter (cm)")
ax.set_title("Pipe Diameter vs. Flow Rate and Velocity")
ax.set_xlabel("Flow Rate (m³/s)")
ax.set_ylabel("Velocity (m/s)")
ax.grid(True)

st.pyplot(fig)

# Footer with additional info
st.sidebar.markdown("---")
st.sidebar.markdown(
    """
    **Note:** Ensure flow rate and velocity are within realistic ranges for best results. 🌟  
    Created with ❤️ using [Streamlit](https://streamlit.io).
    """
)