inference/metrics_visualization.py

"""
SYNTAX predictions visualization:
- points (SYNTAX ground truth vs model predictions) for multiple datasets;
- risk zones (low / high risk);
- ±σ and ±2σ bands around the diagonal;
- logistic trends for each dataset.

The script is independent of PyTorch/Lightning and is used at inference time.
Output is saved to the `visualizations/` folder inside the project.
"""

import os
import numpy as np
import plotly.graph_objects as go
from scipy.optimize import curve_fit  # type: ignore


# ================= GLOBAL STYLE CONSTANTS =================

DATA_MIN = 0.0
DATA_MAX = 60.0
PADDING = 0.5

SIGMA_SLOPE = 0.15
SIGMA_BASE = 1.4
SIGMA_POINTS = 400
TREND_POINTS = 500

PLOT_WIDTH = 980
PLOT_HEIGHT = 980

# Fonts
FONT_FAMILY = "Inter, Roboto, Helvetica Neue, Arial, sans-serif"
BASE_FONT_SIZE = 20
TITLE_FONT_SIZE = 26
AXIS_TITLE_FONT_SIZE = 32
AXIS_TICK_FONT_SIZE = 30
LEGEND_FONT_SIZE = 20

# Markers / lines
MARKER_SIZE = 15
MARKER_LINE_WIDTH = 1.5
LINE_WIDTH = 3
TREND_LINE_WIDTH = 3.5

# Colors
PLOT_BG_COLOR = "rgba(235,238,245,1)"
PAPER_BG_COLOR = "white"
LEGEND_BG_COLOR = "rgba(255,255,255,0.45)"
GRID_COLOR = "rgba(100,116,139,0.18)"

# Layout
MARGIN_LEFT = 100
MARGIN_RIGHT = 15
MARGIN_TOP = 0
MARGIN_BOTTOM = 100

LEGEND_X = 0.008
LEGEND_Y = 0.985

COLORS = ["#1E88E5", "#8E24AA", "#A0D137", "#EA1D1D", "#06EE0D", "#FB8C00"]
SYMBOLS = ["circle", "x", "square", "diamond", "triangle-up", "star"]


def _logistic_time(t, R0, Rmax, t50, k):
    """Logistic function over SYNTAX score."""
    t = np.asarray(t, dtype=float)
    t_safe = np.where(t <= 0, 1e-3, t)
    return R0 + (Rmax - R0) / (1.0 + (t50 / t_safe) ** k)


def _fit_logistic(x, y, domain, n=TREND_POINTS):
    """
    Fit a logistic curve.
    Returns X, Y or (None, None) if the fit fails.
    """
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)
    m = np.isfinite(x) & np.isfinite(y)
    if m.sum() < 4:
        return None, None

    x_m, y_m = x[m], y[m]
    x_min = max(float(np.min(x_m)), float(domain[0]))
    x_max = min(float(np.max(x_m)), float(domain[1]))
    if not np.isfinite(x_min) or not np.isfinite(x_max) or x_max <= x_min:
        return None, None

    x_pos = x_m[x_m > 0]
    if x_pos.size == 0:
        return None, None

    R0_init = float(np.percentile(y_m, 10))
    Rmax_init = float(np.percentile(y_m, 90))
    t50_init = float(np.median(x_pos))
    k_init = 1.0

    lower = [-10.0, 0.0, 1e-3, 0.01]
    upper = [60.0, 80.0, 60.0, 10.0]

    try:
        popt, _ = curve_fit(
            _logistic_time,
            x_m,
            y_m,
            p0=[R0_init, Rmax_init, t50_init, k_init],
            bounds=(lower, upper),
            maxfev=20000,
        )
    except Exception:
        return None, None

    X = np.linspace(x_min, x_max, n)
    Y = _logistic_time(X, *popt)
    return X, Y


def visualize_final_syntax_plotly_multi(
    datasets,
    r2_values,
    gt_row,
    postfix=None,
    threshold: float = 22.0,
    recall_values=None,
    backbone: bool = False,
    show_title: bool = False,
):
    """
    Unified SYNTAX visualization: points, risk zones and logistic trends.

    Parameters
    ----------
    datasets : dict[str, tuple[list[float], list[float]]]
        {dataset_name: (syntax_true_list, syntax_pred_list)}.
    r2_values : dict[str, float]
        Pearson correlation per dataset.
    gt_row : str
        String for the plot title (e.g. "ENSEMBLE" or "BOTH").
    postfix : str | None
        Suffix for the saved file name.
    threshold : float
        SYNTAX threshold (typically 22.0) to separate risk zones.
    recall_values : dict[str, float] | None
        Mean recall per dataset (may be None).
    backbone : bool
        If True, saves into `visualizations/backbone`, else into `visualizations/`.
    """
    fig = go.Figure()

    line_min = DATA_MIN - PADDING
    line_max = DATA_MAX + PADDING
    domain = (line_min, line_max)

    base_font = dict(
        family=FONT_FAMILY,
        size=BASE_FONT_SIZE,
    )

    # ---------- Risk zones and bands (legendrank=0) ----------
    fig.add_trace(
        go.Scatter(
            x=[line_min, threshold, threshold, line_min],
            y=[line_min, line_min, threshold, threshold],
            fill="toself",
            fillcolor="rgba(255, 82, 82, 0.12)",
            line=dict(color="rgba(0,0,0,0)"),
            name="Low-risk zone",
            legendgroup="zones",
            legendgrouptitle_text="Thresholds & lines",
            showlegend=True,
            hoverinfo="skip",
            legendrank=0,
        )
    )
    fig.add_trace(
        go.Scatter(
            x=[threshold, line_max, line_max, threshold],
            y=[threshold, threshold, line_max, line_max],
            fill="toself",
            fillcolor="rgba(76, 175, 80, 0.14)",
            line=dict(color="rgba(0,0,0,0)"),
            name="High-risk zone",
            legendgroup="zones",
            showlegend=True,
            hoverinfo="skip",
            legendrank=0,
        )
    )

    fig.add_trace(
        go.Scatter(
            x=[threshold, threshold, None, line_min, line_max],
            y=[line_min, line_max, None, threshold, threshold],
            mode="lines",
            name=f"SYNTAX = {threshold}",
            legendgroup="zones",
            showlegend=True,
            line=dict(color="rgba(46,125,50,0.85)", width=LINE_WIDTH, dash="dash"),
            legendrank=0,
            hoverinfo="skip",
        )
    )

    x_vals = np.linspace(line_min, line_max, SIGMA_POINTS)
    sigma_upper = x_vals + SIGMA_BASE + SIGMA_SLOPE * x_vals
    sigma_lower = x_vals - SIGMA_BASE - SIGMA_SLOPE * x_vals
    two_sigma_upper = x_vals + 2 * SIGMA_BASE + 2 * SIGMA_SLOPE * x_vals
    two_sigma_lower = x_vals - 2 * SIGMA_BASE - 2 * SIGMA_SLOPE * x_vals

    fig.add_trace(
        go.Scatter(
            x=np.concatenate([x_vals, x_vals[::-1]]),
            y=np.concatenate([two_sigma_lower, two_sigma_upper[::-1]]),
            fill="toself",
            fillcolor="rgba(255,193,7,0.18)",
            line=dict(color="rgba(0,0,0,0)"),
            name="± 2σ",
            legendgroup="zones",
            showlegend=True,
            hoverinfo="skip",
            legendrank=0,
        )
    )
    fig.add_trace(
        go.Scatter(
            x=np.concatenate([x_vals, x_vals[::-1]]),
            y=np.concatenate([sigma_lower, sigma_upper[::-1]]),
            fill="toself",
            fillcolor="rgba(255,152,0,0.30)",
            line=dict(color="rgba(0,0,0,0)"),
            name="± σ",
            legendgroup="zones",
            showlegend=True,
            hoverinfo="skip",
            legendrank=0,
        )
    )

    fig.add_trace(
        go.Scatter(
            x=[line_min, line_max],
            y=[line_min, line_max],
            mode="lines",
            name="Perfect prediction",
            legendgroup="zones",
            showlegend=True,
            line=dict(color="rgba(30,30,30,0.85)", width=LINE_WIDTH),
            legendrank=0,
        )
    )

    # ---------- Datasets (legendrank=20) ----------
    first_dataset = True
    for i, (label, (syntax_true, syntax_pred)) in enumerate(datasets.items()):
        x = np.array(syntax_true, dtype=float)
        y = np.array(syntax_pred, dtype=float)
        if x.size == 0 or y.size == 0:
            continue

        pearson = r2_values.get(label, None)
        recall = recall_values.get(label, None) if recall_values else None
        hover_lines = [f"<b>{label}</b>"]
        if pearson is not None:
            hover_lines.append(f"Pearson = {pearson:.3f}")
        if recall is not None:
            hover_lines.append(f"Mean recall = {recall:.3f}")
        hovertemplate = (
            "<br>".join(hover_lines)
            + "<br>Ground truth: %{x:.3f}<br>Prediction: %{y:.3f}<extra></extra>"
        )

        fig.add_trace(
            go.Scatter(
                x=x,
                y=y,
                mode="markers",
                name=label,
                legendgroup="datasets",
                legendgrouptitle_text=("Datasets" if first_dataset else None),
                showlegend=True,
                marker=dict(
                    color=COLORS[i % len(COLORS)],
                    size=MARKER_SIZE,
                    opacity=0.96,
                    symbol=SYMBOLS[i % len(SYMBOLS)],
                    line=dict(
                        width=MARKER_LINE_WIDTH,
                        color="rgba(255,255,255,0.95)",
                    ),
                ),
                hovertemplate=hovertemplate,
                legendrank=20,
            )
        )
        first_dataset = False

    # ---------- Logistic trends (legendrank=30) ----------
    first_trend = True
    for i, (label, (syntax_true, syntax_pred)) in enumerate(datasets.items()):
        x = np.array(syntax_true, dtype=float)
        y = np.array(syntax_pred, dtype=float)
        if x.size == 0 or y.size == 0:
            continue

        Xc, Yc = _fit_logistic(x, y, domain=domain)
        if Xc is not None:
            fig.add_trace(
                go.Scatter(
                    x=Xc,
                    y=Yc,
                    mode="lines",
                    name=label,
                    legendgroup="trends",
                    legendgrouptitle_text=("Logistic trends" if first_trend else None),
                    showlegend=True,
                    line=dict(
                        color=COLORS[i % len(COLORS)],
                        width=TREND_LINE_WIDTH,
                    ),
                    hoverinfo="skip",
                    legendrank=30,
                )
            )
            first_trend = False

    # ---------- Layout ----------
    title_text = f"SYNTAX predictions ({gt_row})"
    if postfix:
        title_text += f" {postfix}"

    layout_kwargs = dict(
        font=dict(
            family=FONT_FAMILY,
            size=BASE_FONT_SIZE,
        ),
        width=PLOT_WIDTH,
        height=PLOT_HEIGHT,
        plot_bgcolor=PLOT_BG_COLOR,
        paper_bgcolor=PAPER_BG_COLOR,
        legend=dict(
            x=LEGEND_X,
            y=LEGEND_Y,
            bgcolor=LEGEND_BG_COLOR,
            bordercolor="rgba(203,213,225,0.7)",
            borderwidth=1,
            font=dict(size=LEGEND_FONT_SIZE, family=FONT_FAMILY),
            tracegroupgap=8,
            itemclick="toggle",
            itemdoubleclick="toggleothers",
            groupclick="toggleitem",
        ),
        xaxis=dict(
            title=dict(
                text="SYNTAX ground truth",
                font=dict(
                    size=AXIS_TITLE_FONT_SIZE,
                    family=FONT_FAMILY,
                    color="rgba(15,23,42,1)",
                ),
            ),
            showgrid=True,
            gridcolor=GRID_COLOR,
            gridwidth=1,
            zeroline=False,
            tickfont=dict(
                size=AXIS_TICK_FONT_SIZE,
                family=FONT_FAMILY,
            ),
            range=[line_min, line_max],
            constrain="domain",
        ),
        yaxis=dict(
            title=dict(
                text="SYNTAX predictions",
                font=dict(
                    size=AXIS_TITLE_FONT_SIZE,
                    family=FONT_FAMILY,
                    color="rgba(15,23,42,1)",
                ),
            ),
            showgrid=True,
            gridcolor=GRID_COLOR,
            gridwidth=1,
            zeroline=False,
            tickfont=dict(
                size=AXIS_TICK_FONT_SIZE,
                family=FONT_FAMILY,
            ),
            range=[line_min, line_max],
            scaleanchor="x",
            scaleratio=1,
            constrain="domain",
        ),
        margin=dict(
            l=MARGIN_LEFT,
            r=MARGIN_RIGHT,
            t=MARGIN_TOP,
            b=MARGIN_BOTTOM,
        ),
    )

    if show_title:
        layout_kwargs["title"] = dict(
            text=title_text,
            x=0.5,
            xanchor="center",
            font=dict(
                size=TITLE_FONT_SIZE,
                family=FONT_FAMILY,
                color="rgba(15,23,42,1)",
            ),
        )

    fig.update_layout(**layout_kwargs)

    # ---------- Saving ----------
    save_dir = "visualizations"
    if backbone:
        save_dir = os.path.join(save_dir, "backbone")
    os.makedirs(save_dir, exist_ok=True)

    postfix_html = f"{postfix}" if postfix else "syntax"
    save_path_html = os.path.join(save_dir, f"{postfix_html}.html")
    fig.write_html(save_path_html, include_mathjax="cdn")
    print(f"Saved visualization with logistic trends: {save_path_html}")