app.py

import math
from dataclasses import dataclass, asdict
from typing import Dict, List, Tuple, Callable

import numpy as np
import pandas as pd
import streamlit as st
import plotly.express as px

try:
    import graphviz
    GRAPHVIZ_OK = True
except Exception:
    GRAPHVIZ_OK = False


# =========================
# Domain model
# =========================
@dataclass
class SampleState:
    rna_ng: float
    dna_ng: float
    protein_units: float
    inhibitors_units: float

    degradation: float  # 0..1
    rnase_risk: float   # 0..1
    volume_ul: float

    a260_280: float = 0.0
    a260_230: float = 0.0
    integrity_score: float = 0.0  # 0..10
    warnings: str = ""


@dataclass
class StepResult:
    step_name: str
    duration_min: float
    rna_ng: float
    dna_ng: float
    protein_units: float
    inhibitors_units: float
    degradation: float
    note: str


# =========================
# Helpers
# =========================
def clamp(x: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, x))

def jitter_multiplier(rng: np.random.Generator, cv: float) -> float:
    cv = max(0.0, float(cv))
    if cv == 0.0:
        return 1.0
    sigma2 = math.log(1.0 + cv * cv)
    mu = -0.5 * sigma2
    return float(rng.lognormal(mean=mu, sigma=math.sqrt(sigma2)))

def compute_metrics(state: SampleState, ethanol_carryover: float, dna_contam_risk: float) -> SampleState:
    protein_penalty = 0.6 * (state.protein_units / (state.protein_units + 1.0))
    a260_280 = 2.05 - protein_penalty

    inhib_penalty = 1.2 * (state.inhibitors_units / (state.inhibitors_units + 1.0))
    ethanol_penalty = 0.8 * clamp(ethanol_carryover, 0.0, 1.0)
    a260_230 = 2.20 - inhib_penalty - ethanol_penalty

    integrity = 10.0 * (1.0 - clamp(state.degradation, 0.0, 1.0))
    integrity = clamp(integrity, 0.0, 10.0)

    warnings = []
    if state.rna_ng < 50:
        warnings.append("Low yield (RNA < 50 ng).")
    if a260_280 < 1.8:
        warnings.append("Possible protein contamination (A260/280 low).")
    if a260_230 < 1.8:
        warnings.append("Possible salt/phenol/ethanol/inhibitor carryover (A260/230 low).")
    if integrity < 6.0:
        warnings.append("RNA integrity may be poor (low integrity score).")
    if dna_contam_risk > 0.25:
        warnings.append("High DNA contamination risk (consider DNase).")

    state.a260_280 = float(a260_280)
    state.a260_230 = float(a260_230)
    state.integrity_score = float(integrity)
    state.warnings = " | ".join(warnings) if warnings else "None"
    return state


# =========================
# Step functions
# =========================
StepFn = Callable[[SampleState, Dict, np.random.Generator], Tuple[SampleState, str, Dict]]

def step_lysis(state: SampleState, params: Dict, rng: np.random.Generator):
    base = params["lysis_eff"]
    time_min = params["lysis_time_min"]
    mixing = params["mixing_quality"]

    eff = base * (0.85 + 0.15 * mixing) * (0.90 + 0.10 * clamp(time_min / 10.0, 0.0, 1.0))
    eff *= jitter_multiplier(rng, params["pipetting_cv"])

    state.rna_ng *= clamp(eff, 0.2, 1.05)

    rnase = state.rnase_risk
    temp_factor = params["temp_risk"]
    deg_increase = 0.02 + 0.08 * rnase * temp_factor * clamp(time_min / 15.0, 0.0, 1.0)
    deg_increase *= jitter_multiplier(rng, 0.25)
    state.degradation = clamp(state.degradation + deg_increase, 0.0, 1.0)

    state.protein_units *= 0.7
    note = f"Lysis eff ~{eff:.2f}; degradation +{deg_increase:.2f}"
    return state, note, {}

def step_phase_separation(state: SampleState, params: Dict, rng: np.random.Generator):
    sep = params["phase_sep_quality"]
    rna_ret = (0.80 + 0.18 * sep) * jitter_multiplier(rng, params["pipetting_cv"])
    protein_rem = (0.30 + 0.60 * sep)
    inhibitor_carry = (1.0 - sep) * 0.8 * jitter_multiplier(rng, 0.3)

    state.rna_ng *= clamp(rna_ret, 0.4, 1.0)
    state.protein_units *= clamp(1.0 - protein_rem, 0.05, 0.7)
    state.inhibitors_units += inhibitor_carry

    note = f"RNA retained ~{rna_ret:.2f}; phenol carryover +{inhibitor_carry:.2f}"
    return state, note, {}

def step_binding(state: SampleState, params: Dict, rng: np.random.Generator):
    base = params["binding_eff"]
    inhibitor_effect = 1.0 - 0.35 * (state.inhibitors_units / (state.inhibitors_units + 1.0))
    eff = base * inhibitor_effect * jitter_multiplier(rng, params["pipetting_cv"])

    state.rna_ng *= clamp(eff, 0.3, 1.0)

    dna_cobind = params["dna_cobind"] * jitter_multiplier(rng, 0.2)
    state.dna_ng *= clamp(0.60 + 0.40 * dna_cobind, 0.3, 1.0)

    note = f"Binding eff ~{eff:.2f}; DNA co-bind ~{dna_cobind:.2f}"
    return state, note, {"dna_contam_risk": float(dna_cobind)}

def step_wash(state: SampleState, params: Dict, rng: np.random.Generator):
    wash_strength = params["wash_strength"]
    wash_count = params["wash_count"]

    inhib_removal = 1.0 - math.exp(-0.9 * wash_strength * wash_count)
    prot_removal = 1.0 - math.exp(-0.6 * wash_strength * wash_count)
    rna_loss = (0.02 + 0.02 * wash_strength) * wash_count * jitter_multiplier(rng, 0.25)

    state.inhibitors_units *= clamp(1.0 - inhib_removal, 0.05, 0.9)
    state.protein_units *= clamp(1.0 - prot_removal, 0.05, 0.9)
    state.rna_ng *= clamp(1.0 - rna_loss, 0.5, 0.99)

    dry_time = params["dry_time_min"]
    ethanol_carryover = clamp(0.9 - 0.15 * dry_time, 0.0, 0.9) * jitter_multiplier(rng, 0.3)

    note = f"RNA loss ~{rna_loss:.2f}; ethanol carryover ~{ethanol_carryover:.2f}"
    return state, note, {"ethanol_carryover": float(ethanol_carryover)}

def step_dnase(state: SampleState, params: Dict, rng: np.random.Generator):
    if not params["dnase_on"]:
        return state, "DNase skipped", {}

    eff = (0.75 + 0.20 * params["dnase_quality"]) * jitter_multiplier(rng, params["pipetting_cv"])
    state.dna_ng *= clamp(1.0 - eff, 0.02, 0.5)
    state.rna_ng *= clamp(1.0 - (0.01 + 0.01 * (1.0 - params["dnase_quality"])), 0.95, 0.99)

    note = f"DNase removed ~{eff:.2f} of DNA; small RNA loss"
    return state, note, {"dna_contam_risk": float(1.0 - eff)}

def step_elution(state: SampleState, params: Dict, rng: np.random.Generator):
    vol = params["elution_ul"]
    elute_eff = (0.70 + 0.28 * params["elution_quality"]) * jitter_multiplier(rng, params["pipetting_cv"])
    state.rna_ng *= clamp(elute_eff, 0.4, 1.0)
    state.volume_ul = float(vol)

    if params["heat_elution"]:
        state.degradation = clamp(state.degradation + 0.01 * jitter_multiplier(rng, 0.4), 0.0, 1.0)
        note = f"Elution eff ~{elute_eff:.2f} (heated); volume {vol:.0f} µL"
    else:
        note = f"Elution eff ~{elute_eff:.2f}; volume {vol:.0f} µL"
    return state, note, {}


# =========================
# Protocols
# =========================
def build_protocol(protocol_name: str) -> List[Tuple[str, float, StepFn]]:
    if protocol_name == "Silica Column":
        return [
            ("Lysis", 10, step_lysis),
            ("Binding", 8, step_binding),
            ("Wash", 12, step_wash),
            ("DNase (optional)", 10, step_dnase),
            ("Elution", 5, step_elution),
        ]
    if protocol_name == "Magnetic Beads":
        return [
            ("Lysis", 10, step_lysis),
            ("Binding", 10, step_binding),
            ("Wash", 10, step_wash),
            ("DNase (optional)", 10, step_dnase),
            ("Elution", 5, step_elution),
        ]
    return [
        ("Lysis", 10, step_lysis),
        ("Phase Separation", 12, step_phase_separation),
        ("Binding", 8, step_binding),
        ("Wash", 12, step_wash),
        ("DNase (optional)", 10, step_dnase),
        ("Elution", 5, step_elution),
    ]


def init_sample(params: Dict, rng: np.random.Generator) -> SampleState:
    sample_type = params["sample_type"]
    input_amount = params["input_amount"]

    if sample_type == "Cultured cells":
        base_rna = 600 * input_amount
        base_dna = 300 * input_amount
        protein = 1.5 * input_amount
        inhib = 0.2
    elif sample_type == "Blood":
        base_rna = 250 * input_amount
        base_dna = 400 * input_amount
        protein = 2.0 * input_amount
        inhib = 0.6
    else:  # Tissue
        base_rna = 800 * input_amount
        base_dna = 500 * input_amount
        protein = 2.5 * input_amount
        inhib = 0.5

    base_rna *= jitter_multiplier(rng, 0.15)
    base_dna *= jitter_multiplier(rng, 0.15)

    return SampleState(
        rna_ng=float(base_rna),
        dna_ng=float(base_dna),
        protein_units=float(protein),
        inhibitors_units=float(inhib),
        degradation=float(params["initial_degradation"]),
        rnase_risk=float(params["rnase_risk"]),
        volume_ul=float(params["elution_ul"]),
    )


def run_once(protocol_name: str, params: Dict, seed: int) -> Tuple[SampleState, List[StepResult]]:
    rng = np.random.default_rng(int(seed))
    protocol = build_protocol(protocol_name)
    state = init_sample(params, rng)

    ethanol_carryover = 0.0
    dna_contam_risk = 0.0

    history: List[StepResult] = []
    for (step_name, base_dur, fn) in protocol:
        duration = float(base_dur * params["time_scale"])
        state, note, side = fn(state, params, rng)

        if "ethanol_carryover" in side:
            ethanol_carryover = max(ethanol_carryover, float(side["ethanol_carryover"]))
        if "dna_contam_risk" in side:
            dna_contam_risk = float(side["dna_contam_risk"])

        history.append(StepResult(
            step_name=step_name,
            duration_min=duration,
            rna_ng=float(state.rna_ng),
            dna_ng=float(state.dna_ng),
            protein_units=float(state.protein_units),
            inhibitors_units=float(state.inhibitors_units),
            degradation=float(state.degradation),
            note=note
        ))

    state = compute_metrics(state, ethanol_carryover=ethanol_carryover, dna_contam_risk=dna_contam_risk)
    return state, history


def run_monte_carlo(protocol_name: str, params: Dict, n: int, seed: int) -> pd.DataFrame:
    rows = []
    seed = int(seed)
    n = int(n)
    for i in range(n):
        s, _ = run_once(protocol_name, params, seed + i)
        conc = s.rna_ng / max(1e-6, s.volume_ul)
        rows.append({
            "run": i + 1,
            "rna_ng": s.rna_ng,
            "rna_conc_ng_per_ul": conc,
            "a260_280": s.a260_280,
            "a260_230": s.a260_230,
            "integrity_score": s.integrity_score,
            "degradation": s.degradation,
            "warnings": s.warnings
        })
    return pd.DataFrame(rows)


# =========================
# UI Styling
# =========================
st.set_page_config(page_title="RNA Isolation Simulator", page_icon="🧬", layout="wide")

CUSTOM_CSS = """
<style>
/* tighter page padding */
.block-container { padding-top: 1.2rem; padding-bottom: 2rem; }

/* nicer metric cards */
[data-testid="stMetric"] {
  background: rgba(255,255,255,0.03);
  border: 1px solid rgba(255,255,255,0.08);
  padding: 16px;
  border-radius: 16px;
}

/* sidebar spacing */
section[data-testid="stSidebar"] .block-container { padding-top: 1rem; }

/* subtle separators */
hr { border: none; border-top: 1px solid rgba(255,255,255,0.08); margin: 0.75rem 0; }

/* table rounding */
[data-testid="stDataFrame"] {
  border-radius: 16px;
  overflow: hidden;
  border: 1px solid rgba(255,255,255,0.08);
}
</style>
"""
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)


# =========================
# Sidebar Controls
# =========================
with st.sidebar:
    st.markdown("## 🧬 RNA Isolation Simulator")
    st.caption("Interactive workflow + QC simulator (heuristic model).")
    st.markdown("---")

    protocol_name = st.selectbox(
        "Protocol",
        ["Silica Column", "Magnetic Beads", "TRIzol-like + cleanup"],
        index=0,
        help="Choose a workflow structure."
    )

    st.markdown("### Sample")
    sample_type = st.selectbox("Sample type", ["Cultured cells", "Tissue", "Blood"], index=0)
    input_amount = st.slider("Input amount", 0.1, 5.0, 1.0, 0.1, help="Abstract units. Higher usually gives higher yield.")
    rnase_risk = st.slider("RNase risk", 0.0, 1.0, 0.2, 0.05)
    initial_degradation = st.slider("Initial degradation", 0.0, 1.0, 0.05, 0.01)

    st.markdown("### Operator / Environment")
    mixing_quality = st.slider("Mixing quality", 0.0, 1.0, 0.7, 0.05)
    temp_risk = st.slider("Handling temperature risk", 0.0, 1.0, 0.3, 0.05)
    pipetting_cv = st.slider("Pipetting variability (CV)", 0.0, 0.25, 0.08, 0.01)
    time_scale = st.slider("Time scale", 0.6, 1.6, 1.0, 0.05, help=">1 means slower/longer steps.")

    with st.expander("Method parameters", expanded=False):
        st.markdown("**Lysis**")
        lysis_eff = st.slider("Base lysis efficiency", 0.4, 1.0, 0.85, 0.01)
        lysis_time_min = st.slider("Lysis time (min)", 3, 20, 10, 1)

        st.markdown("**TRIzol phase separation (if applicable)**")
        phase_sep_quality = st.slider("Phase separation quality", 0.0, 1.0, 0.7, 0.05)

        st.markdown("**Binding**")
        binding_eff = st.slider("Base binding efficiency", 0.4, 1.0, 0.85, 0.01)
        dna_cobind = st.slider("DNA co-binding tendency", 0.0, 1.0, 0.4, 0.05)

        st.markdown("**Wash**")
        wash_strength = st.slider("Wash strength", 0.3, 1.0, 0.75, 0.05)
        wash_count = st.slider("Number of washes", 1, 4, 2, 1)
        dry_time_min = st.slider("Dry time after wash (min)", 0, 8, 3, 1)

        st.markdown("**DNase**")
        dnase_on = st.checkbox("Include DNase step", value=True)
        dnase_quality = st.slider("DNase effectiveness", 0.0, 1.0, 0.7, 0.05)

        st.markdown("**Elution**")
        elution_ul = st.slider("Elution volume (µL)", 10, 100, 30, 5)
        elution_quality = st.slider("Elution efficiency factor", 0.0, 1.0, 0.7, 0.05)
        heat_elution = st.checkbox("Heat elution", value=False)
    st.markdown("---")

    st.markdown("### Simulation")
    n_runs = st.slider("Monte Carlo runs", 20, 1000, 200, 20)
    base_seed = st.number_input("Random seed", min_value=1, max_value=10_000_000, value=42, step=1)


params = dict(
    sample_type=sample_type,
    input_amount=float(input_amount),
    rnase_risk=float(rnase_risk),
    initial_degradation=float(initial_degradation),
    mixing_quality=float(mixing_quality),
    temp_risk=float(temp_risk),
    pipetting_cv=float(pipetting_cv),
    time_scale=float(time_scale),
    lysis_eff=float(lysis_eff),
    lysis_time_min=float(lysis_time_min),
    phase_sep_quality=float(phase_sep_quality),
    binding_eff=float(binding_eff),
    dna_cobind=float(dna_cobind),
    wash_strength=float(wash_strength),
    wash_count=int(wash_count),
    dry_time_min=float(dry_time_min),
    dnase_on=bool(dnase_on),
    dnase_quality=float(dnase_quality),
    elution_ul=float(elution_ul),
    elution_quality=float(elution_quality),
    heat_elution=bool(heat_elution),
)


# =========================
# Header + main run
# =========================
st.markdown("# 🧬 RNA Isolation Simulator")
st.caption("A clean UI for exploring how protocol choices & handling can affect yield/purity/integrity (heuristic).")

final_state, history = run_once(protocol_name, params, seed=int(base_seed))
hist_df = pd.DataFrame([asdict(h) for h in history])

# KPI row
k1, k2, k3, k4, k5 = st.columns([1, 1, 1, 1, 1.3])
k1.metric("Yield (ng)", f"{final_state.rna_ng:.1f}")
k2.metric("Conc (ng/µL)", f"{(final_state.rna_ng / max(1e-6, final_state.volume_ul)):.2f}")
k3.metric("A260/280", f"{final_state.a260_280:.2f}")
k4.metric("A260/230", f"{final_state.a260_230:.2f}")
k5.metric("Integrity (0–10)", f"{final_state.integrity_score:.1f}")

# Warnings banner
if final_state.warnings != "None":
    st.warning(final_state.warnings)
else:
    st.success("QC looks OK for this run (no warnings).")


# =========================
# Tabs
# =========================
tab_overview, tab_steps, tab_mc, tab_compare = st.tabs(
    ["📌 Overview", "🧩 Step-by-step", "📈 Monte Carlo", "⚖️ Compare Protocols"]
)


# ---- Overview
with tab_overview:
    left, right = st.columns([1.05, 1.25], gap="large")

    with left:
        st.subheader("Protocol Flow")
        protocol_steps = build_protocol(protocol_name)

        if GRAPHVIZ_OK:
            dot = graphviz.Digraph()
            dot.attr(rankdir="LR")
            for idx, (name, dur, _fn) in enumerate(protocol_steps):
                dot.node(f"s{idx}", f"{name}\n~{dur*params['time_scale']:.0f} min")
                if idx > 0:
                    dot.edge(f"s{idx-1}", f"s{idx}")
            st.graphviz_chart(dot)
        else:
            st.info("Install system Graphviz for diagram: `brew install graphviz`")
            st.write(" → ".join([s[0] for s in protocol_steps]))

        st.subheader("Timeline")
        timeline_df = hist_df[["step_name", "duration_min"]].copy()
        timeline_df["start_min"] = timeline_df["duration_min"].cumsum().shift(1).fillna(0.0)
        timeline_df["end_min"] = timeline_df["duration_min"].cumsum()

        fig_tl = px.timeline(timeline_df, x_start="start_min", x_end="end_min", y="step_name")
        fig_tl.update_layout(height=320, xaxis_title="Minutes", yaxis_title="", margin=dict(l=10, r=10, t=10, b=10))
        st.plotly_chart(fig_tl, use_container_width=True)

    with right:
        st.subheader("Single-run QC snapshot")
        qc_df = pd.DataFrame([{
            "yield_ng": final_state.rna_ng,
            "conc_ng_per_ul": final_state.rna_ng / max(1e-6, final_state.volume_ul),
            "a260_280": final_state.a260_280,
            "a260_230": final_state.a260_230,
            "integrity_score": final_state.integrity_score,
            "degradation": final_state.degradation,
        }])
        st.dataframe(qc_df, use_container_width=True, hide_index=True)

        st.subheader("What drives the outputs (in this model)")
        st.markdown(
            """
- **Yield** is mainly impacted by lysis/binding efficiency and wash losses.
- **A260/280** drops when **protein contamination** remains.
- **A260/230** drops with **inhibitors/phenol/ethanol carryover**.
- **Integrity** drops with **RNase risk**, temperature risk, and longer exposure.
            """
        )


# ---- Step-by-step
with tab_steps:
    st.subheader("Step trace table")
    st.dataframe(hist_df, use_container_width=True, height=300)

    st.subheader("Step inspector")
    step_names = hist_df["step_name"].tolist()
    pick = st.selectbox("Pick a step to inspect", step_names, index=0)
    row = hist_df[hist_df["step_name"] == pick].iloc[0].to_dict()

    c1, c2, c3, c4 = st.columns(4)
    c1.metric("RNA (ng)", f"{row['rna_ng']:.1f}")
    c2.metric("DNA (ng)", f"{row['dna_ng']:.1f}")
    c3.metric("Inhibitors", f"{row['inhibitors_units']:.2f}")
    c4.metric("Degradation", f"{row['degradation']:.2f}")
    st.info(row["note"])

    # quick trend chart
    st.subheader("RNA across steps")
    fig_line = px.line(hist_df, x="step_name", y="rna_ng", markers=True)
    fig_line.update_layout(height=320, xaxis_title="", yaxis_title="RNA (ng)", margin=dict(l=10, r=10, t=10, b=10))
    st.plotly_chart(fig_line, use_container_width=True)


# ---- Monte Carlo
with tab_mc:
    st.subheader("Monte Carlo simulation")
    df = run_monte_carlo(protocol_name, params, n=int(n_runs), seed=int(base_seed))

    top = st.columns([1.2, 1.2, 1.2, 1.0])
    top[0].metric("Median yield (ng)", f"{df['rna_ng'].median():.1f}")
    top[1].metric("Median A260/280", f"{df['a260_280'].median():.2f}")
    top[2].metric("Median A260/230", f"{df['a260_230'].median():.2f}")
    top[3].metric("Median integrity", f"{df['integrity_score'].median():.1f}")

    st.markdown("#### Distributions")
    cA, cB = st.columns(2, gap="large")
    with cA:
        st.plotly_chart(px.histogram(df, x="rna_ng", nbins=40, title="Yield distribution (ng)"),
                        use_container_width=True)
        st.plotly_chart(px.histogram(df, x="integrity_score", nbins=30, title="Integrity distribution (0–10)"),
                        use_container_width=True)

    with cB:
        st.plotly_chart(px.scatter(df, x="a260_280", y="a260_230", hover_data=["rna_ng"], title="Purity scatter"),
                        use_container_width=True)
        st.plotly_chart(px.histogram(df, x="a260_230", nbins=40, title="A260/230 distribution"),
                        use_container_width=True)

    st.markdown("#### Results table + download")
    st.dataframe(df, use_container_width=True, height=280)

    csv_bytes = df.to_csv(index=False).encode("utf-8")
    st.download_button(
        "Download results (CSV)",
        data=csv_bytes,
        file_name="rna_isolation_sim_results.csv",
        mime="text/csv",
        use_container_width=True
    )


# ---- Compare protocols
with tab_compare:
    st.subheader("Compare protocols (same inputs, same seed)")
    base_protocols = ["Silica Column", "Magnetic Beads", "TRIzol-like + cleanup"]

    compare_runs = st.slider("Runs per protocol", 50, 600, 200, 50)
    comp_rows = []
    for p in base_protocols:
        comp = run_monte_carlo(p, params, n=int(compare_runs), seed=int(base_seed))
        comp["protocol"] = p
        comp_rows.append(comp)
    comp_df = pd.concat(comp_rows, ignore_index=True)

    st.markdown("#### Median metrics by protocol")
    med = comp_df.groupby("protocol")[["rna_ng", "a260_280", "a260_230", "integrity_score"]].median().reset_index()
    st.dataframe(med, use_container_width=True, hide_index=True)

    st.markdown("#### Yield comparison")
    st.plotly_chart(px.box(comp_df, x="protocol", y="rna_ng", points="outliers", title="Yield (ng)"),
                    use_container_width=True)

    st.markdown("#### Integrity comparison")
    st.plotly_chart(px.box(comp_df, x="protocol", y="integrity_score", points="outliers", title="Integrity (0–10)"),
                    use_container_width=True)

    st.markdown("#### Purity comparison")
    st.plotly_chart(px.scatter(comp_df, x="a260_280", y="a260_230", color="protocol", title="Purity scatter by protocol"),
                    use_container_width=True)


st.markdown("---")
st.caption(
    "This is a heuristic simulator for UI/prototyping. If you collect real NanoDrop/RIN data, "
    "you can fit these parameters to become more realistic."
)