"""Multi-track Interactive Residue Inspector (formerly sequence_viewer)."""
import streamlit as st
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# 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 parse_b_factors(pdb_text: str) -> np.ndarray:
    """Extract CA B-factors from PDB text."""
    b_factors = []
    if not pdb_text:
        return None
    for line in pdb_text.splitlines():
        if line.startswith("ATOM  ") and line[12:16].strip() == "CA":
            try:
                bf = float(line[60:66].strip())
                b_factors.append(bf)
            except ValueError:
                b_factors.append(0.0)
    return np.array(b_factors) if b_factors else None


def render_sequence_viewer(seq: str, displacements: np.ndarray, coverage: np.ndarray = None, mode_label: str = "Mode 0"):
    """Legacy HTML sequence viewer kept for compatibility, but enhanced."""
    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; }}
    </style>
    <div class="seq-container">
        <div class="seq-header">{mode_label} — Per-residue displacement ({n} residues)</div>
    """

    for row_start in range(0, n, 80):
        row_end = min(row_start + 80, n)
        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>"
        
        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
            r, g, b = int(30 + 225 * t), int(27 + 20 * (1 - t)), int(75 - 50 * t)
            txt_color = "white" if t > 0.3 else "#a5b4fc"
            border_w = max(0, int(3 * (1 - c)))
            border = f"{border_w}px solid #ef4444" if border_w > 0 else "none"
            tooltip = f"{aa}{i+1} | {d:.3f}Å | cov={c:.2f} | {AA_PROPS.get(aa, 'unknown')}"
            html += f'<div class="res" style="background:rgb({r},{g},{b});color:{txt_color};border:{border}" data-tooltip="{tooltip}">{aa}</div>'
        html += "</div>"
    html += "</div>"
    st.markdown(html, unsafe_allow_html=True)


def render_residue_inspector(displacements: dict, seq: str = "", coverage: np.ndarray = None, pdb_text: str = None, conservation: np.ndarray = None):
    """
    Render powerful multi-track Residue Inspector.
    Tracks:
    1. Sequence / Coverage
    2. Displacement by Mode
    3. Structural B-Factors (from PDB)
    4. Conservation (Evolutionary) -> Placeholder/Mock if None
    """
    n_res = len(list(displacements.values())[0])
    residues = np.arange(1, n_res + 1)
    hover_text = [f"{seq[i] if i < len(seq) else '?'}{i+1}" for i in range(n_res)]

    # Parse B-factors if pdb is provided
    b_factors = parse_b_factors(pdb_text) if pdb_text else None
    
    # Determine how many rows we need
    rows = 2
    titles = ["Predicted Mobility (Å)", "Coverage"]
    if b_factors is not None and len(b_factors) == n_res:
        rows += 1
        titles.append("Experimental B-Factors")
    if conservation is not None:
        rows += 1
        titles.append("Known Natural Variants (Mutations)")
        
    row_heights = [0.5] + [(0.5 / (rows - 1))] * (rows - 1)

    fig = make_subplots(
        rows=rows, cols=1, row_heights=row_heights,
        shared_xaxes=True, vertical_spacing=0.06,
        subplot_titles=titles
    )

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

    # ── Track 1: Displacement modes ──
    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}",
            line=dict(color=colors[k % len(colors)], width=2),
            fill="tozeroy",
            fillcolor=colors[k % len(colors)].replace(")", ",0.05)").replace("rgb", "rgba"),
            text=hover_text, hovertemplate="%{text}<br>Mobility: %{y:.3f}Å<extra>Mode " + str(k) + "</extra>",
        ), row=1, col=1)

    # ── Track 2: 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)
        
    current_row = 3
    
    # ── Track 3: B-Factors ──
    if b_factors is not None and len(b_factors) == n_res:
        fig.add_trace(go.Scatter(
            x=residues, y=b_factors, mode="lines",
            name="B-Factor", line=dict(color="#ec4899", width=1.5),
            fill="tozeroy", fillcolor="rgba(236,72,153,0.15)",
            hovertemplate="%{x}<br>B-Factor: %{y:.1f}<extra></extra>",
        ), row=current_row, col=1)
        current_row += 1

    # ── Track 4: Mutations (Conservation arg) ──
    if conservation is not None:
        fig.add_trace(go.Scatter(
            x=residues, y=conservation, mode="lines",
            name="Variants", line=dict(color="#14b8a6", width=1.5),
            fill="tozeroy", fillcolor="rgba(20,184,166,0.15)",
            hovertemplate="%{x}<br>Variants: %{y:.0f}<extra></extra>",
        ), row=current_row, col=1)

    # Decorate X-axis with sequence for zoom
    # We add rug plot at bottom if sequence is provided
    fig.update_layout(
        template="plotly_dark", height=200 + (rows * 150),
        paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(30,27,75,0.4)",
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        margin=dict(l=50, r=20, t=40, b=30),
        hovermode="x unified",
    )
    fig.update_xaxes(title_text="Residue Index", row=rows, col=1)
    
    st.plotly_chart(fig, use_container_width=True, key=f"residue_inspector")


# Aliasing the old function name so we don't break existing imports, 
# but calling the new powerful multi-track one.
def render_displacement_chart(displacements: dict, seq: str = "", coverage: np.ndarray = None, pdb_text: str = None, conservation: np.ndarray = None):
    render_residue_inspector(displacements, seq, coverage, pdb_text, conservation)