import os
import numpy as np
import plotly.graph_objects as go
import gradio as gr

def _cos_blend(x):
    x = np.clip(x, 0.0, 1.0)
    return 0.5 - 0.5*np.cos(np.pi*x)

def generate_cycle(HR=72, preload=1.0, afterload=1.0, inotropy=1.0):
    T = 60.0 / HR
    n = 1600
    t = np.linspace(0, T, n)
    tn = t / T

    EDV0, ESV0 = 120.0, 50.0
    EDV = EDV0 * (0.9 + 0.25*preload)
    ESV = max(10.0, ESV0 * (1.1 + 0.35*afterload - 0.45*inotropy))
    SV  = max(5.0, EDV - ESV)

    P_lv_peak0 = 120.0
    P_lv_peak  = P_lv_peak0 * (0.85 + 0.25*inotropy + 0.10*afterload)

    P_ao_sys0 = 120.0
    P_ao_sys  = P_ao_sys0 * (0.9 + 0.25*afterload + 0.10*inotropy)
    runoff = 0.75 + 0.25*(HR/72.0)
    P_ao_dia0 = 80.0
    P_ao_dia  = max(45.0, P_ao_dia0 * (0.8 + 0.4*afterload) * (0.9 + 0.2*runoff))

    P_la_mean0 = 10.0
    P_la_mean  = P_la_mean0 * (0.9 + 0.2*preload)

    dur_isoC = max(0.05, 0.08 - 0.015*(HR-72)/72)
    dur_ej   = max(0.18, 0.26 - 0.02*(afterload-1) + 0.03*(inotropy-1))
    dur_isoR = max(0.05, 0.08 - 0.015*(HR-72)/72)
    dur_fill = max(0.25, 0.42 - (dur_isoC+dur_ej+dur_isoR))
    dur_as   = min(0.18, 0.12 + 0.04*(preload-1))
    total = dur_isoC + dur_ej + dur_isoR + dur_fill
    s = 1.0 / total
    dur_isoC *= s; dur_ej *= s; dur_isoR *= s; dur_fill *= s
    dur_fill_early = max(0.05, dur_fill - dur_as)

    t0 = 0.0
    t_fill_early_end = t0 + dur_fill_early
    t_as_end   = t_fill_early_end + dur_as
    t_isoC_end = t_as_end + dur_isoC
    t_ej_end   = t_isoC_end + dur_ej
    t_isoR_end = t_ej_end + dur_isoR

    t_AV_open  = (t_isoR_end) % 1.0
    t_AV_close = t_as_end
    t_SL_open  = t_isoC_end
    t_SL_close = t_ej_end

    m_fillE = (tn >= 0.0) & (tn < t_fill_early_end)
    m_as    = (tn >= t_fill_early_end) & (tn < t_as_end)
    m_isoC  = (tn >= t_as_end) & (tn < t_isoC_end)
    m_ej    = (tn >= t_isoC_end) & (tn < t_ej_end)
    m_isoR  = (tn >= t_ej_end) & (tn < t_isoR_end)

    Vlv = np.full_like(tn, EDV)
    if m_ej.any():
        xe = (tn[m_ej] - t_isoC_end) / max(1e-6, (t_ej_end - t_isoC_end))
        Vlv[m_ej] = EDV - SV * _cos_blend(xe)
    if m_fillE.any():
        xf = (tn[m_fillE] - 0.0) / max(1e-6, t_fill_early_end - 0.0)
        Vstart, Vend = ESV, EDV - 0.12*SV
        Vlv[m_fillE] = Vstart + (Vend - Vstart) * _cos_blend(xf)
    if m_as.any():
        xa = (tn[m_as] - t_fill_early_end) / max(1e-6, dur_as)
        Vstart, Vend = EDV - 0.12*SV, EDV
        Vlv[m_as] = Vstart + (Vend - Vstart) * _cos_blend(xa)

    Plv = 5 + 7*(Vlv-ESV)/max(1.0, (EDV-ESV))
    if m_isoC.any():
        xc = (tn[m_isoC] - t_as_end) / max(1e-6, (t_isoC_end - t_as_end))
        Plv[m_isoC] = np.maximum(Plv[m_isoC], 20 + (P_lv_peak-20)*_cos_blend(xc))
    if m_ej.any():
        xe = (tn[m_ej] - t_isoC_end) / max(1e-6, (t_ej_end - t_isoC_end))
        Plv[m_ej] = np.maximum(Plv[m_ej], P_lv_peak * (1 - 0.35*xe))
    if m_isoR.any():
        xr = (tn[m_isoR] - t_ej_end) / max(1e-6, (t_isoR_end - t_ej_end))
        Plv[m_isoR] = np.maximum(5 + 7*(Vlv[m_isoR]-ESV)/max(1.0,(EDV-ESV)),
                                 P_lv_peak*(1 - _cos_blend(xr)) + 5*_cos_blend(xr))

    Pao = np.full_like(tn, P_ao_dia)
    if m_ej.any():
        Pao[m_ej] = np.minimum(Plv[m_ej] - 2.0, P_ao_sys)
    Pao[~m_ej] = np.maximum(Pao[~m_ej], P_ao_dia)

    Pla = np.full_like(tn, P_la_mean)
    if m_as.any():
        xa = (tn[m_as] - t_fill_early_end) / max(1e-6, dur_as)
        Pla[m_as] += 3.0*np.sin(np.pi*xa)**2
    if m_isoC.any():
        xc = (tn[m_isoC] - t_as_end) / max(1e-6, (t_isoC_end - t_as_end))
        Pla[m_isoC] += 1.4*np.sin(np.pi*xc)**2
    if m_ej.any():
        xv = (tn[m_ej] - t_isoC_end) / max(1e-6, (t_ej_end - t_isoC_end))
        Pla[m_ej] += 3.0*(xv**2)

    events_s = {
        "AV_open":  (t_AV_open*T) % T,
        "AV_close": (t_AV_close*T) % T,
        "SL_open":  (t_SL_open*T) % T,
        "SL_close": (t_SL_close*T) % T,
    }
    return t, Plv, Pao, Pla, events_s

def make_plot(HR, preload, afterload, inotropy, show_events):
    t, Plv, Pao, Pla, ev = generate_cycle(HR, preload, afterload, inotropy)
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=t, y=Pao, mode="lines", line=dict(color="black", width=3)))
    fig.add_trace(go.Scatter(x=t, y=Plv, mode="lines", line=dict(color="red", width=3)))
    fig.add_trace(go.Scatter(x=t, y=Pla, mode="lines", line=dict(color="blue", width=3)))

    if show_events:
        markers = [("1", ev["AV_close"]), ("2", ev["SL_open"]), ("3", ev["SL_close"]), ("4", ev["AV_open"])]
        ymax = max(Pao.max(), Plv.max(), Pla.max())
        for label, x in markers:
            fig.add_vline(x=x, line=dict(width=1, dash="dot", color="#777"))
            fig.add_annotation(x=x, y=ymax*1.02, text=label, showarrow=False, font=dict(size=14, color="#444"), yanchor="bottom")

    ymax = max(Pao.max(), Plv.max(), Pla.max())
    fig.update_yaxes(range=[0, max(140, ymax*1.1)], tickfont=dict(size=12), showline=True, mirror=True)
    fig.update_xaxes(tickfont=dict(size=12), showline=True, mirror=True)
    fig.update_layout(template="simple_white", height=420, margin=dict(l=40, r=10, t=10, b=40), showlegend=False)
    return fig

with gr.Blocks(title="Wiggers Diagram (Minimal)") as demo:
    gr.Markdown("### Wiggers Diagram (Minimal)")
    with gr.Row():
        HR = gr.Slider(40, 180, value=72, step=1, label="HR (bpm)")
        preload = gr.Slider(0.6, 1.6, value=1.0, step=0.05, label="Preload")
        afterload = gr.Slider(0.6, 1.6, value=1.0, step=0.05, label="Afterload")
        inotropy = gr.Slider(0.5, 1.8, value=1.0, step=0.05, label="Contractility")
        show_events = gr.Checkbox(value=False, label="Show event numbers (1–4)")

    plot = gr.Plot(value=make_plot(72, 1.0, 1.0, 1.0, False))
    for w in (HR, preload, afterload, inotropy, show_events):
        w.change(make_plot, [HR, preload, afterload, inotropy, show_events], plot)

if __name__ == "__main__":
    # Required for Spaces
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))