import { RingBuffer, COLORS, EMA } from "./utils.js";
import { ENTITY_TYPES, measureDiversity } from "./genetics.js";

const HISTORY_LENGTH = 200;

export class StatsTracker {
  constructor() {
    this.generationCount = 0;
    this.totalBirths = 0;
    this.totalDeaths = 0;
    this.totalTime = 0;
    this.peakPopulation = 0;

    this.popHistory = {};
    this.totalPopHistory = new RingBuffer(HISTORY_LENGTH);
    for (const type of Object.keys(ENTITY_TYPES)) {
      this.popHistory[type] = new RingBuffer(HISTORY_LENGTH);
    }

    this.avgFitnessHistory = new RingBuffer(HISTORY_LENGTH);
    this.maxFitnessHistory = new RingBuffer(HISTORY_LENGTH);

    this.diversityHistory = new RingBuffer(HISTORY_LENGTH);

    this.fpsEMA = new EMA(0.1);
    this.entityCountEMA = new EMA(0.2);

    this.current = {
      population: {},
      totalPop: 0,
      avgFitness: 0,
      maxFitness: 0,
      diversity: 0,
      fps: 60,
      entityCount: 0,
      particleCount: 0,
    };

    this.heatmapCols = 40;
    this.heatmapRows = 30;
    this.heatmap = new Float32Array(this.heatmapCols * this.heatmapRows);

    this._recordTimer = 0;
    this._recordInterval = 0.5;
  }

  record(entities, world, dt, fps, particleCount) {
    this.totalTime += dt;
    this._recordTimer += dt;

    this.current.fps = this.fpsEMA.update(fps);
    this.current.entityCount = entities.length;
    this.current.particleCount = particleCount;

    if (this._recordTimer < this._recordInterval) return;
    this._recordTimer = 0;

    const popCounts = {};
    for (const type of Object.keys(ENTITY_TYPES)) {
      popCounts[type] = 0;
    }
    let totalFitness = 0;
    let maxFitness = 0;
    const genomes = [];

    this.heatmap.fill(0);
    const cellW = world.width / this.heatmapCols;
    const cellH = world.height / this.heatmapRows;

    for (const e of entities) {
      popCounts[e.type] = (popCounts[e.type] || 0) + 1;
      totalFitness += e.fitness;
      if (e.fitness > maxFitness) maxFitness = e.fitness;
      genomes.push(e.genome);

      const col = Math.floor(e.x / cellW);
      const row = Math.floor(e.y / cellH);
      if (
        col >= 0 &&
        col < this.heatmapCols &&
        row >= 0 &&
        row < this.heatmapRows
      ) {
        this.heatmap[row * this.heatmapCols + col] += 1;
      }
    }

    const totalPop = entities.length;
    this.current.population = popCounts;
    this.current.totalPop = totalPop;
    this.current.avgFitness = totalPop > 0 ? totalFitness / totalPop : 0;
    this.current.maxFitness = maxFitness;
    this.current.diversity = measureDiversity(genomes);

    if (totalPop > this.peakPopulation) this.peakPopulation = totalPop;

    this.totalPopHistory.push(totalPop);
    for (const type of Object.keys(ENTITY_TYPES)) {
      this.popHistory[type].push(popCounts[type] || 0);
    }
    this.avgFitnessHistory.push(this.current.avgFitness);
    this.maxFitnessHistory.push(maxFitness);
    this.diversityHistory.push(this.current.diversity);
  }

  onBirth() {
    this.totalBirths++;
  }
  onDeath() {
    this.totalDeaths++;
  }
  onGeneration() {
    this.generationCount++;
  }

  renderSparkline(ctx, buffer, x, y, w, h, color) {
    const data = buffer.toArray();
    if (data.length < 2) return;

    let max = -Infinity,
      min = Infinity;
    for (const v of data) {
      if (v > max) max = v;
      if (v < min) min = v;
    }
    const range = max - min || 1;

    ctx.beginPath();
    ctx.strokeStyle = color;
    ctx.lineWidth = 1.5;
    ctx.lineJoin = "round";

    for (let i = 0; i < data.length; i++) {
      const px = x + (i / (data.length - 1)) * w;
      const py = y + h - ((data[i] - min) / range) * h;
      if (i === 0) ctx.moveTo(px, py);
      else ctx.lineTo(px, py);
    }
    ctx.stroke();

    ctx.lineTo(x + w, y + h);
    ctx.lineTo(x, y + h);
    ctx.closePath();
    const grad = ctx.createLinearGradient(x, y, x, y + h);
    const rgbMatch = color.match(/(\d+),\s*(\d+),\s*(\d+)/);
    if (rgbMatch) {
      grad.addColorStop(
        0,
        `rgba(${rgbMatch[1]},${rgbMatch[2]},${rgbMatch[3]},0.15)`,
      );
      grad.addColorStop(
        1,
        `rgba(${rgbMatch[1]},${rgbMatch[2]},${rgbMatch[3]},0.0)`,
      );
    }
    ctx.fillStyle = grad;
    ctx.fill();
  }

  renderCharts(ctx, x, y, w, h) {
    const chartH = h / 3 - 8;
    const pad = 4;

    ctx.fillStyle = "rgba(100, 200, 255, 0.4)";
    ctx.font = "9px Courier New";
    ctx.fillText("POPULATION", x + 2, y + 8);
    this.renderSparkline(
      ctx,
      this.totalPopHistory,
      x,
      y + 12,
      w,
      chartH,
      "rgb(0, 240, 255)",
    );

    const types = Object.keys(ENTITY_TYPES);
    for (const type of types) {
      const c = COLORS[type];
      if (c) {
        this.renderSparkline(
          ctx,
          this.popHistory[type],
          x,
          y + 12,
          w,
          chartH,
          `rgb(${c.r},${c.g},${c.b})`,
        );
      }
    }

    const fy = y + chartH + 20;
    ctx.fillStyle = "rgba(100, 200, 255, 0.4)";
    ctx.fillText("FITNESS", x + 2, fy + 8);
    this.renderSparkline(
      ctx,
      this.avgFitnessHistory,
      x,
      fy + 12,
      w,
      chartH,
      "rgb(0, 255, 136)",
    );
    this.renderSparkline(
      ctx,
      this.maxFitnessHistory,
      x,
      fy + 12,
      w,
      chartH,
      "rgb(255, 170, 0)",
    );

    const dy = fy + chartH + 20;
    ctx.fillStyle = "rgba(100, 200, 255, 0.4)";
    ctx.fillText("DIVERSITY", x + 2, dy + 8);
    this.renderSparkline(
      ctx,
      this.diversityHistory,
      x,
      dy + 12,
      w,
      chartH,
      "rgb(204, 102, 255)",
    );
  }

  renderHeatmap(ctx, worldW, worldH) {
    const cellW = worldW / this.heatmapCols;
    const cellH = worldH / this.heatmapRows;

    let maxDensity = 0;
    for (let i = 0; i < this.heatmap.length; i++) {
      if (this.heatmap[i] > maxDensity) maxDensity = this.heatmap[i];
    }
    if (maxDensity === 0) return;

    for (let row = 0; row < this.heatmapRows; row++) {
      for (let col = 0; col < this.heatmapCols; col++) {
        const density = this.heatmap[row * this.heatmapCols + col];
        if (density === 0) continue;
        const intensity = density / maxDensity;
        const alpha = intensity * 0.06;
        ctx.fillStyle = `rgba(0, 240, 255, ${alpha})`;
        ctx.fillRect(col * cellW, row * cellH, cellW + 1, cellH + 1);
      }
    }
  }

  toJSON() {
    return {
      generationCount: this.generationCount,
      totalBirths: this.totalBirths,
      totalDeaths: this.totalDeaths,
      totalTime: this.totalTime,
      peakPopulation: this.peakPopulation,
    };
  }

  static fromJSON(data) {
    const s = new StatsTracker();
    s.generationCount = data.generationCount || 0;
    s.totalBirths = data.totalBirths || 0;
    s.totalDeaths = data.totalDeaths || 0;
    s.totalTime = data.totalTime || 0;
    s.peakPopulation = data.peakPopulation || 0;
    return s;
  }
}