"""Aquarium 2D visualization helpers.

Builds fish records (with full genome attributes + ancestor lineage) and
renders a self-contained HTML/JS canvas animation, embeddable in Streamlit
via ``streamlit.components.v1.html``. Includes a click handler that opens
an info panel showing genome details and BFS ancestor levels.
"""

from __future__ import annotations

import json
import math
from typing import Any

import pandas as pd  # type: ignore[import-untyped]

# Color palette per cognitive style. Default fallback for unknown styles is grey.
STYLE_COLORS: dict[str, str] = {
    "physicist": "#4cc9f0",
    "biologist": "#52b788",
    "historian": "#e76f51",
    "meteorologist": "#ffd166",
    "ecologist": "#a78bfa",
    "engineer": "#fb6f92",
}
DEFAULT_COLOR: str = "#94a3b8"


def _is_nan(v: Any) -> bool:
    try:
        return bool(pd.isna(v))
    except (TypeError, ValueError):
        return False


def _safe_float(v: Any, default: float = 0.0) -> float:
    if v is None or _is_nan(v):
        return default
    try:
        return float(v)
    except (TypeError, ValueError):
        return default


def _safe_int(v: Any, default: int = 0) -> int:
    if v is None or _is_nan(v):
        return default
    try:
        return int(v)
    except (TypeError, ValueError):
        return default


def _safe_str(v: Any, default: str = "") -> str:
    if v is None or _is_nan(v):
        return default
    return str(v)


def _safe_list(v: Any) -> list[Any]:
    if v is None:
        return []
    if isinstance(v, list):
        return list(v)
    # pandas may store python lists in object cells; if it's e.g. a numpy array,
    # falling back to list() is fine. NaN scalar is excluded by _is_nan.
    if _is_nan(v):
        return []
    try:
        return list(v)
    except TypeError:
        return []


def build_lineage_index(
    genomes_df: pd.DataFrame, evals_df: pd.DataFrame
) -> dict[str, dict[str, Any]]:
    """Build ``{genome_id: attrs}`` for every genome in the run.

    ``genomes_df`` must come from ``genomes_df(repo, run_id)`` (no gen filter):
    columns include ``id``, ``generation_idx``, ``system_prompt``,
    ``feature_access``, ``temperature``, ``top_p``, ``model_tier``,
    ``lookback_window``, ``cognitive_style``, ``parent_ids``, ``generation``.

    ``evals_df`` must come from ``evaluations_df(repo, run_id)``: columns
    include ``genome_id``, ``fitness``, ``dsr``, ``sharpe``, ``max_dd``,
    ``n_trades``.
    """
    if genomes_df.empty:
        return {}

    if evals_df is None or evals_df.empty:
        merged = genomes_df.copy()
        for col in ("fitness", "dsr", "sharpe", "max_dd", "n_trades"):
            if col not in merged.columns:
                merged[col] = 0.0 if col != "n_trades" else 0
    else:
        merged = genomes_df.merge(
            evals_df,
            left_on="id",
            right_on="genome_id",
            how="left",
            suffixes=("", "_eval"),
        )

    index: dict[str, dict[str, Any]] = {}
    for _, row in merged.iterrows():
        gid = _safe_str(row.get("id"))
        if not gid:
            continue
        # ``generation`` is the genome's evolutionary generation (from payload).
        # If absent, fall back to ``generation_idx`` (column added by the
        # repository). Defensive: both may be missing in edge cases.
        gen_val: Any = row.get("generation")
        if gen_val is None or _is_nan(gen_val):
            gen_val = row.get("generation_idx", 0)
        index[gid] = {
            "id": gid,
            "generation": _safe_int(gen_val, 0),
            "fitness": _safe_float(row.get("fitness"), 0.0),
            "dsr": _safe_float(row.get("dsr"), 0.0),
            "sharpe": _safe_float(row.get("sharpe"), 0.0),
            "max_dd": _safe_float(row.get("max_dd"), 0.0),
            "n_trades": _safe_int(row.get("n_trades"), 0),
            "cognitive_style": _safe_str(row.get("cognitive_style"), ""),
            "system_prompt": _safe_str(row.get("system_prompt"), ""),
            "temperature": _safe_float(row.get("temperature"), 0.0),
            "lookback_window": _safe_int(row.get("lookback_window"), 0),
            "feature_access": _safe_list(row.get("feature_access")),
            "model_tier": _safe_str(row.get("model_tier"), ""),
            "parent_ids": _safe_list(row.get("parent_ids")),
        }
    return index


def trace_ancestors(
    genome_id: str,
    lineage_index: dict[str, dict[str, Any]],
    max_levels: int = 5,
) -> list[list[dict[str, Any]]]:
    """BFS over ``parent_ids`` returning levels of ancestors.

    ``levels[0]`` = direct parents, ``levels[1]`` = grandparents, etc. Each
    entry is a small dict (no ``system_prompt``, to keep JSON payload light):
    ``{id, generation, fitness, cognitive_style}``. Cycles are guarded via a
    ``seen`` set; missing parents (not in this run) are stubbed with sentinel
    values so the lineage display still renders the relationship.
    """
    levels: list[list[dict[str, Any]]] = []
    root = lineage_index.get(genome_id, {})
    current_ids: list[str] = list(root.get("parent_ids", []))
    seen: set[str] = {genome_id}
    for _ in range(max_levels):
        if not current_ids:
            break
        level_entries: list[dict[str, Any]] = []
        next_ids: list[str] = []
        for pid in current_ids:
            if pid in seen:
                continue
            seen.add(pid)
            entry = lineage_index.get(pid)
            if entry is None:
                level_entries.append(
                    {
                        "id": pid,
                        "generation": -1,
                        "fitness": 0.0,
                        "cognitive_style": "",
                    }
                )
                continue
            level_entries.append(
                {
                    "id": entry["id"],
                    "generation": entry["generation"],
                    "fitness": entry["fitness"],
                    "cognitive_style": entry["cognitive_style"],
                }
            )
            next_ids.extend(entry.get("parent_ids", []))
        if not level_entries:
            break
        levels.append(level_entries)
        current_ids = next_ids
    return levels


def build_fish_dataset(
    active_df: pd.DataFrame,
    lineage_index: dict[str, dict[str, Any]] | None = None,
    max_lineage_levels: int = 5,
) -> list[dict[str, Any]]:
    """Build full fish records for each active genome.

    For every row in ``active_df`` the matching entry in ``lineage_index`` is
    looked up by ``genome_id`` (or ``id``) and attached together with the BFS
    ancestor levels. Rows whose id is not in the index are skipped.

    Backward-compat: if ``lineage_index`` is ``None`` (legacy call site, e.g.
    test fixtures with simple merged DataFrames) we synthesize a minimal
    lineage from ``active_df`` itself so the function still returns useful
    fish records.
    """
    if active_df.empty:
        return []

    if lineage_index is None:
        # Legacy path: build a tiny index from the active df only.
        synth: dict[str, dict[str, Any]] = {}
        for _, row in active_df.iterrows():
            gid = _safe_str(row.get("genome_id") or row.get("id"))
            if not gid:
                continue
            fitness_val = _safe_float(row.get("fitness"), float("nan"))
            if math.isnan(fitness_val):
                continue
            synth[gid] = {
                "id": gid,
                "generation": _safe_int(row.get("generation"), 0),
                "fitness": fitness_val,
                "dsr": _safe_float(row.get("dsr"), 0.0),
                "sharpe": _safe_float(row.get("sharpe"), 0.0),
                "max_dd": _safe_float(row.get("max_dd"), 0.0),
                "n_trades": _safe_int(row.get("n_trades"), 0),
                "cognitive_style": _safe_str(row.get("cognitive_style"), "unknown"),
                "system_prompt": _safe_str(row.get("system_prompt"), ""),
                "temperature": _safe_float(row.get("temperature"), 0.0),
                "lookback_window": _safe_int(row.get("lookback_window"), 0),
                "feature_access": _safe_list(row.get("feature_access")),
                "model_tier": _safe_str(row.get("model_tier"), ""),
                "parent_ids": _safe_list(row.get("parent_ids")),
            }
        lineage_index = synth

    fish: list[dict[str, Any]] = []
    for _, row in active_df.iterrows():
        gid = _safe_str(row.get("genome_id") or row.get("id"))
        if not gid:
            continue
        attrs = lineage_index.get(gid)
        if attrs is None:
            continue
        if math.isnan(attrs.get("fitness", 0.0)):
            continue
        ancestors = trace_ancestors(gid, lineage_index, max_lineage_levels)
        record = {**attrs, "ancestors": ancestors}
        fish.append(record)
    return fish


def build_aquarium_html(
    fish: list[dict[str, Any]],
    canvas_w: int = 1000,
    canvas_h: int = 600,
) -> str:
    """Build the self-contained HTML/JS string for the aquarium canvas.

    The output embeds a click handler: tapping a fish opens an info panel
    (top-right of the canvas) showing its genome attributes and BFS ancestor
    levels. Labels are no longer rendered on the canvas itself.
    """
    fish_json = json.dumps(fish)
    palette_json = json.dumps(STYLE_COLORS)
    default_color = DEFAULT_COLOR

    # All braces inside <style>/<script> are escaped to literals using {{ }}
    # so we can use Python f-string substitution for the few JSON payloads.
    return f"""
<div style="position:relative;width:100%;height:{canvas_h}px;">
  <canvas id="aquarium" width="{canvas_w}" height="{canvas_h}"
          style="width:100%;height:{canvas_h}px;border-radius:12px;
                 background:linear-gradient(180deg,#0a2540 0%,#1a4d80 100%);
                 display:block;cursor:pointer;"></canvas>
  <div id="fish-info-panel" style="
    position:absolute;
    top:12px;
    right:12px;
    width:340px;
    max-height:580px;
    overflow-y:auto;
    background:rgba(8,16,32,0.92);
    color:#e2e8f0;
    border-radius:10px;
    padding:14px 16px;
    font-family:system-ui,-apple-system,sans-serif;
    font-size:12px;
    line-height:1.5;
    border:1px solid rgba(255,255,255,0.1);
    backdrop-filter:blur(6px);
    -webkit-backdrop-filter:blur(6px);
    display:none;
    z-index:10;
  ">
    <div id="fish-info-content"></div>
    <button id="fish-info-close" style="
      position:absolute;top:8px;right:10px;
      background:transparent;color:#94a3b8;border:none;
      cursor:pointer;font-size:16px;
    ">&times;</button>
  </div>
</div>
<script>
(function() {{
  const FISH_DATA = {fish_json};
  const STYLE_COLORS = {palette_json};
  const DEFAULT_COLOR = {json.dumps(default_color)};

  const canvas = document.getElementById('aquarium');
  if (!canvas) return;
  const ctx = canvas.getContext('2d');
  const W = canvas.width;
  const H = canvas.height;

  const panel = document.getElementById('fish-info-panel');
  const panelContent = document.getElementById('fish-info-content');
  const closeBtn = document.getElementById('fish-info-close');
  if (closeBtn) {{
    closeBtn.addEventListener('click', function() {{
      panel.style.display = 'none';
    }});
  }}

  // Normalize fitness for sizing.
  let maxFit = 0.0;
  for (const f of FISH_DATA) {{
    if (typeof f.fitness === 'number' && f.fitness > maxFit) maxFit = f.fitness;
  }}

  function lerp(a, b, t) {{ return a + (b - a) * t; }}

  function radiusFor(fitness) {{
    if (maxFit <= 0) return 8;
    const t = Math.max(0.05, Math.min(1.0, fitness / maxFit));
    return lerp(8, 40, t);
  }}

  function colorFor(style) {{
    return STYLE_COLORS[style] || DEFAULT_COLOR;
  }}

  // Init fish state. Each entry keeps a reference to the original data dict
  // so the click handler can show full attributes + ancestors.
  const fishState = FISH_DATA.map(function(f, idx) {{
    const r = radiusFor(f.fitness);
    return {{
      data: f,
      color: colorFor(f.cognitive_style),
      radius: r,
      x: Math.random() * (W - 2 * r) + r,
      y: Math.random() * (H - 2 * r) + r,
      vx: (Math.random() - 0.5) * 1.5,
      vy: (Math.random() - 0.5) * 1.0,
      rank: idx,
    }};
  }});

  // Bubbles for ambience.
  const N_BUBBLES = 25;
  const bubbles = Array.from({{length: N_BUBBLES}}, function() {{ return {{
    x: Math.random() * W,
    y: Math.random() * H,
    r: 1 + Math.random() * 3,
    vy: 0.3 + Math.random() * 0.7,
  }}; }});

  function drawBubble(b) {{
    ctx.beginPath();
    ctx.arc(b.x, b.y, b.r, 0, Math.PI * 2);
    ctx.fillStyle = 'rgba(255,255,255,0.18)';
    ctx.fill();
  }}

  function updateBubble(b) {{
    b.y -= b.vy;
    if (b.y < -10) {{
      b.y = H + 5;
      b.x = Math.random() * W;
    }}
  }}

  function drawFish(f) {{
    const facingLeft = f.vx < 0;
    ctx.save();
    ctx.translate(f.x, f.y);
    if (facingLeft) ctx.scale(-1, 1);

    // Halo for top-3 fish.
    if (f.rank < 3) {{
      const grad = ctx.createRadialGradient(0, 0, f.radius * 0.5, 0, 0, f.radius * 2.0);
      grad.addColorStop(0, f.color + 'aa');
      grad.addColorStop(1, f.color + '00');
      ctx.fillStyle = grad;
      ctx.beginPath();
      ctx.arc(0, 0, f.radius * 2.0, 0, Math.PI * 2);
      ctx.fill();
    }}

    // Body (ellipse).
    ctx.fillStyle = f.color;
    ctx.beginPath();
    ctx.ellipse(0, 0, f.radius, f.radius * 0.6, 0, 0, Math.PI * 2);
    ctx.fill();

    // Tail.
    ctx.beginPath();
    ctx.moveTo(-f.radius, 0);
    ctx.lineTo(-f.radius * 1.6, -f.radius * 0.5);
    ctx.lineTo(-f.radius * 1.6, f.radius * 0.5);
    ctx.closePath();
    ctx.fill();

    // Eye.
    ctx.fillStyle = '#ffffff';
    ctx.beginPath();
    ctx.arc(f.radius * 0.45, -f.radius * 0.15, Math.max(1.5, f.radius * 0.12), 0, Math.PI * 2);
    ctx.fill();
    ctx.fillStyle = '#1a1a1a';
    ctx.beginPath();
    ctx.arc(f.radius * 0.50, -f.radius * 0.15, Math.max(0.8, f.radius * 0.06), 0, Math.PI * 2);
    ctx.fill();

    ctx.restore();
  }}

  function updateFish(f) {{
    f.vx += (Math.random() - 0.5) * 0.05;
    f.vy += (Math.random() - 0.5) * 0.05;

    const speed = Math.hypot(f.vx, f.vy);
    const maxSpeed = 1.5;
    if (speed > maxSpeed) {{
      f.vx = (f.vx / speed) * maxSpeed;
      f.vy = (f.vy / speed) * maxSpeed;
    }}

    f.x += f.vx;
    f.y += f.vy;

    if (f.x < f.radius) {{ f.x = f.radius; f.vx = -f.vx; }}
    if (f.x > W - f.radius) {{ f.x = W - f.radius; f.vx = -f.vx; }}
    if (f.y < f.radius) {{ f.y = f.radius; f.vy = -f.vy; }}
    if (f.y > H - f.radius) {{ f.y = H - f.radius; f.vy = -f.vy; }}
  }}

  function frame() {{
    ctx.clearRect(0, 0, W, H);

    ctx.strokeStyle = 'rgba(255,255,255,0.04)';
    ctx.lineWidth = 1;
    for (let i = 0; i < 6; i++) {{
      const y = (H / 6) * i + (Date.now() / 50 % (H / 6));
      ctx.beginPath();
      ctx.moveTo(0, y);
      ctx.lineTo(W, y);
      ctx.stroke();
    }}

    for (const b of bubbles) {{
      updateBubble(b);
      drawBubble(b);
    }}
    for (const f of fishState) {{
      updateFish(f);
      drawFish(f);
    }}
    requestAnimationFrame(frame);
  }}

  // CLICK HANDLER: hit-test in canvas pixel space (account for CSS scaling).
  canvas.addEventListener('click', function(e) {{
    const rect = canvas.getBoundingClientRect();
    const scaleX = canvas.width / rect.width;
    const scaleY = canvas.height / rect.height;
    const cx = (e.clientX - rect.left) * scaleX;
    const cy = (e.clientY - rect.top) * scaleY;
    let best = null;
    let bestDist = Infinity;
    for (const f of fishState) {{
      const dx = cx - f.x;
      const dy = cy - f.y;
      const d = Math.sqrt(dx * dx + dy * dy);
      const hit = Math.max(f.radius + 6, 14);
      if (d < hit && d < bestDist) {{
        best = f;
        bestDist = d;
      }}
    }}
    if (best) showFishInfo(best);
  }});

  const ROW_STYLE = 'display:flex;justify-content:space-between;'
                  + 'padding:2px 0;border-bottom:1px solid rgba(255,255,255,0.05);';
  const PROMPT_STYLE = 'margin-top:10px;padding:8px;'
                     + 'background:rgba(255,255,255,0.04);border-radius:6px;'
                     + 'font-size:11px;font-style:italic;color:#cbd5e1;';
  const ANC_HEAD_STYLE = 'margin:14px 0 6px 0;color:#94a3b8;'
                       + 'text-transform:uppercase;font-size:10px;letter-spacing:1px;';
  const ANC_ROW_STYLE = 'display:flex;align-items:center;padding:4px 6px;'
                      + 'margin-bottom:2px;background:rgba(255,255,255,0.03);'
                      + 'border-radius:4px;border-left:3px solid ';
  const NO_ANC_STYLE = 'margin-top:10px;font-size:10px;color:#64748b;';
  const DASH = '\\u2014';

  function metricRow(label, value) {{
    return '<div style="' + ROW_STYLE + '">'
         + '<span style="color:#94a3b8;">' + label + '</span>'
         + '<span style="color:#e2e8f0;">' + value + '</span></div>';
  }}

  function escapeHtml(s) {{
    const div = document.createElement('div');
    div.appendChild(document.createTextNode(String(s)));
    return div.innerHTML;
  }}

  function fmt(v, dp) {{
    if (v === null || v === undefined || typeof v !== 'number' || isNaN(v)) {{
      return DASH;
    }}
    return v.toFixed(dp);
  }}

  function showFishInfo(fish) {{
    const data = fish.data;
    const styleColor = STYLE_COLORS[data.cognitive_style] || DEFAULT_COLOR;
    let html = '';
    const idShort = String(data.id || '').slice(0, 8);
    html += '<h4 style="margin:0 0 10px 0;color:' + styleColor + ';">';
    html += escapeHtml(idShort)
         + ' <span style="color:#94a3b8;font-weight:normal;font-size:11px;">'
         + 'gen ' + escapeHtml(data.generation) + '</span>';
    html += '</h4>';
    html += metricRow('fitness', fmt(data.fitness, 3));
    html += metricRow('DSR', fmt(data.dsr, 3));
    html += metricRow('Sharpe', fmt(data.sharpe, 3));
    html += metricRow('max DD', fmt(data.max_dd, 3));
    const trades = data.n_trades == null ? 0 : data.n_trades;
    html += metricRow('trades', escapeHtml(trades));
    html += metricRow('style', escapeHtml(data.cognitive_style || DASH));
    html += metricRow('tier', escapeHtml(data.model_tier || DASH));
    html += metricRow('temp', fmt(data.temperature, 2));
    const lookback = data.lookback_window == null ? DASH : data.lookback_window;
    html += metricRow('lookback', escapeHtml(lookback));
    const feats = (data.feature_access || []).join(', ');
    html += metricRow('features', escapeHtml(feats || DASH));
    if (data.system_prompt) {{
      html += '<div style="' + PROMPT_STYLE + '">'
           + escapeHtml(data.system_prompt) + '</div>';
    }}
    if (data.ancestors && data.ancestors.length > 0) {{
      html += '<h5 style="' + ANC_HEAD_STYLE + '">Discendenza</h5>';
      data.ancestors.forEach(function(level, idx) {{
        html += '<div style="margin-bottom:8px;">';
        html += '<div style="font-size:10px;color:#64748b;margin-bottom:4px;">'
             + 'Gen N\\u2212' + (idx + 1) + '</div>';
        level.forEach(function(ancestor) {{
          const c = STYLE_COLORS[ancestor.cognitive_style] || DEFAULT_COLOR;
          const aShort = String(ancestor.id || '').slice(0, 8);
          html += '<div style="' + ANC_ROW_STYLE + c + ';">';
          html += '<code style="color:' + c + ';font-size:10px;">'
               + escapeHtml(aShort) + '</code>';
          const af = ancestor.fitness;
          const fitTxt = (typeof af === 'number' && !isNaN(af))
                       ? af.toFixed(2) : DASH;
          html += '<span style="margin-left:auto;color:#94a3b8;font-size:10px;">'
               + 'fit ' + fitTxt + '</span>';
          html += '</div>';
        }});
        html += '</div>';
      }});
    }} else {{
      html += '<div style="' + NO_ANC_STYLE + '">'
           + 'Genoma di prima generazione (no ancestors)</div>';
    }}
    panelContent.innerHTML = html;
    panel.style.display = 'block';
  }}

  if (fishState.length === 0) {{
    ctx.fillStyle = 'rgba(255,255,255,0.7)';
    ctx.font = '16px sans-serif';
    ctx.textAlign = 'center';
    ctx.fillText('Acquario vuoto: nessun genoma da mostrare.', W / 2, H / 2);
  }} else {{
    requestAnimationFrame(frame);
  }}
}})();
</script>
"""