app.py

import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Arc, Rectangle
from matplotlib.lines import Line2D

st.set_page_config(page_title="Flow Down Gradients — Simple Explorer", layout="wide")

# ---------------- Pressure presets (mmHg) ----------------
PRESSURE_PRESETS = {
    "Demo Preset": {
        "cap_o2": 52, "cap_co2": 49,
        "alv_o2": 55, "alv_co2": 45,
        "thoracic": 761,
    },
    "Preset 1": {
        "cap_o2": 50, "cap_co2": 45,
        "alv_o2": 100, "alv_co2": 40,
        "thoracic": 763,
    },
    "Preset 2": {
        "cap_o2": 60, "cap_co2": 45,
        "alv_o2": 58, "alv_co2": 40,
        "thoracic": 757,
    },
    "Preset 3": {
        "cap_o2": 48, "cap_co2": 45,
        "alv_o2": 53, "alv_co2": 8,
        "thoracic": 768,
    },
}

def get_pressure_preset_from_query():
    qp = st.query_params
    key = qp.get("preset") if "preset" in qp else None
    if key in PRESSURE_PRESETS:
        return key
    return "Demo Preset"

if "pressure_preset" not in st.session_state:
    st.session_state.pressure_preset = get_pressure_preset_from_query()

# ---------------- Action Potential presets ----------------
AP_PRESETS = {
    "Preset 1": {"more_na": "Extracellular", "more_k": "Intracellular", "more_neg": "Intracellular"},
    "Preset 2": {"more_na": "Intracellular", "more_k": "Extracellular", "more_neg": "Extracellular"},
    "Preset 3": {"more_na": "Extracellular", "more_k": "Intracellular", "more_neg": "Extracellular"},
}

if "ap_preset" not in st.session_state:
    st.session_state.ap_preset = "Preset 1"

# ---------------- UI ----------------
st.title("Flow Down Gradients — Simple Explorer")

tab1, tab2 = st.tabs([
    "Pressure (Alveolus ↔ Capillary)",
    "Concentration: Action Potential (extracellular ↔ intracellular)",
])

# ---------------- Tab 1: Pressure ----------------
with tab1:
    st.subheader("Pressure Gradients (mmHg)")
    # Preset buttons with unique keys
    bcols = st.columns(4)
    for i, (name, col) in enumerate(zip(PRESSURE_PRESETS.keys(), bcols)):
        if col.button(name, key=f"pressure_btn_{i}_{name}", use_container_width=True):
            st.session_state.pressure_preset = name

    # Current preset
    preset_name = st.session_state.pressure_preset
    p = PRESSURE_PRESETS[preset_name]

    # Values
    st.markdown(f"**Current preset:** {preset_name}")
    st.markdown(
        f"- **Intracapillary O₂:** {p['cap_o2']} mmHg  \n"
        f"- **Intracapillary CO₂:** {p['cap_co2']} mmHg  \n"
        f"- **Intraalveolar O₂:** {p['alv_o2']} mmHg  \n"
        f"- **Intraalveolar CO₂:** {p['alv_co2']} mmHg  \n"
        f"- **Thoracic pressure:** {p['thoracic']} mmHg"
    )

    # Diagram (enlarged, placed below values)
    fig, ax = plt.subplots(figsize=(11.5, 5.6))
    ax.set_xlim(0, 12); ax.set_ylim(0, 6); ax.axis("off")

    # Alveolus: round sac
    alve_center = (4.0, 3.0); alve_r = 2.1
    alve = Circle(alve_center, alve_r, fill=False, lw=2)
    ax.add_patch(alve)
    ax.text(alve_center[0], alve_center[1] + alve_r + 0.35, "Alveolus", ha="center", fontsize=13)

    # Capillary: thin vessel hugging right side (two arcs + short straight segments)
    cap_outer = Arc((6.2, 3.0), 3.6, 3.6, angle=0, theta1=-65, theta2=65, lw=6, capstyle="round")
    cap_inner = Arc((6.2, 3.0), 3.0, 3.0, angle=0, theta1=-65, theta2=65, lw=6, capstyle="round")
    ax.add_patch(cap_outer); ax.add_patch(cap_inner)
    # Straight ends to right
    ax.plot([7.8, 11.0], [4.2, 4.2], lw=6, solid_capstyle="round")
    ax.plot([7.8, 11.0], [1.8, 1.8], lw=6, solid_capstyle="round")
    ax.text(11.2, 3.0, "Capillary", va="center", fontsize=13, rotation=90)

    # Accessibility-friendly markers/colors
    color_o2, marker_o2 = "tab:blue", "o"   # O2 = blue circle
    color_co2, marker_co2 = "tab:orange", "s"  # CO2 = orange square

    # Dot counts scaled from mmHg (cap at 150 points per region)
    maxdots = 150
    ao2 = int(np.clip(p["alv_o2"], 0, maxdots))
    aco2 = int(np.clip(p["alv_co2"], 0, maxdots))
    co2 = int(np.clip(p["cap_o2"], 0, maxdots))
    cco2 = int(np.clip(p["cap_co2"], 0, maxdots))

    # Alveolar dots: random inside circle
    rngA = np.random.default_rng(101)
    need = ao2 + aco2
    if need > 0:
        xs, ys = [], []
        while len(xs) < need:
            x = rngA.uniform(alve_center[0]-alve_r+0.12, alve_center[0]+alve_r-0.12)
            y = rngA.uniform(alve_center[1]-alve_r+0.12, alve_center[1]+alve_r-0.12)
            if (x-alve_center[0])**2 + (y-alve_center[1])**2 <= (alve_r-0.18)**2:
                xs.append(x); ys.append(y)
        xs = np.array(xs); ys = np.array(ys)
        if ao2 > 0:
            ax.scatter(xs[:ao2], ys[:ao2], s=36, c=color_o2, marker=marker_o2)
        if aco2 > 0:
            ax.scatter(xs[ao2:], ys[ao2:], s=36, c=color_co2, marker=marker_co2)

    # Capillary dots: right rectangular segment for clarity
    rngC = np.random.default_rng(202)
    cx = rngC.uniform(8.2, 10.8, co2 + cco2)
    cy = rngC.uniform(1.95, 4.05, co2 + cco2)
    if co2 > 0:
        ax.scatter(cx[:co2], cy[:co2], s=36, c=color_o2, marker=marker_o2)
    if cco2 > 0:
        ax.scatter(cx[co2:], cy[co2:], s=36, c=color_co2, marker=marker_co2)

    # Clean, accessible legend using Line2D handles (not cramped)
    legend_elements = [
        Line2D([0], [0], marker=marker_o2, color='w', label='O₂ (Blue • Circle)',
               markerfacecolor=color_o2, markersize=10),
        Line2D([0], [0], marker=marker_co2, color='w', label='CO₂ (Orange • Square)',
               markerfacecolor=color_co2, markersize=10),
    ]
    ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(0.02, 0.98), frameon=False)

    # Stamp preset name
    ax.text(0.98, 0.05, f"{preset_name}", transform=ax.transAxes, ha="right", va="bottom", fontsize=11, alpha=0.8)

    st.pyplot(fig, clear_figure=True)

# ---------------- Tab 2: Action Potential (presets, no sliders) ----------------
with tab2:
    st.subheader("Static starting conditions (no animation)")
    # Preset buttons with unique keys
    bcols = st.columns(3)
    for i, (name, col) in enumerate(zip(AP_PRESETS.keys(), bcols)):
        if col.button(name, key=f"ap_btn_{i}_{name}", use_container_width=True):
            st.session_state.ap_preset = name

    ap = AP_PRESETS[st.session_state.ap_preset]
    st.markdown(f"**Current preset:** {st.session_state.ap_preset}")
    st.markdown(
        f"- More **Na⁺**: {ap['more_na']}  \n"
        f"- More **K⁺**: {ap['more_k']}  \n"
        f"- More **negative**: {ap['more_neg']}"
    )

    fig, ax = plt.subplots(figsize=(11.5, 5.0))
    ax.set_xlim(0, 12); ax.set_ylim(0, 6); ax.axis("off")

    # Compartments
    ax.add_patch(Rectangle((0.8, 1.0), 4.5, 4.0, fill=False, lw=2))  # Extracellular
    ax.add_patch(Rectangle((6.7, 1.0), 4.5, 4.0, fill=False, lw=2))  # Intracellular
    ax.plot([6.2, 6.2], [1.0, 5.0], lw=8)  # Membrane

    ax.text(3.05, 5.25, "Extracellular", ha="center", fontsize=12)
    ax.text(8.95, 5.25, "Intracellular", ha="center", fontsize=12)

    # Polarity signs ABOVE labels
    if ap["more_neg"] == "Intracellular":
        ax.text(3.05, 5.85, "+", ha="center", va="center", fontsize=24, alpha=0.6)
        ax.text(8.95, 5.85, "−", ha="center", va="center", fontsize=24, alpha=0.9)
    else:
        ax.text(3.05, 5.85, "−", ha="center", va="center", fontsize=24, alpha=0.9)
        ax.text(8.95, 5.85, "+", ha="center", va="center", fontsize=24, alpha=0.6)

    # Colors + shapes (accessibility): Na+ = blue circle, K+ = orange square
    color_na, marker_na = "tab:blue", "o"
    color_k,  marker_k  = "tab:orange", "s"

    # Dot counts: base + extra on "more" side
    base_amt = 12
    extra_amt = 12
    na_out = base_amt + (extra_amt if ap["more_na"] == "Extracellular" else 0)
    na_in  = base_amt + (extra_amt if ap["more_na"] == "Intracellular" else 0)
    k_out  = base_amt + (extra_amt if ap["more_k"]  == "Extracellular" else 0)
    k_in   = base_amt + (extra_amt if ap["more_k"]  == "Intracellular" else 0)

    rng2 = np.random.default_rng(808)
    xo = rng2.uniform(1.1, 5.0, na_out + k_out); yo = rng2.uniform(1.2, 4.8, na_out + k_out)
    xi = rng2.uniform(7.0, 11.4, na_in  + k_in ); yi = rng2.uniform(1.2, 4.8, na_in  + k_in )

    # Draw ions with fixed color+shape per type, regardless of side
    if na_out > 0:
        ax.scatter(xo[:na_out], yo[:na_out], s=45, c=color_na, marker=marker_na, label="Na⁺")
    if k_out > 0:
        ax.scatter(xo[na_out:], yo[na_out:], s=45, c=color_k, marker=marker_k, label="K⁺")
    if na_in > 0:
        ax.scatter(xi[:na_in], yi[:na_in], s=45, c=color_na, marker=marker_na)
    if k_in > 0:
        ax.scatter(xi[na_in:], yi[na_in:], s=45, c=color_k, marker=marker_k)

    # Legend
    legend_elements = [
        Line2D([0], [0], marker=marker_na, color='w', label='Na⁺ (Blue • Circle)',
               markerfacecolor=color_na, markersize=10),
        Line2D([0], [0], marker=marker_k, color='w', label='K⁺ (Orange • Square)',
               markerfacecolor=color_k, markersize=10),
    ]
    ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(0.01, 0.98), frameon=False)

    # Stamp preset
    ax.text(0.98, 0.05, f"{st.session_state.ap_preset}", transform=ax.transAxes,
            ha="right", va="bottom", fontsize=11, alpha=0.8)

    st.pyplot(fig, clear_figure=True)