import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

# ── style ──────────────────────────────────────────────────────────────────
BG       = '#f8fafc'
PANEL    = '#ffffff'
GRID     = '#e2e8f0'
ACCENT   = '#2563eb'
C_INC    = '#d97706'
C_ZS     = '#7c3aed'
C_RR     = '#059669'
C_TEXT   = '#1e293b'
C_DIM    = '#64748b'
C_BORDER = '#cbd5e1'

plt.rcParams.update({
    'figure.facecolor': BG, 'axes.facecolor': PANEL,
    'axes.edgecolor': C_BORDER, 'axes.labelcolor': C_DIM,
    'xtick.color': C_DIM, 'ytick.color': C_DIM,
    'grid.color': GRID, 'grid.linewidth': 0.8,
    'text.color': C_TEXT, 'font.family': 'monospace',
})

EP_RANGE = np.linspace(0, 15, 600)

def calc_paf(episodes, incidence, zscore, rr):
    return 1 - 1 / np.exp(episodes * incidence * zscore * rr)

# ── plot function ──────────────────────────────────────────────────────────
def draw(ep, inc, zs, rr):
    paf_curve = calc_paf(EP_RANGE, inc, zs, rr) * 100
    cur_paf   = calc_paf(ep, inc, zs, rr)
    exponent  = ep * inc * zs * rr

    fig, axes = plt.subplots(1, 2, figsize=(15, 6),
                             gridspec_kw={'width_ratios': [3, 1]})
    fig.patch.set_facecolor(BG)

    # ── left: curve ───────────────────────────────────────────────────────
    ax = axes[0]
    ax.set_facecolor(PANEL)
    for spine in ax.spines.values():
        spine.set_color(C_BORDER)

    # shadow + main curve
    ax.plot(EP_RANGE, paf_curve, color=ACCENT, linewidth=5,   alpha=0.15)
    ax.plot(EP_RANGE, paf_curve, color=ACCENT, linewidth=2.2)

    # crosshairs
    ax.axvline(ep,          color=ACCENT, linestyle='--', linewidth=1.2, alpha=0.6)
    ax.axhline(cur_paf*100, color=ACCENT, linestyle='--', linewidth=1.2, alpha=0.3)

    # dot at current point
    ax.scatter([ep], [cur_paf*100], color=ACCENT, s=80, zorder=5,
               edgecolors='white', linewidths=1.5)

    # shaded area under curve up to current episode
    mask = EP_RANGE <= ep
    ax.fill_between(EP_RANGE[mask], paf_curve[mask], color=ACCENT, alpha=0.08)

    ax.set_xlim(0, 15)
    ax.set_ylim(0, 102)
    ax.set_xlabel('Episodes',  fontsize=11, labelpad=8)
    ax.set_ylabel('PAF (%)',   fontsize=11, labelpad=8)
    ax.set_title('Population Attributable Fraction',
                 color=C_TEXT, fontsize=13, pad=12, fontweight='normal')
    ax.grid(True, linewidth=0.8)

    # annotation
    ax.annotate(
        f'  {cur_paf*100:.1f}%',
        xy=(ep, cur_paf*100),
        xytext=(min(ep + 1, 13), cur_paf*100 + 4),
        color=ACCENT, fontsize=15, fontweight='bold',
        arrowprops=dict(arrowstyle='->', color=ACCENT, lw=1.2)
    )

    # formula watermark
    ax.text(0.98, 0.05,
            'PAF = 1 − exp(−ep × inc × Δz × RR)',
            transform=ax.transAxes, fontsize=11,
            color=C_DIM, ha='right', va='bottom')

    # ── right: breakdown panel ────────────────────────────────────────────
    ax2 = axes[1]
    ax2.set_facecolor('#f1f5f9')
    for spine in ax2.spines.values():
        spine.set_color(C_BORDER)
    ax2.set_xticks([])
    ax2.set_yticks([])

    rows = [
        ('episodes',  f'{ep:.2f}',                ACCENT),
        ('incidence', f'{inc:.3f}',                C_INC),
        ('Δz / ep.',  f'{zs:.3f}',                 C_ZS),
        ('RR',        f'{rr:.3f}',                 C_RR),
        ('─' * 10,    '─' * 6,                    C_BORDER),
        ('exponent',  f'{exponent:.4f}',           '#475569'),
        ('exp(−x)',   f'{1/np.exp(exponent):.4f}', '#475569'),
        ('─' * 10,    '─' * 6,                    C_BORDER),
        ('PAF',       f'{cur_paf*100:.2f}%',       ACCENT),
    ]

    ax2.text(0.5, 0.97, 'BREAKDOWN', transform=ax2.transAxes,
             ha='center', fontsize=13, color=ACCENT, fontweight='bold')

    y_start = 0.88
    for i, (k, v, c) in enumerate(rows):
        y = y_start - i * 0.095
        ax2.text(0.08, y, k, transform=ax2.transAxes,
                 fontsize=11, color=C_DIM, va='center')
        ax2.text(0.92, y, v, transform=ax2.transAxes,
                 fontsize=11, color=c, va='center', ha='right', fontweight='bold')

    ax2.set_title('Parameters', color=C_TEXT, fontsize=13, pad=8)

    plt.tight_layout(pad=1.5)
    return fig

# ── Gradio UI ─────────────────────────────────────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(), title="PAF Visualizer") as demo:

    gr.HTML("""
    <div style="background:#eff6ff;border:1px solid #bfdbfe;border-radius:10px;
                padding:12px 20px;margin-bottom:12px;font-family:monospace;">
      <span style="color:#1e293b;font-size:20px;font-weight:600;">
        Population Attributable Fraction Visualizer
      </span><br>
      <span style="color:#64748b;font-size:13px;">
        PAF = 1 − exp(−episodes × incidence × Δz × RR)
      </span>
    </div>""")

    with gr.Row():
        with gr.Column():
            ep  = gr.Slider(0,    10,  value=4.2,  step=0.1,  label="Episodes")
            inc = gr.Slider(0.01,  1,  value=0.2,  step=0.01, label="Pathogen Incidence")
            zs  = gr.Slider(0.01,  1,  value=0.2,  step=0.01, label="Δz per Episode")
            rr  = gr.Slider(1.01,  3,  value=1.5,  step=0.05, label="Relative Risk (RR)")

            reset_btn = gr.Button("Reset Defaults", variant="secondary")

    plot = gr.Plot(label="")

    # live update on any slider change
    for slider in [ep, inc, zs, rr]:
        slider.change(fn=draw, inputs=[ep, inc, zs, rr], outputs=plot)

    # reset button
    def reset():
        return 5.0, 0.3, 1.2, 2.5

    reset_btn.click(fn=reset, outputs=[ep, inc, zs, rr])

    # draw on load
    demo.load(fn=draw, inputs=[ep, inc, zs, rr], outputs=plot)

demo.launch()