app/components/sequence_viewer.py

"""Interactive sequence viewer with per-residue displacement heatmap."""
import streamlit as st
import numpy as np


# Amino acid property classification
AA_PROPS = {
    "A": "hydrophobic", "I": "hydrophobic", "L": "hydrophobic", "M": "hydrophobic",
    "F": "hydrophobic", "W": "hydrophobic", "V": "hydrophobic", "P": "hydrophobic",
    "D": "charged", "E": "charged", "K": "charged", "R": "charged", "H": "charged",
    "S": "polar", "T": "polar", "N": "polar", "Q": "polar", "C": "polar", "Y": "polar",
    "G": "special", "X": "unknown",
}

PROP_COLORS = {
    "hydrophobic": "#f59e0b",
    "charged": "#ef4444",
    "polar": "#10b981",
    "special": "#94a3b8",
    "unknown": "#64748b",
}


def render_sequence_viewer(
    seq: str,
    displacements: np.ndarray,
    coverage: np.ndarray = None,
    mode_label: str = "Mode 0",
    max_chars_per_row: int = 80,
):
    """Render interactive HTML sequence viewer with displacement coloring.

    Each residue is displayed as a colored cell where:
    - Background: displacement magnitude (white → red gradient)
    - Border: coverage (thick = low coverage)
    - Tooltip: residue info

    Args:
        seq: Amino acid sequence (1-letter codes)
        displacements: Per-residue displacement magnitudes (N,)
        coverage: Per-residue coverage (N,) in [0, 1]
        mode_label: Label for the current mode
        max_chars_per_row: Characters per row before wrapping
    """
    n = len(seq)
    if coverage is None:
        coverage = np.ones(n)

    max_d = displacements.max() + 1e-8

    html = f"""
    <style>
        .seq-container {{
            font-family: 'Consolas', 'Monaco', monospace;
            background: #1e1b4b;
            border-radius: 8px;
            padding: 12px;
            margin: 8px 0;
        }}
        .seq-header {{
            color: #a5b4fc;
            font-size: 13px;
            margin-bottom: 8px;
            font-weight: bold;
        }}
        .seq-row {{
            display: flex;
            flex-wrap: wrap;
            gap: 1px;
            margin-bottom: 2px;
        }}
        .res {{
            display: inline-flex;
            align-items: center;
            justify-content: center;
            width: 14px;
            height: 22px;
            font-size: 10px;
            font-weight: bold;
            border-radius: 2px;
            cursor: pointer;
            transition: transform 0.1s;
            position: relative;
        }}
        .res:hover {{
            transform: scale(1.8);
            z-index: 10;
            box-shadow: 0 0 8px rgba(99, 102, 241, 0.8);
        }}
        .res:hover::after {{
            content: attr(data-tooltip);
            position: absolute;
            top: -38px;
            left: 50%;
            transform: translateX(-50%);
            background: #312e81;
            color: white;
            padding: 4px 8px;
            border-radius: 4px;
            font-size: 10px;
            white-space: nowrap;
            z-index: 100;
            border: 1px solid #6366f1;
        }}
        .seq-ruler {{
            display: flex;
            gap: 1px;
            margin-bottom: 1px;
        }}
        .ruler-mark {{
            width: 14px;
            font-size: 7px;
            color: #64748b;
            text-align: center;
        }}
        .legend {{
            display: flex;
            gap: 16px;
            margin-top: 8px;
            font-size: 11px;
            color: #94a3b8;
        }}
        .legend-item {{
            display: flex;
            align-items: center;
            gap: 4px;
        }}
        .legend-swatch {{
            width: 12px;
            height: 12px;
            border-radius: 2px;
        }}
    </style>
    <div class="seq-container">
        <div class="seq-header">{mode_label} — Per-residue displacement ({n} residues)</div>
    """

    # Build rows
    for row_start in range(0, n, max_chars_per_row):
        row_end = min(row_start + max_chars_per_row, n)

        # Ruler
        html += '<div class="seq-ruler">'
        for i in range(row_start, row_end):
            if (i + 1) % 10 == 0:
                html += f'<div class="ruler-mark">{i + 1}</div>'
            elif (i + 1) % 5 == 0:
                html += '<div class="ruler-mark">·</div>'
            else:
                html += '<div class="ruler-mark"></div>'
        html += "</div>"

        # Residues
        html += '<div class="seq-row">'
        for i in range(row_start, row_end):
            aa = seq[i] if i < len(seq) else "X"
            d = displacements[i]
            c = coverage[i] if i < len(coverage) else 1.0
            t = d / max_d  # Normalized displacement

            # Background: displacement heatmap (dark purple → bright red)
            r = int(30 + 225 * t)
            g = int(27 + 20 * (1 - t))
            b = int(75 - 50 * t)
            bg = f"rgb({r},{g},{b})"

            # Text color: white for high displacement, light for low
            txt_color = "white" if t > 0.3 else "#a5b4fc"

            # Border: thicker = lower coverage
            border_w = max(0, int(3 * (1 - c)))
            border = f"{border_w}px solid #ef4444" if border_w > 0 else "none"

            prop = AA_PROPS.get(aa, "unknown")
            tooltip = f"{aa}{i+1} | {d:.3f}Å | cov={c:.2f} | {prop}"

            html += (
                f'<div class="res" style="background:{bg};color:{txt_color};'
                f'border:{border}" data-tooltip="{tooltip}">{aa}</div>'
            )
        html += "</div>"

    # Legend
    html += """
    <div class="legend">
        <div class="legend-item">
            <div class="legend-swatch" style="background:linear-gradient(90deg,#1e1b4b,#ff3030)"></div>
            Low → High displacement
        </div>
        <div class="legend-item">
            <div class="legend-swatch" style="border:2px solid #ef4444;background:none"></div>
            Red border = low coverage
        </div>
    </div>
    </div>
    """

    st.markdown(html, unsafe_allow_html=True)


def render_displacement_chart(
    displacements: dict,
    seq: str = "",
    coverage: np.ndarray = None,
):
    """Render interactive displacement profile chart using Plotly.

    Args:
        displacements: {mode_idx: np.ndarray of per-residue magnitudes}
        seq: Amino acid sequence
        coverage: Per-residue coverage
    """
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots

    n_modes = len(displacements)
    n_res = len(list(displacements.values())[0])
    residues = np.arange(1, n_res + 1)

    # Hover text with AA identity
    hover_text = [f"{seq[i] if i < len(seq) else '?'}{i+1}" for i in range(n_res)]

    fig = make_subplots(
        rows=2, cols=1, row_heights=[0.75, 0.25],
        shared_xaxes=True, vertical_spacing=0.08,
        subplot_titles=["Displacement by Mode", "Coverage"]
    )

    colors = ["#6366f1", "#ef4444", "#10b981", "#f59e0b", "#ec4899", "#8b5cf6"]

    for k, d in displacements.items():
        mags = np.linalg.norm(d, axis=1) if d.ndim == 2 else d
        fig.add_trace(go.Scatter(
            x=residues, y=mags,
            mode="lines",
            name=f"Mode {k} (μ={mags.mean():.3f}Å)",
            line=dict(color=colors[k % len(colors)], width=1.5),
            fill="tozeroy",
            fillcolor=colors[k % len(colors)].replace(")", ",0.1)").replace("rgb", "rgba"),
            text=hover_text,
            hovertemplate="%{text}<br>Displacement: %{y:.3f}Å<extra>Mode " + str(k) + "</extra>",
        ), row=1, col=1)

    # Coverage
    if coverage is not None:
        fig.add_trace(go.Scatter(
            x=residues, y=coverage[:n_res],
            mode="lines",
            name="Coverage",
            line=dict(color="#94a3b8", width=1.5),
            fill="tozeroy",
            fillcolor="rgba(148,163,184,0.15)",
            hovertemplate="%{x}<br>Coverage: %{y:.3f}<extra></extra>",
        ), row=2, col=1)

    fig.update_layout(
        template="plotly_dark",
        height=400,
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(30,27,75,0.5)",
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        margin=dict(l=50, r=20, t=40, b=30),
    )
    fig.update_xaxes(title_text="Residue", row=2, col=1)
    fig.update_yaxes(title_text="Displacement (Å)", row=1, col=1)
    fig.update_yaxes(title_text="Coverage", range=[0, 1.1], row=2, col=1)

    st.plotly_chart(fig, use_container_width=True, key=f"disp_chart")