{% extends "base.html" %}

{% block title %}Contextual Bandits — Fractal Memory{% endblock %}

{% block content %}
<div class="page-header" style="display:flex;align-items:flex-start;justify-content:space-between;">
  <div>
    <h1>Contextual Bandits</h1>
    <p class="subtitle">Thompson Sampling — arm selection, Beta posteriors, and context bins</p>
  </div>
  <button class="btn-secondary" onclick="openHelp('bandits')" title="Help">? Help</button>
</div>

<div class="bandits-layout">

  <!-- ── Current Arm (full width) ───────────────────────────────── -->
  <div class="chart-card full-width" id="current-arm-card">
    <h2>Current Arm</h2>
    <div style="display:flex;align-items:flex-start;gap:2rem;flex-wrap:wrap;">
      <div>
        <div id="current-arm-badge" class="arm-badge arm-standard">—</div>
        <div style="margin-top:0.5rem;font-size:0.8rem;color:var(--text-dim);">
          Context key: <span id="current-context-key" style="color:var(--text);font-family:monospace;">—</span>
        </div>
      </div>
      <div id="arm-params" style="display:flex;flex-wrap:wrap;gap:0.5rem;align-items:flex-start;"></div>
    </div>
  </div>

  <!-- ── Arm Posterior Distribution (full width) ────────────────── -->
  <div class="chart-card full-width">
    <h2>Arm Posterior Distribution
      <span style="font-size:0.75rem;color:var(--text-dim);font-weight:400;margin-left:0.75rem;text-transform:none;letter-spacing:0;">Alpha=successes (blue), Beta=failures (red). Line=mean success rate.</span>
    </h2>
    <div style="position:relative;height:260px;">
      <canvas id="posterior-chart"></canvas>
    </div>
  </div>

  <!-- ── Arm Stats Table (full width) ───────────────────────────── -->
  <div class="chart-card full-width">
    <h2>Arm Statistics</h2>
    <div style="overflow-x:auto;">
      <table class="data-table" id="arm-stats-table">
        <thead>
          <tr>
            <th>Arm</th>
            <th>Alpha</th>
            <th>Beta</th>
            <th>Mean Success Rate</th>
            <th>Pull Count</th>
            <th>Win Rate</th>
          </tr>
        </thead>
        <tbody id="arm-stats-tbody">
          <tr><td colspan="6" style="color:var(--text-dim);text-align:center;padding:2rem;">Loading arm stats…</td></tr>
        </tbody>
      </table>
    </div>
  </div>

  <!-- ── Selection History (full width) ────────────────────────── -->
  <div class="chart-card full-width">
    <h2>Arm Selection History
      <span style="font-size:0.75rem;color:var(--text-dim);font-weight:400;margin-left:0.75rem;text-transform:none;letter-spacing:0;">Last 20 sessions — stacked area by arm.</span>
    </h2>
    <div style="position:relative;height:240px;">
      <canvas id="history-chart"></canvas>
    </div>
  </div>

  <!-- ── Context Bin Distribution (full width) ─────────────────── -->
  <div class="chart-card full-width">
    <h2>Context Bin Distribution</h2>
    <div style="position:relative;height:220px;">
      <canvas id="context-chart"></canvas>
    </div>
    <p style="font-size:0.75rem;color:var(--text-dim);margin-top:0.75rem;">
      Context key = richness_bin-pressure_bin-density_bin (0=low, 1=med, 2=high)
    </p>
  </div>

</div><!-- /.bandits-layout -->
{% endblock %}


{% block scripts %}
<script>
/* ================================================================
   Contextual Bandits — Thompson Sampling dashboard
   ================================================================ */

// ── Utility ───────────────────────────────────────────────────────
function escHtml(str) {
  return String(str)
    .replace(/&/g, '&amp;')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;');
}

// ── Chart.js defaults ─────────────────────────────────────────────
Chart.defaults.color = '#6b7280';
Chart.defaults.borderColor = '#1e293b';
Chart.defaults.font.family = '"Inter", "Segoe UI", system-ui, sans-serif';

const ARM_COLORS = ['#4a9eff', '#a855f7', '#34d399', '#ff6b35', '#f5c518', '#ff2d55', '#22d3ee'];

// ── Charts ────────────────────────────────────────────────────────
let posteriorChart = null;
let historyChart   = null;
let contextChart   = null;

function initPosteriorChart() {
  const ctx = document.getElementById('posterior-chart');
  if (!ctx) return;
  posteriorChart = new Chart(ctx, {
    type: 'bar',
    data: { labels: [], datasets: [] },
    options: {
      responsive: true,
      maintainAspectRatio: false,
      animation: { duration: 300 },
      scales: {
        x: { grid: { color: '#1e293b' }, ticks: { color: '#6b7280' } },
        y: {
          id: 'y',
          position: 'left',
          title: { display: true, text: 'Alpha / Beta', color: '#6b7280', font: { size: 11 } },
          grid: { color: '#1e293b' }, ticks: { color: '#6b7280' },
        },
        y2: {
          id: 'y2',
          position: 'right',
          min: 0, max: 1,
          title: { display: true, text: 'Mean Success Rate', color: '#6b7280', font: { size: 11 } },
          grid: { drawOnChartArea: false }, ticks: { color: '#6b7280', stepSize: 0.2 },
        },
      },
      plugins: {
        legend: { position: 'top', labels: { color: '#c8d6e5', padding: 14, font: { size: 11 } } },
        tooltip: { backgroundColor: '#111118', borderColor: '#1e293b', borderWidth: 1, titleColor: '#c8d6e5', bodyColor: '#6b7280' },
      },
    },
  });
}

function initHistoryChart() {
  const ctx = document.getElementById('history-chart');
  if (!ctx) return;
  historyChart = new Chart(ctx, {
    type: 'line',
    data: { labels: [], datasets: [] },
    options: {
      responsive: true,
      maintainAspectRatio: false,
      animation: { duration: 300 },
      scales: {
        x: { title: { display: true, text: 'Session', color: '#6b7280', font: { size: 11 } }, grid: { color: '#1e293b' }, ticks: { color: '#6b7280', maxTicksLimit: 10 } },
        y: { title: { display: true, text: 'Selections', color: '#6b7280', font: { size: 11 } }, stacked: true, grid: { color: '#1e293b' }, ticks: { color: '#6b7280' } },
      },
      plugins: {
        legend: { position: 'top', labels: { color: '#c8d6e5', padding: 14, font: { size: 11 } } },
        tooltip: { backgroundColor: '#111118', borderColor: '#1e293b', borderWidth: 1, titleColor: '#c8d6e5', bodyColor: '#6b7280' },
        filler: {},
      },
    },
  });
}

function initContextChart() {
  const ctx = document.getElementById('context-chart');
  if (!ctx) return;
  contextChart = new Chart(ctx, {
    type: 'bar',
    data: { labels: [], datasets: [] },
    options: {
      responsive: true,
      maintainAspectRatio: false,
      animation: { duration: 300 },
      scales: {
        x: { grid: { color: '#1e293b' }, ticks: { color: '#6b7280', maxRotation: 45 } },
        y: { title: { display: true, text: 'Times Selected', color: '#6b7280', font: { size: 11 } }, grid: { color: '#1e293b' }, ticks: { color: '#6b7280' } },
      },
      plugins: {
        legend: { position: 'top', labels: { color: '#c8d6e5', padding: 12, font: { size: 11 } } },
        tooltip: { backgroundColor: '#111118', borderColor: '#1e293b', borderWidth: 1, titleColor: '#c8d6e5', bodyColor: '#6b7280' },
      },
    },
  });
}

// ── Arm name -> badge colour class ───────────────────────────────
function armClass(name) {
  const n = String(name).toLowerCase();
  if (n === 'aggressive') return 'arm-aggressive';
  if (n === 'conservative') return 'arm-conservative';
  if (n === 'explorer') return 'arm-explorer';
  if (n === 'precision') return 'arm-precision';
  return 'arm-standard';
}

// ── Win-rate progress bar ─────────────────────────────────────────
function winRateBar(rate) {
  const pct = Math.max(0, Math.min(100, rate * 100));
  const color = pct >= 60 ? 'var(--green)' : pct >= 40 ? '#f5c518' : 'var(--red)';
  return `<div style="display:flex;align-items:center;gap:0.4rem;">
    <div style="flex:1;height:6px;background:var(--border);border-radius:3px;overflow:hidden;min-width:60px;">
      <div style="width:${pct.toFixed(1)}%;height:100%;background:${color};border-radius:3px;transition:width .4s;"></div>
    </div>
    <span style="font-size:0.78rem;font-weight:600;color:var(--text);width:36px;text-align:right;">${pct.toFixed(0)}%</span>
  </div>`;
}

// ── Render current arm ────────────────────────────────────────────
function renderCurrentArm(data) {
  const badge     = document.getElementById('current-arm-badge');
  const ctxKey    = document.getElementById('current-context-key');
  const paramsEl  = document.getElementById('arm-params');

  const arm = data.current_arm || data.selected_arm || {};
  const name = arm.name || arm.arm || data.current_arm_name || '—';

  if (badge) {
    badge.textContent = escHtml(String(name).toUpperCase());
    badge.className = 'arm-badge ' + armClass(name);
  }
  if (ctxKey) {
    ctxKey.textContent = escHtml(String(data.context_key || data.current_context || '—'));
  }
  if (paramsEl) {
    const params = arm.parameters || arm.params || {};
    const entries = Object.entries(params);
    if (entries.length === 0) {
      paramsEl.innerHTML = '<span style="color:var(--text-dim);font-size:0.82rem;">No arm parameters.</span>';
    } else {
      paramsEl.innerHTML = entries.map(([k, v]) => `
        <div style="background:var(--bg3);border:1px solid var(--border);border-radius:6px;padding:0.4rem 0.75rem;font-size:0.8rem;">
          <span style="color:var(--text-dim);">${escHtml(k)}</span>
          <span style="color:var(--text);font-weight:600;margin-left:0.4rem;">${escHtml(String(v))}</span>
        </div>`).join('');
    }
  }
}

// ── Render arm stats table ────────────────────────────────────────
function renderArmStatsTable(arms) {
  const tbody = document.getElementById('arm-stats-tbody');
  if (!tbody) return;
  if (!arms || arms.length === 0) {
    tbody.innerHTML = '<tr><td colspan="6" style="color:var(--text-dim);text-align:center;padding:2rem;">No arm data available.</td></tr>';
    return;
  }
  tbody.innerHTML = arms.map(a => {
    const alpha = a.alpha !== undefined ? a.alpha : (a.successes || 0);
    const beta  = a.beta  !== undefined ? a.beta  : (a.failures  || 0);
    const mean  = (alpha + beta) > 0 ? alpha / (alpha + beta) : 0;
    const pulls = a.pull_count || a.pulls || a.n || (alpha + beta);
    return `<tr class="memory-row">
      <td style="font-weight:600;">${escHtml(String(a.name || a.arm || '—'))}</td>
      <td style="color:var(--blue);font-variant-numeric:tabular-nums;">${Number(alpha).toFixed(1)}</td>
      <td style="color:var(--red);font-variant-numeric:tabular-nums;">${Number(beta).toFixed(1)}</td>
      <td style="font-variant-numeric:tabular-nums;">${(mean * 100).toFixed(1)}%</td>
      <td style="color:var(--text-dim);font-variant-numeric:tabular-nums;">${escHtml(String(pulls))}</td>
      <td>${winRateBar(mean)}</td>
    </tr>`;
  }).join('');
}

// ── Update posterior chart ────────────────────────────────────────
function updatePosteriorChart(arms) {
  if (!posteriorChart || !arms || arms.length === 0) return;
  const labels = arms.map(a => a.name || a.arm || '?');
  const alphas = arms.map(a => a.alpha !== undefined ? a.alpha : (a.successes || 0));
  const betas  = arms.map(a => a.beta  !== undefined ? a.beta  : (a.failures  || 0));
  const means  = alphas.map((al, i) => {
    const total = al + betas[i];
    return total > 0 ? al / total : 0;
  });

  posteriorChart.data.labels = labels;
  posteriorChart.data.datasets = [
    {
      label: 'Alpha (successes)',
      data: alphas,
      backgroundColor: 'rgba(74,158,255,0.6)',
      borderColor: '#4a9eff',
      borderWidth: 1,
      yAxisID: 'y',
    },
    {
      label: 'Beta (failures)',
      data: betas,
      backgroundColor: 'rgba(255,45,85,0.5)',
      borderColor: '#ff2d55',
      borderWidth: 1,
      yAxisID: 'y',
    },
    {
      label: 'Mean Success Rate',
      data: means,
      type: 'line',
      borderColor: '#34d399',
      backgroundColor: 'rgba(52,211,153,0.15)',
      borderWidth: 2,
      pointRadius: 4,
      tension: 0.3,
      fill: false,
      yAxisID: 'y2',
    },
  ];
  posteriorChart.update();
}

// ── Update history chart ──────────────────────────────────────────
function updateHistoryChart(history) {
  if (!historyChart) return;
  const sessions = history || [];
  if (sessions.length === 0) {
    historyChart.data.labels = [];
    historyChart.data.datasets = [];
    historyChart.update();
    return;
  }

  // Collect all arm names
  const armSet = new Set();
  sessions.forEach(s => {
    if (s.arm_counts) Object.keys(s.arm_counts).forEach(k => armSet.add(k));
    else if (s.arm) armSet.add(s.arm);
  });
  const armNames = Array.from(armSet);

  const labels = sessions.map((s, i) => s.session_label || ('S' + (s.session || i + 1)));

  historyChart.data.labels = labels;
  historyChart.data.datasets = armNames.map((arm, idx) => ({
    label: arm,
    data: sessions.map(s => {
      if (s.arm_counts) return s.arm_counts[arm] || 0;
      return s.arm === arm ? 1 : 0;
    }),
    borderColor: ARM_COLORS[idx % ARM_COLORS.length],
    backgroundColor: ARM_COLORS[idx % ARM_COLORS.length].replace(')', ',0.25)').replace('rgb(', 'rgba(').replace('#', '').slice(0, 0) + ARM_COLORS[idx % ARM_COLORS.length] + '40',
    borderWidth: 2,
    pointRadius: 2,
    tension: 0.3,
    fill: true,
  }));
  historyChart.update();
}

// ── Update context chart ──────────────────────────────────────────
function updateContextChart(data) {
  if (!contextChart) return;
  const bins = data.context_bins || data.context_distribution || {};
  const entries = Object.entries(bins);

  if (entries.length === 0) {
    contextChart.data.labels = ['(no context data)'];
    contextChart.data.datasets = [{ label: 'Selections', data: [0], backgroundColor: 'rgba(74,158,255,0.5)', borderColor: '#4a9eff', borderWidth: 1 }];
    contextChart.update();
    return;
  }

  // Group by arm if nested, or flat context_key -> count
  const armSet = new Set();
  entries.forEach(([, val]) => {
    if (typeof val === 'object' && val !== null) Object.keys(val).forEach(k => armSet.add(k));
  });

  const labels = entries.map(([k]) => k);
  if (armSet.size > 1) {
    const armNames = Array.from(armSet);
    contextChart.data.labels = labels;
    contextChart.data.datasets = armNames.map((arm, idx) => ({
      label: arm,
      data: entries.map(([, val]) => (typeof val === 'object' && val !== null ? (val[arm] || 0) : 0)),
      backgroundColor: ARM_COLORS[idx % ARM_COLORS.length] + '99',
      borderColor: ARM_COLORS[idx % ARM_COLORS.length],
      borderWidth: 1,
    }));
  } else {
    contextChart.data.labels = labels;
    contextChart.data.datasets = [{
      label: 'Times Selected',
      data: entries.map(([, val]) => typeof val === 'object' ? Object.values(val).reduce((a, b) => a + b, 0) : Number(val) || 0),
      backgroundColor: 'rgba(74,158,255,0.55)',
      borderColor: '#4a9eff',
      borderWidth: 1,
    }];
  }
  contextChart.update();
}

// ── Load bandits state ────────────────────────────────────────────
async function loadBanditsState() {
  try {
    const resp = await fetch('/api/bandits-state');
    if (!resp.ok) return;
    const data = await resp.json();

    renderCurrentArm(data);

    const arms = data.arms || data.arm_stats || [];
    const armList = Array.isArray(arms) ? arms : Object.entries(arms).map(([name, stats]) => ({ name, ...stats }));
    renderArmStatsTable(armList);
    updatePosteriorChart(armList);
    updateContextChart(data);
  } catch (e) {
    console.warn('Bandits: bandits-state load failed', e);
  }
}

async function loadBanditsHistory() {
  try {
    const resp = await fetch('/api/bandits-history');
    if (!resp.ok) return;
    const data = await resp.json();
    const history = data.history || data.sessions || data || [];
    const recent = Array.isArray(history) ? history.slice(-20) : [];
    updateHistoryChart(recent);
  } catch (e) {
    console.warn('Bandits: bandits-history load failed', e);
  }
}

// ── Boot ──────────────────────────────────────────────────────────
initPosteriorChart();
initHistoryChart();
initContextChart();

loadBanditsState();
loadBanditsHistory();

setInterval(loadBanditsState,   3000);
setInterval(loadBanditsHistory, 30000);
</script>
{% endblock %}