"""
Shape2Force (S2F) - GUI for force map prediction from bright field microscopy images.
"""
import os
import sys
import io
import cv2
cv2.utils.logging.setLogLevel(cv2.utils.logging.LOG_LEVEL_ERROR)

import numpy as np
import streamlit as st
from PIL import Image
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# Ensure S2F is in path
S2F_ROOT = os.path.dirname(os.path.abspath(__file__))
if S2F_ROOT not in sys.path:
    sys.path.insert(0, S2F_ROOT)

from utils.substrate_settings import list_substrates

st.set_page_config(page_title="Shape2Force (S2F)", page_icon="🦠", layout="centered")
st.markdown("""
    <style>
    section[data-testid="stSidebar"] { width: 380px !important; }
    </style>
    """, unsafe_allow_html=True)
st.title("🦠 Shape2Force (S2F)")
st.caption("Predict traction force maps from bright-field microscopy images of cells or spheroids")

# Folders: checkpoints in subfolders by model type (single_cell / spheroid)
ckp_base = os.path.join(S2F_ROOT, "ckp")
# Fallback: use project root ckp when running from S2F repo (ckp at S2F/ckp/)
if not os.path.isdir(ckp_base):
    project_root = os.path.dirname(S2F_ROOT)
    if os.path.isdir(os.path.join(project_root, "ckp")):
        ckp_base = os.path.join(project_root, "ckp")
ckp_single_cell = os.path.join(ckp_base, "single_cell")
ckp_spheroid = os.path.join(ckp_base, "spheroid")
sample_base = os.path.join(S2F_ROOT, "samples")
sample_single_cell = os.path.join(sample_base, "single_cell")
sample_spheroid = os.path.join(sample_base, "spheroid")

SAMPLE_EXTENSIONS = (".tif", ".tiff", ".png", ".jpg", ".jpeg")


def get_ckp_files_for_model(model_type):
    """Return list of .pth files in the checkpoint folder for the given model type."""
    folder = ckp_single_cell if model_type == "single_cell" else ckp_spheroid
    if os.path.isdir(folder):
        return sorted([f for f in os.listdir(folder) if f.endswith(".pth")])
    return []


def get_sample_files_for_model(model_type):
    """Return list of sample images in the sample folder for the given model type."""
    folder = sample_single_cell if model_type == "single_cell" else sample_spheroid
    if os.path.isdir(folder):
        return sorted([f for f in os.listdir(folder)
                       if f.lower().endswith(SAMPLE_EXTENSIONS)])
    return []

# Sidebar: model configuration
with st.sidebar:
    st.header("Model configuration")
    model_type = st.radio(
        "Model type",
        ["single_cell", "spheroid"],
        format_func=lambda x: "Single cell" if x == "single_cell" else "Spheroid",
        horizontal=False,
        help="Single cell: substrate-aware force prediction. Spheroid: spheroid force maps.",
    )
    st.caption(f"Inference mode: **{'Single cell' if model_type == 'single_cell' else 'Spheroid'}**")

    ckp_files = get_ckp_files_for_model(model_type)
    ckp_folder = ckp_single_cell if model_type == "single_cell" else ckp_spheroid
    ckp_subfolder_name = "single_cell" if model_type == "single_cell" else "spheroid"

    if ckp_files:
        checkpoint = st.selectbox(
            "Checkpoint",
            ckp_files,
            help=f"Select a .pth file from ckp/{ckp_subfolder_name}/",
        )
    else:
        st.warning(f"No .pth files in ckp/{ckp_subfolder_name}/. Add checkpoints to load.")
        checkpoint = None

    substrate_config = None
    substrate_val = "fibroblasts_PDMS"
    use_manual = False
    if model_type == "single_cell":
        try:
            substrates = list_substrates()
            substrate_val = st.selectbox(
                "Substrate (from config)",
                substrates,
                help="Select a preset from config/substrate_settings.json",
            )
            use_manual = st.checkbox("Enter substrate values manually", value=False)
            if use_manual:
                st.caption("Enter pixelsize (µm/px) and Young's modulus (Pa)")
                manual_pixelsize = st.number_input("Pixelsize (µm/px)", min_value=0.1, max_value=50.0,
                                                   value=3.0769, step=0.1, format="%.4f")
                manual_young = st.number_input("Young's modulus (Pa)", min_value=100.0, max_value=100000.0,
                                               value=6000.0, step=100.0, format="%.0f")
                substrate_config = {"pixelsize": manual_pixelsize, "young": manual_young}
            else:
                substrate_config = None
        except FileNotFoundError:
            st.error("config/substrate_settings.json not found")

# Main area: image input
img_source = st.radio("Image source", ["Upload", "Example"], horizontal=True, label_visibility="collapsed")
img = None
uploaded = None
selected_sample = None

if img_source == "Upload":
    uploaded = st.file_uploader(
        "Upload bright-field image",
        type=["tif", "tiff", "png", "jpg", "jpeg"],
        help="Bright-field microscopy image of a cell or spheroid on a substrate (grayscale or RGB). The model will predict traction forces from the cell shape.",
    )
    if uploaded:
        bytes_data = uploaded.read()
        nparr = np.frombuffer(bytes_data, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
        uploaded.seek(0)  # reset for potential re-read
else:
    sample_files = get_sample_files_for_model(model_type)
    sample_folder = sample_single_cell if model_type == "single_cell" else sample_spheroid
    sample_subfolder_name = "single_cell" if model_type == "single_cell" else "spheroid"
    if sample_files:
        selected_sample = st.selectbox(
            f"Select example image (from `samples/{sample_subfolder_name}/`)",
            sample_files,
            format_func=lambda x: x,
            key=f"sample_{model_type}",
        )
        if selected_sample:
            sample_path = os.path.join(sample_folder, selected_sample)
            img = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)
        # Show example thumbnails (filtered by model type)
        n_cols = min(5, len(sample_files))
        cols = st.columns(n_cols)
        for i, fname in enumerate(sample_files[:8]):  # show up to 8
            with cols[i % n_cols]:
                path = os.path.join(sample_folder, fname)
                sample_img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                if sample_img is not None:
                    st.image(sample_img, caption=fname, width='content')
    else:
        st.info(f"No example images in samples/{sample_subfolder_name}/. Add images or use Upload.")

col_btn, col_model, col_path = st.columns([1, 1, 3])
with col_btn:
    run = st.button("Run prediction", type="primary")
with col_model:
    model_label = "Single cell" if model_type == "single_cell" else "Spheroid"
    st.markdown(f"<span style='display: inline-flex; align-items: center; height: 38px;'>{model_label}</span>", unsafe_allow_html=True)
with col_path:
    ckp_path = f"ckp/{ckp_subfolder_name}/{checkpoint}" if checkpoint else f"ckp/{ckp_subfolder_name}/"
    st.markdown(f"<span style='display: inline-flex; align-items: center; height: 38px;'>Checkpoint: <code>{ckp_path}</code></span>", unsafe_allow_html=True)
has_image = img is not None

# Persist results in session state so they survive re-runs (e.g. when clicking Download)
if "prediction_result" not in st.session_state:
    st.session_state["prediction_result"] = None

# Show results if we just ran prediction OR we have cached results from a previous run
just_ran = run and checkpoint and has_image
cached = st.session_state["prediction_result"]
key_img = (uploaded.name if uploaded else None) if img_source == "Upload" else selected_sample
current_key = (model_type, checkpoint, key_img)
has_cached = cached is not None and cached.get("cache_key") == current_key

if just_ran:
    st.session_state["prediction_result"] = None  # Clear before new run
    with st.spinner("Loading model and predicting..."):
        try:
            from predictor import S2FPredictor
            predictor = S2FPredictor(
                model_type=model_type,
                checkpoint_path=checkpoint,
                ckp_folder=ckp_folder,
            )
            if img is not None:
                sub_val = substrate_val if model_type == "single_cell" and not use_manual else "fibroblasts_PDMS"
                heatmap, force, pixel_sum = predictor.predict(
                    image_array=img,
                    substrate=sub_val,
                    substrate_config=substrate_config if model_type == "single_cell" else None,
                )

                st.success("Prediction complete!")

                # Visualization - Plotly with zoom/pan, annotated (titles in Streamlit to avoid clipping)
                tit1, tit2 = st.columns(2)
                with tit1:
                    st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Input: Bright-field image</p>', unsafe_allow_html=True)
                with tit2:
                    st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Output: Predicted traction force map</p>', unsafe_allow_html=True)
                fig_pl = make_subplots(rows=1, cols=2)
                fig_pl.add_trace(go.Heatmap(z=img, colorscale="gray", showscale=False), row=1, col=1)
                fig_pl.add_trace(go.Heatmap(z=heatmap, colorscale="Jet", zmin=0, zmax=1, showscale=True,
                    colorbar=dict(len=0.4, thickness=12)), row=1, col=2)
                fig_pl.update_layout(
                    height=400,
                    margin=dict(l=10, r=10, t=10, b=10),
                    xaxis=dict(scaleanchor="y", scaleratio=1),
                    xaxis2=dict(scaleanchor="y2", scaleratio=1),
                )
                fig_pl.update_xaxes(showticklabels=False)
                fig_pl.update_yaxes(showticklabels=False, autorange="reversed")
                st.plotly_chart(fig_pl, use_container_width=True)

                # Metrics with help (below plot)
                col1, col2, col3, col4 = st.columns(4)
                with col1:
                    st.metric("Sum of all pixels", f"{pixel_sum:.2f}", help="Raw sum of all pixel values in the force map")
                with col2:
                    st.metric("Cell force (scaled)", f"{force:.2f}", help="Total traction force in physical units")
                with col3:
                    st.metric("Heatmap max", f"{np.max(heatmap):.4f}", help="Peak force intensity in the map")
                with col4:
                    st.metric("Heatmap mean", f"{np.mean(heatmap):.4f}", help="Average force intensity")

                # How to read (below numbers)
                with st.expander("ℹ️ How to read the results"):
                    st.markdown("""
**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
This is the raw image you provided—it shows cell shape but not forces.

**Output (right):** Predicted traction force map.  
- **Color** indicates force magnitude: blue = low, red = high  
- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
- Values are normalized to [0, 1] for visualization

**Metrics:**
- **Sum of all pixels:** Total force is the sum of all pixels in the force map. Each pixel represents the magnitude of force at that location; summing them gives the overall traction.
- **Cell force (scaled):** Total traction force in physical units (scaled by substrate stiffness)
- **Heatmap max/mean:** Peak and average force intensity in the map
                    """)

                # Download
                heatmap_uint8 = (np.clip(heatmap, 0, 1) * 255).astype(np.uint8)
                heatmap_rgb = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
                heatmap_rgb = cv2.cvtColor(heatmap_rgb, cv2.COLOR_BGR2RGB)
                pil_heatmap = Image.fromarray(heatmap_rgb)
                buf_hm = io.BytesIO()
                pil_heatmap.save(buf_hm, format="PNG")
                buf_hm.seek(0)
                st.download_button("Download Heatmap", data=buf_hm.getvalue(),
                                   file_name="s2f_heatmap.png", mime="image/png", key="download_heatmap")

                # Store in session state so results persist when user clicks Download
                cache_key = (model_type, checkpoint, key_img)
                st.session_state["prediction_result"] = {
                    "img": img.copy(),
                    "heatmap": heatmap.copy(),
                    "force": force,
                    "pixel_sum": pixel_sum,
                    "cache_key": cache_key,
                }

        except Exception as e:
            st.error(f"Prediction failed: {e}")
            import traceback
            st.code(traceback.format_exc())

elif has_cached:
    # Show cached results (e.g. after clicking Download)
    r = st.session_state["prediction_result"]
    img, heatmap, force, pixel_sum = r["img"], r["heatmap"], r["force"], r["pixel_sum"]
    st.success("Prediction complete!")
    tit1, tit2 = st.columns(2)
    with tit1:
        st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Input: Bright-field image</p>', unsafe_allow_html=True)
    with tit2:
        st.markdown('<p style="font-size: 1.1rem; color: black; font-weight: 600;">Output: Predicted traction force map</p>', unsafe_allow_html=True)
    fig_pl = make_subplots(rows=1, cols=2)
    fig_pl.add_trace(go.Heatmap(z=img, colorscale="gray", showscale=False), row=1, col=1)
    fig_pl.add_trace(go.Heatmap(z=heatmap, colorscale="Jet", zmin=0, zmax=1, showscale=True,
        colorbar=dict(len=0.4, thickness=12)), row=1, col=2)
    fig_pl.update_layout(height=400, margin=dict(l=10, r=10, t=10, b=10),
                         xaxis=dict(scaleanchor="y", scaleratio=1),
                         xaxis2=dict(scaleanchor="y2", scaleratio=1))
    fig_pl.update_xaxes(showticklabels=False)
    fig_pl.update_yaxes(showticklabels=False, autorange="reversed")
    st.plotly_chart(fig_pl, use_container_width=True)
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("Sum of all pixels", f"{pixel_sum:.2f}", help="Raw sum of all pixel values in the force map")
    with col2:
        st.metric("Cell force (scaled)", f"{force:.2f}", help="Total traction force in physical units")
    with col3:
        st.metric("Heatmap max", f"{np.max(heatmap):.4f}", help="Peak force intensity in the map")
    with col4:
        st.metric("Heatmap mean", f"{np.mean(heatmap):.4f}", help="Average force intensity")
    with st.expander("ℹ️ How to read the results"):
        st.markdown("""
**Input (left):** Bright-field microscopy image of a cell or spheroid on a substrate.  
This is the raw image you provided—it shows cell shape but not forces.

**Output (right):** Predicted traction force map.  
- **Color** indicates force magnitude: blue = low, red = high  
- **Brighter/warmer colors** = stronger forces exerted by the cell on the substrate  
- Values are normalized to [0, 1] for visualization

**Metrics:**
- **Sum of all pixels:** Total force is the sum of all pixels in the force map. Each pixel represents the magnitude of force at that location; summing them gives the overall traction.
- **Cell force (scaled):** Total traction force in physical units (scaled by substrate stiffness)
- **Heatmap max/mean:** Peak and average force intensity in the map
        """)
    heatmap_uint8 = (np.clip(heatmap, 0, 1) * 255).astype(np.uint8)
    heatmap_rgb = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
    heatmap_rgb = cv2.cvtColor(heatmap_rgb, cv2.COLOR_BGR2RGB)
    pil_heatmap = Image.fromarray(heatmap_rgb)
    buf_hm = io.BytesIO()
    pil_heatmap.save(buf_hm, format="PNG")
    buf_hm.seek(0)
    st.download_button("Download Heatmap", data=buf_hm.getvalue(),
                       file_name="s2f_heatmap.png", mime="image/png", key="download_cached")

elif run and not checkpoint:
    st.warning("Please add checkpoint files to the ckp/ folder and select one.")
elif run and not has_image:
    st.warning("Please upload an image or select an example.")

# Footer
st.sidebar.divider()
st.sidebar.caption(f"Examples: `samples/{ckp_subfolder_name}/`")
st.sidebar.caption("If you find this software useful, please cite:")
st.sidebar.caption(
    "Lautaro Baro, Kaveh Shahhosseini, Amparo Andrés-Bordería, Claudio Angione, and Maria Angeles Juanes. "
    "**\"Shape-to-force (S2F): Predicting Cell Traction Forces from LabelFree Imaging\"**, 2026."
)