neural-network-study 0.5.2

use neural_network_study::{ActivationFunction, NeuralNetwork};
use rand::{Rng, SeedableRng, rngs::StdRng, seq::SliceRandom};
use serde::Serialize;
use std::{
    error::Error,
    fs,
    path::{Path, PathBuf},
    time::Instant,
};

const OUTPUT_PATH: &str = "target/depth-benchmark.html";
const GRID_SIZE: usize = 160;
const TRAIN_SPLIT: f64 = 0.8;
const MAX_PLOT_SAMPLES: usize = 280;
const HISTORY_SAMPLE_LIMIT: usize = 180;

#[derive(Clone)]
struct Sample {
    x: f64,
    y: f64,
    label: f64,
}

#[derive(Clone)]
struct DatasetSpec {
    slug: &'static str,
    name: &'static str,
    description: &'static str,
    learning_rate: f64,
    epochs: usize,
    train_samples: Vec<Sample>,
    validation_samples: Vec<Sample>,
    plot_samples: Vec<Sample>,
}

#[derive(Clone)]
struct ArchitectureSpec {
    slug: &'static str,
    name: &'static str,
    description: &'static str,
    layer_sizes: Vec<usize>,
}

#[derive(Serialize)]
struct DatasetView {
    slug: &'static str,
    name: &'static str,
    description: &'static str,
    learning_rate: f64,
    epochs: usize,
    train_count: usize,
    validation_count: usize,
    depth_gain: f64,
    grid_size: usize,
    samples: Vec<SampleView>,
    models: Vec<ModelView>,
}

#[derive(Serialize)]
struct SampleView {
    x: f64,
    y: f64,
    label: f64,
}

#[derive(Serialize)]
struct ModelView {
    slug: &'static str,
    name: &'static str,
    description: &'static str,
    layer_sizes: Vec<usize>,
    parameter_count: usize,
    duration_ms: u64,
    train_accuracy: f64,
    validation_accuracy: f64,
    train_loss: f64,
    validation_loss: f64,
    history: Vec<HistoryPoint>,
    boundary: Vec<u8>,
}

#[derive(Serialize)]
struct HistoryPoint {
    epoch: usize,
    train_loss: f64,
    validation_loss: f64,
    train_accuracy: f64,
    validation_accuracy: f64,
}

fn main() -> Result<(), Box<dyn Error>> {
    let architectures = vec![
        ArchitectureSpec {
            slug: "shallow",
            name: "1 Hidden Layer",
            description: "2 -> 16 -> 1",
            layer_sizes: vec![2, 16, 1],
        },
        ArchitectureSpec {
            slug: "deep-2",
            name: "2 Hidden Layers",
            description: "2 -> 16 -> 16 -> 1",
            layer_sizes: vec![2, 16, 16, 1],
        },
        ArchitectureSpec {
            slug: "deep-3",
            name: "3 Hidden Layers",
            description: "2 -> 16 -> 16 -> 16 -> 1",
            layer_sizes: vec![2, 16, 16, 16, 1],
        },
    ];

    let datasets = vec![make_ring_dataset(), make_spiral_dataset()];
    let views: Vec<DatasetView> = datasets
        .into_iter()
        .enumerate()
        .map(|(index, dataset)| benchmark_dataset(dataset, &architectures, 1_000 + index as u64))
        .collect::<Result<_, _>>()?;

    let output_path = PathBuf::from(OUTPUT_PATH);
    write_playground(&output_path, &views)?;

    println!(
        "Depth benchmark playground written to {}",
        output_path.canonicalize().unwrap_or(output_path).display()
    );

    Ok(())
}

fn benchmark_dataset(
    dataset: DatasetSpec,
    architectures: &[ArchitectureSpec],
    seed: u64,
) -> Result<DatasetView, Box<dyn Error>> {
    let models: Vec<ModelView> = architectures
        .iter()
        .enumerate()
        .map(|(idx, architecture)| train_model(&dataset, architecture, seed + idx as u64))
        .collect::<Result<_, _>>()?;

    let depth_gain = if models.len() >= 2 {
        let shallow = models
            .first()
            .map(|model| model.validation_accuracy)
            .unwrap_or(0.0);
        let best_deep = models
            .iter()
            .skip(1)
            .map(|model| model.validation_accuracy)
            .fold(f64::NEG_INFINITY, f64::max);
        best_deep - shallow
    } else {
        0.0
    };

    let samples = dataset
        .plot_samples
        .iter()
        .map(|sample| SampleView {
            x: sample.x,
            y: sample.y,
            label: sample.label,
        })
        .collect();

    Ok(DatasetView {
        slug: dataset.slug,
        name: dataset.name,
        description: dataset.description,
        learning_rate: dataset.learning_rate,
        epochs: dataset.epochs,
        train_count: dataset.train_samples.len(),
        validation_count: dataset.validation_samples.len(),
        depth_gain,
        grid_size: GRID_SIZE,
        samples,
        models,
    })
}

fn train_model(
    dataset: &DatasetSpec,
    architecture: &ArchitectureSpec,
    seed: u64,
) -> Result<ModelView, Box<dyn Error>> {
    let mut rng = StdRng::seed_from_u64(seed);
    let mut network = NeuralNetwork::new(architecture.layer_sizes.clone(), Some(&mut rng))?;
    network.set_activation_function(ActivationFunction::Tanh);
    network.set_learning_rate(dataset.learning_rate);

    let mut shuffled = dataset.train_samples.clone();
    let checkpoint_stride = (dataset.epochs / 24).max(1);
    let mut history = Vec::new();
    let history_train_len = dataset.train_samples.len().min(HISTORY_SAMPLE_LIMIT);
    let history_validation_len = dataset.validation_samples.len().min(HISTORY_SAMPLE_LIMIT);

    let start = Instant::now();
    for epoch in 1..=dataset.epochs {
        shuffled.shuffle(&mut rng);
        for sample in &shuffled {
            // Match tanh's output range to the probability mapping used in the report.
            let target = 2.0 * sample.label - 1.0;
            network.train(vec![sample.x, sample.y], vec![target])?;
        }

        if epoch == 1 || epoch % checkpoint_stride == 0 || epoch == dataset.epochs {
            let (train_accuracy, train_loss) =
                evaluate_model(&network, &dataset.train_samples[..history_train_len])?;
            let (validation_accuracy, validation_loss) = evaluate_model(
                &network,
                &dataset.validation_samples[..history_validation_len],
            )?;
            history.push(HistoryPoint {
                epoch,
                train_loss,
                validation_loss,
                train_accuracy,
                validation_accuracy,
            });
        }
    }
    let duration_ms = (start.elapsed().as_secs_f64() * 1000.0).round() as u64;

    let (train_accuracy, train_loss) = evaluate_model(&network, &dataset.train_samples)?;
    let (validation_accuracy, validation_loss) =
        evaluate_model(&network, &dataset.validation_samples)?;
    let boundary = sample_decision_boundary(&network)?;

    Ok(ModelView {
        slug: architecture.slug,
        name: architecture.name,
        description: architecture.description,
        layer_sizes: architecture.layer_sizes.clone(),
        parameter_count: parameter_count(&architecture.layer_sizes),
        duration_ms,
        train_accuracy,
        validation_accuracy,
        train_loss,
        validation_loss,
        history,
        boundary,
    })
}

fn evaluate_model(
    network: &NeuralNetwork,
    samples: &[Sample],
) -> Result<(f64, f64), Box<dyn Error>> {
    if samples.is_empty() {
        return Ok((0.0, 0.0));
    }

    let mut correct = 0usize;
    let mut total_loss = 0.0;
    for sample in samples {
        let probability = probability_from_tanh_output(network.predict(vec![sample.x, sample.y])?)
            .clamp(1e-7, 1.0 - 1e-7);
        if classify(probability) == classify(sample.label) {
            correct += 1;
        }
        total_loss +=
            -(sample.label * probability.ln() + (1.0 - sample.label) * (1.0 - probability).ln());
    }

    Ok((
        correct as f64 / samples.len() as f64,
        total_loss / samples.len() as f64,
    ))
}

fn sample_decision_boundary(network: &NeuralNetwork) -> Result<Vec<u8>, Box<dyn Error>> {
    let mut boundary = Vec::with_capacity(GRID_SIZE * GRID_SIZE);
    for grid_y in 0..GRID_SIZE {
        for grid_x in 0..GRID_SIZE {
            let x = grid_x as f64 / (GRID_SIZE - 1) as f64;
            let y = 1.0 - grid_y as f64 / (GRID_SIZE - 1) as f64;
            let prediction = network.predict(vec![x, y])?[0];
            boundary.push((tanh_output_to_probability(prediction) * 255.0).round() as u8);
        }
    }
    Ok(boundary)
}

fn write_playground(path: &Path, datasets: &[DatasetView]) -> Result<(), Box<dyn Error>> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }

    let dataset_json = serde_json::to_string(datasets)?;
    let html = format!(
        r#"<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>Depth Benchmark Playground</title>
  <style>
    :root {{
      --bg: #eff3ea;
      --panel: rgba(254, 255, 252, 0.88);
      --ink: #1f2630;
      --muted: #5d6a77;
      --accent: #1a8f72;
      --accent-2: #c26a1a;
      --border: rgba(31, 38, 48, 0.12);
      --shadow: 0 18px 48px rgba(32, 58, 44, 0.12);
    }}

    * {{
      box-sizing: border-box;
    }}

    body {{
      margin: 0;
      min-height: 100vh;
      font-family: "Avenir Next", "Trebuchet MS", sans-serif;
      color: var(--ink);
      background:
        radial-gradient(circle at 15% 20%, rgba(26, 143, 114, 0.14), transparent 28%),
        radial-gradient(circle at 78% 88%, rgba(194, 106, 26, 0.16), transparent 24%),
        linear-gradient(165deg, #f2f6ef 0%, #e7efe0 44%, #e6ebea 100%);
      padding: 30px 14px 48px;
    }}

    .shell {{
      max-width: 1280px;
      margin: 0 auto;
      display: grid;
      gap: 20px;
    }}

    .hero {{
      display: grid;
      gap: 8px;
    }}

    .eyebrow {{
      text-transform: uppercase;
      letter-spacing: 0.18em;
      font-size: 0.75rem;
      color: var(--accent);
      font-weight: 700;
    }}

    h1 {{
      margin: 0;
      font-size: clamp(2rem, 4vw, 4.2rem);
      line-height: 0.95;
      max-width: 16ch;
    }}

    .lede {{
      margin: 0;
      max-width: 72ch;
      color: var(--muted);
      line-height: 1.6;
    }}

    .panel {{
      background: var(--panel);
      border: 1px solid var(--border);
      border-radius: 24px;
      box-shadow: var(--shadow);
      backdrop-filter: blur(8px);
    }}

    .controls {{
      display: grid;
      gap: 16px;
      padding: 20px;
    }}

    .control-row {{
      display: grid;
      gap: 12px;
      grid-template-columns: minmax(0, 340px) minmax(0, 1fr);
      align-items: end;
    }}

    .label {{
      display: grid;
      gap: 8px;
      font-size: 0.9rem;
      color: var(--muted);
    }}

    select {{
      width: 100%;
      border: 1px solid rgba(31, 38, 48, 0.2);
      border-radius: 14px;
      padding: 12px 14px;
      font: inherit;
      color: var(--ink);
      background: #fff;
    }}

    .dataset-copy {{
      display: grid;
      gap: 6px;
    }}

    .dataset-copy h2 {{
      margin: 0;
      font-size: 1.6rem;
      line-height: 1.08;
    }}

    .dataset-copy p {{
      margin: 0;
      color: var(--muted);
      line-height: 1.55;
    }}

    .stats {{
      display: grid;
      gap: 12px;
      grid-template-columns: repeat(4, minmax(0, 1fr));
    }}

    .stat {{
      padding: 12px;
      border-radius: 16px;
      background: rgba(255, 255, 255, 0.72);
      border: 1px solid rgba(31, 38, 48, 0.08);
    }}

    .stat-label {{
      margin: 0 0 4px;
      font-size: 0.73rem;
      color: var(--muted);
      text-transform: uppercase;
      letter-spacing: 0.08em;
    }}

    .stat-value {{
      margin: 0;
      font-size: 1.05rem;
      font-weight: 700;
    }}

    .rankings {{
      padding: 14px;
      border-radius: 16px;
      border: 1px solid rgba(31, 38, 48, 0.09);
      background: rgba(255, 255, 255, 0.62);
      color: var(--muted);
      font-size: 0.94rem;
      line-height: 1.5;
    }}

    .rankings strong {{
      color: var(--ink);
    }}

    .models {{
      display: grid;
      gap: 14px;
      grid-template-columns: repeat(3, minmax(0, 1fr));
    }}

    .model {{
      padding: 16px;
      display: grid;
      gap: 12px;
    }}

    .model h3 {{
      margin: 0;
      font-size: 1.2rem;
    }}

    .model .subtitle {{
      margin: 0;
      font-size: 0.9rem;
      color: var(--muted);
    }}

    .model-stats {{
      display: grid;
      gap: 8px;
      grid-template-columns: repeat(2, minmax(0, 1fr));
    }}

    .chip {{
      border: 1px solid rgba(31, 38, 48, 0.1);
      border-radius: 12px;
      padding: 8px 10px;
      background: rgba(255, 255, 255, 0.72);
    }}

    .chip-label {{
      display: block;
      font-size: 0.7rem;
      text-transform: uppercase;
      letter-spacing: 0.07em;
      color: var(--muted);
      margin-bottom: 3px;
    }}

    .chip-value {{
      font-size: 0.92rem;
      font-weight: 600;
    }}

    .canvas-stack {{
      display: grid;
      gap: 8px;
    }}

    .canvas-frame {{
      position: relative;
      width: 100%;
      aspect-ratio: 1;
      border: 1px solid rgba(31, 38, 48, 0.12);
      border-radius: 16px;
      overflow: hidden;
      background:
        linear-gradient(rgba(255,255,255,0.15) 1px, transparent 1px),
        linear-gradient(90deg, rgba(255,255,255,0.15) 1px, transparent 1px),
        linear-gradient(180deg, rgba(16, 27, 39, 0.04), rgba(16, 27, 39, 0.08));
      background-size: 10% 10%, 10% 10%, 100% 100%;
    }}

    canvas {{
      display: block;
      width: 100%;
      height: 100%;
    }}

    .curve-frame {{
      width: 100%;
      height: 146px;
      border: 1px solid rgba(31, 38, 48, 0.12);
      border-radius: 12px;
      overflow: hidden;
      background: rgba(255, 255, 255, 0.86);
    }}

    .legend {{
      margin: 0;
      font-size: 0.8rem;
      color: var(--muted);
      display: flex;
      gap: 12px;
      flex-wrap: wrap;
    }}

    .legend b {{
      color: var(--ink);
    }}

    @media (max-width: 1050px) {{
      .models {{
        grid-template-columns: 1fr;
      }}

      .stats {{
        grid-template-columns: repeat(2, minmax(0, 1fr));
      }}

      .control-row {{
        grid-template-columns: 1fr;
      }}
    }}
  </style>
</head>
<body>
  <main class="shell">
    <section class="hero">
      <span class="eyebrow">Neural Network Study</span>
      <h1>Depth Benchmark Playground</h1>
      <p class="lede">
        Same optimizer, same dataset, same neurons per hidden layer. Deeper networks also
        have more parameters, and each architecture uses a single random seed.
        Pick a dataset and compare how boundary shape, loss trajectory, and validation accuracy
        shift as we stack more hidden layers.
      </p>
    </section>

    <section class="panel controls">
      <div class="control-row">
        <label class="label" for="dataset-select">
          Dataset
          <select id="dataset-select"></select>
        </label>
        <div class="dataset-copy">
          <h2 id="dataset-title"></h2>
          <p id="dataset-description"></p>
        </div>
      </div>

      <div class="stats">
        <div class="stat">
          <p class="stat-label">Train / Validation Samples</p>
          <p class="stat-value" id="sample-counts"></p>
        </div>
        <div class="stat">
          <p class="stat-label">Epochs</p>
          <p class="stat-value" id="epochs"></p>
        </div>
        <div class="stat">
          <p class="stat-label">Learning Rate</p>
          <p class="stat-value" id="learning-rate"></p>
        </div>
        <div class="stat">
          <p class="stat-label">Best Deep Gain</p>
          <p class="stat-value" id="depth-gain"></p>
        </div>
      </div>

      <div class="rankings" id="rankings"></div>
    </section>

    <section class="models" id="models"></section>
  </main>

  <script>
    const datasets = {dataset_json};
    const select = document.getElementById('dataset-select');
    const title = document.getElementById('dataset-title');
    const description = document.getElementById('dataset-description');
    const sampleCounts = document.getElementById('sample-counts');
    const epochs = document.getElementById('epochs');
    const learningRate = document.getElementById('learning-rate');
    const depthGain = document.getElementById('depth-gain');
    const rankings = document.getElementById('rankings');
    const modelsRoot = document.getElementById('models');

    const classColors = [
      [245, 220, 124],
      [24, 128, 152],
    ];

    for (const dataset of datasets) {{
      const option = document.createElement('option');
      option.value = dataset.slug;
      option.textContent = dataset.name;
      select.appendChild(option);
    }}

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

    function formatPercent(value) {{
      return `${{(value * 100).toFixed(1)}}%`;
    }}

    function drawBoundary(canvas, dataset, model) {{
      const ctx = canvas.getContext('2d');
      const width = canvas.width;
      const height = canvas.height;
      const image = ctx.createImageData(width, height);

      for (let y = 0; y < height; y++) {{
        const sampleY = Math.min(dataset.grid_size - 1, Math.floor((y / height) * dataset.grid_size));
        for (let x = 0; x < width; x++) {{
          const sampleX = Math.min(dataset.grid_size - 1, Math.floor((x / width) * dataset.grid_size));
          const t = model.boundary[sampleY * dataset.grid_size + sampleX] / 255;
          const red = Math.round(lerp(classColors[0][0], classColors[1][0], t));
          const green = Math.round(lerp(classColors[0][1], classColors[1][1], t));
          const blue = Math.round(lerp(classColors[0][2], classColors[1][2], t));
          const offset = (y * width + x) * 4;
          image.data[offset] = red;
          image.data[offset + 1] = green;
          image.data[offset + 2] = blue;
          image.data[offset + 3] = 255;
        }}
      }}

      ctx.putImageData(image, 0, 0);

      for (const sample of dataset.samples) {{
        const x = sample.x * width;
        const y = (1 - sample.y) * height;
        const color = sample.label >= 0.5 ? classColors[1] : classColors[0];
        ctx.beginPath();
        ctx.arc(x, y, 3.6, 0, Math.PI * 2);
        ctx.fillStyle = `rgb(${{color[0]}}, ${{color[1]}}, ${{color[2]}})`;
        ctx.fill();
        ctx.lineWidth = 1;
        ctx.strokeStyle = 'rgba(31, 38, 48, 0.32)';
        ctx.stroke();
      }}
    }}

    function drawLossCurve(canvas, history) {{
      const ctx = canvas.getContext('2d');
      const width = canvas.width;
      const height = canvas.height;
      const pad = {{ left: 36, right: 14, top: 14, bottom: 24 }};

      ctx.clearRect(0, 0, width, height);

      if (history.length < 2) {{
        return;
      }}

      let minLoss = Infinity;
      let maxLoss = -Infinity;
      let maxEpoch = 1;
      for (const point of history) {{
        minLoss = Math.min(minLoss, point.train_loss, point.validation_loss);
        maxLoss = Math.max(maxLoss, point.train_loss, point.validation_loss);
        maxEpoch = Math.max(maxEpoch, point.epoch);
      }}

      if (Math.abs(maxLoss - minLoss) < 1e-9) {{
        maxLoss += 1;
      }}

      function mapX(epoch) {{
        return pad.left + ((epoch / maxEpoch) * (width - pad.left - pad.right));
      }}

      function mapY(loss) {{
        const t = (loss - minLoss) / (maxLoss - minLoss);
        return height - pad.bottom - t * (height - pad.top - pad.bottom);
      }}

      ctx.strokeStyle = 'rgba(31, 38, 48, 0.22)';
      ctx.lineWidth = 1;
      ctx.beginPath();
      ctx.moveTo(pad.left, pad.top);
      ctx.lineTo(pad.left, height - pad.bottom);
      ctx.lineTo(width - pad.right, height - pad.bottom);
      ctx.stroke();

      function plot(accessor, color) {{
        ctx.strokeStyle = color;
        ctx.lineWidth = 2.2;
        ctx.beginPath();
        history.forEach((point, index) => {{
          const x = mapX(point.epoch);
          const y = mapY(accessor(point));
          if (index === 0) {{
            ctx.moveTo(x, y);
          }} else {{
            ctx.lineTo(x, y);
          }}
        }});
        ctx.stroke();
      }}

      plot((point) => point.train_loss, '#1a8f72');
      plot((point) => point.validation_loss, '#c26a1a');

      ctx.fillStyle = 'rgba(31, 38, 48, 0.72)';
      ctx.font = '11px "Avenir Next", "Trebuchet MS", sans-serif';
      ctx.fillText(`loss ${{maxLoss.toFixed(2)}}`, 6, pad.top + 4);
      ctx.fillText(`loss ${{minLoss.toFixed(2)}}`, 6, height - pad.bottom - 2);
      ctx.fillText(`epoch ${{maxEpoch}}`, width - pad.right - 66, height - 6);
    }}

    function makeCard(dataset, model) {{
      const card = document.createElement('article');
      card.className = 'panel model';

      card.innerHTML = `
        <div>
          <h3>${{model.name}}</h3>
          <p class="subtitle">${{model.description}}</p>
        </div>
        <div class="model-stats">
          <div class="chip">
            <span class="chip-label">Validation Accuracy</span>
            <span class="chip-value">${{formatPercent(model.validation_accuracy)}}</span>
          </div>
          <div class="chip">
            <span class="chip-label">Validation Loss</span>
            <span class="chip-value">${{model.validation_loss.toFixed(3)}}</span>
          </div>
          <div class="chip">
            <span class="chip-label">Parameters</span>
            <span class="chip-value">${{model.parameter_count.toLocaleString()}}</span>
          </div>
          <div class="chip">
            <span class="chip-label">Training Time</span>
            <span class="chip-value">${{(model.duration_ms / 1000).toFixed(2)}}s</span>
          </div>
        </div>
        <div class="canvas-stack">
          <div class="canvas-frame">
            <canvas width="320" height="320"></canvas>
          </div>
          <div class="curve-frame">
            <canvas width="320" height="146"></canvas>
          </div>
        </div>
        <p class="legend">
          <span><b>Boundary</b>: background confidence + sampled points</span>
          <span><b>Curves</b>: green train loss, orange validation loss</span>
        </p>
      `;

      const canvases = card.querySelectorAll('canvas');
      drawBoundary(canvases[0], dataset, model);
      drawLossCurve(canvases[1], model.history);
      return card;
    }}

    function renderDataset(dataset) {{
      title.textContent = dataset.name;
      description.textContent = dataset.description;
      sampleCounts.textContent = `${{dataset.train_count}} / ${{dataset.validation_count}}`;
      epochs.textContent = dataset.epochs.toLocaleString();
      learningRate.textContent = dataset.learning_rate.toFixed(3);
      const gain = dataset.depth_gain * 100;
      const sign = gain >= 0 ? '+' : '';
      depthGain.textContent = `${{sign}}${{gain.toFixed(1)}} pts`;

      const sorted = [...dataset.models].sort((a, b) => b.validation_accuracy - a.validation_accuracy);
      rankings.innerHTML = `<strong>Validation ranking:</strong> ${{sorted.map((model, index) => `${{index + 1}}) ${{model.name}} (${{formatPercent(model.validation_accuracy)}})`).join(' | ')}}`;

      modelsRoot.innerHTML = '';
      for (const model of dataset.models) {{
        modelsRoot.appendChild(makeCard(dataset, model));
      }}
    }}

    select.addEventListener('change', () => {{
      const selected = datasets.find((dataset) => dataset.slug === select.value);
      if (selected) {{
        renderDataset(selected);
      }}
    }});

    select.value = datasets[0].slug;
    renderDataset(datasets[0]);
  </script>
</body>
</html>
"#
    );

    fs::write(path, html)?;
    Ok(())
}

fn classify(value: f64) -> u8 {
    if value >= 0.5 { 1 } else { 0 }
}

fn probability_from_tanh_output(output: Vec<f64>) -> f64 {
    tanh_output_to_probability(output[0])
}

fn tanh_output_to_probability(value: f64) -> f64 {
    ((value + 1.0) * 0.5).clamp(0.0, 1.0)
}

fn parameter_count(layer_sizes: &[usize]) -> usize {
    layer_sizes
        .windows(2)
        .map(|pair| {
            let fan_in = pair[0];
            let fan_out = pair[1];
            fan_in * fan_out + fan_out
        })
        .sum()
}

fn split_dataset(
    mut all_samples: Vec<Sample>,
    rng: &mut StdRng,
) -> (Vec<Sample>, Vec<Sample>, Vec<Sample>) {
    all_samples.shuffle(rng);

    let train_count = ((all_samples.len() as f64 * TRAIN_SPLIT).round() as usize)
        .clamp(1, all_samples.len().saturating_sub(1));
    let train_samples = all_samples[..train_count].to_vec();
    let validation_samples = all_samples[train_count..].to_vec();

    let mut plot_samples = all_samples;
    plot_samples.shuffle(rng);
    plot_samples.truncate(MAX_PLOT_SAMPLES.min(plot_samples.len()));

    (train_samples, validation_samples, plot_samples)
}

fn make_ring_dataset() -> DatasetSpec {
    let mut rng = StdRng::seed_from_u64(29);
    let mut samples = Vec::with_capacity(300);
    for _ in 0..300 {
        let x = rng.random::<f64>();
        let y = rng.random::<f64>();
        let dx = x - 0.5;
        let dy = y - 0.5;
        let radius = (dx * dx + dy * dy).sqrt();
        let label = if radius > 0.2 && radius < 0.34 {
            1.0
        } else {
            0.0
        };
        samples.push(Sample { x, y, label });
    }

    let (train_samples, validation_samples, plot_samples) = split_dataset(samples, &mut rng);

    DatasetSpec {
        slug: "ring",
        name: "Concentric Ring",
        description: "Class 1 wraps around class 0 in a thin annulus. Deeper models usually track the loop more cleanly and avoid disconnected artifacts.",
        learning_rate: 0.12,
        epochs: 760,
        train_samples,
        validation_samples,
        plot_samples,
    }
}

fn make_spiral_dataset() -> DatasetSpec {
    let mut rng = StdRng::seed_from_u64(47);
    let mut samples = Vec::with_capacity(400);
    let points_per_class = 200;
    for i in 0..points_per_class {
        let t = i as f64 / (points_per_class - 1) as f64;
        let radius = 0.06 + 0.42 * t;
        let angle = 1.2 + 3.8 * std::f64::consts::PI * t;

        for class in 0..2 {
            let phase = if class == 0 {
                0.0
            } else {
                std::f64::consts::PI
            };
            let jitter_angle = rng.random_range(-0.16..0.16);
            let jitter_radius = rng.random_range(-0.018..0.018);
            let noise_x = rng.random_range(-0.012..0.012);
            let noise_y = rng.random_range(-0.012..0.012);
            let r = (radius + jitter_radius).clamp(0.02, 0.49);
            let theta = angle + phase + jitter_angle;

            let x = (0.5 + r * theta.cos() + noise_x).clamp(0.0, 1.0);
            let y = (0.5 + r * theta.sin() + noise_y).clamp(0.0, 1.0);
            let label = class as f64;
            samples.push(Sample { x, y, label });
        }
    }

    let (train_samples, validation_samples, plot_samples) = split_dataset(samples, &mut rng);

    DatasetSpec {
        slug: "spirals",
        name: "Two Spirals",
        description: "The depth stress test. Shallow networks tend to smear this shape, while extra hidden layers can model the repeated twists more faithfully.",
        learning_rate: 0.08,
        epochs: 1_100,
        train_samples,
        validation_samples,
        plot_samples,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn binary_cross_entropy(label: f64, probability: f64) -> f64 {
        let probability = probability.clamp(1e-7, 1.0 - 1e-7);
        -(label * probability.ln() + (1.0 - label) * (1.0 - probability).ln())
    }

    #[test]
    fn training_learns_confident_predictions_for_both_classes() -> Result<(), Box<dyn Error>> {
        let samples = vec![
            Sample {
                x: 0.0,
                y: 0.0,
                label: 0.0,
            },
            Sample {
                x: 1.0,
                y: 1.0,
                label: 1.0,
            },
        ];
        let dataset = DatasetSpec {
            slug: "separable",
            name: "Separable",
            description: "Two-class training regression",
            learning_rate: 0.1,
            epochs: 500,
            train_samples: samples.clone(),
            validation_samples: samples.clone(),
            plot_samples: samples,
        };
        let architecture = ArchitectureSpec {
            slug: "probe",
            name: "Probe",
            description: "2 -> 1",
            layer_sizes: vec![2, 1],
        };

        let model = train_model(&dataset, &architecture, 5)?;
        assert_eq!(model.validation_accuracy, 1.0);
        assert!(
            model.validation_loss < 0.1,
            "both classes should learn confident predictions; loss was {}",
            model.validation_loss
        );
        assert!(
            model.boundary[GRID_SIZE - 1] > 229,
            "class one should exceed 90% probability"
        );
        assert!(
            model.boundary[(GRID_SIZE - 1) * GRID_SIZE] < 26,
            "class zero should be below 10% probability"
        );
        Ok(())
    }

    #[test]
    fn tanh_outputs_are_mapped_to_probabilities_before_reporting_loss() -> Result<(), Box<dyn Error>>
    {
        let seed = 5;
        let sample = Sample {
            x: 0.8,
            y: 0.2,
            label: 1.0,
        };
        let architecture = ArchitectureSpec {
            slug: "probe",
            name: "Probe",
            description: "2 -> 1",
            layer_sizes: vec![2, 1],
        };
        let dataset = DatasetSpec {
            slug: "probe",
            name: "Probe",
            description: "Single-sample loss probe",
            learning_rate: 0.1,
            epochs: 0,
            train_samples: vec![sample.clone()],
            validation_samples: vec![sample.clone()],
            plot_samples: vec![sample.clone()],
        };

        let mut rng = StdRng::seed_from_u64(seed);
        let mut network = NeuralNetwork::new(architecture.layer_sizes.clone(), Some(&mut rng))?;
        network.set_activation_function(ActivationFunction::Tanh);
        let raw_output = network.predict(vec![sample.x, sample.y])?[0];
        let probability = (raw_output + 1.0) * 0.5;
        let expected_loss = binary_cross_entropy(sample.label, probability);

        let model = train_model(&dataset, &architecture, seed)?;

        assert!(
            !(0.0..=1.0).contains(&raw_output),
            "probe seed should produce a tanh output outside the probability range"
        );
        assert!(
            (model.validation_loss - expected_loss).abs() < 1e-12,
            "reported loss should use tanh output mapped from [-1, 1] into [0, 1]"
        );

        Ok(())
    }
}