const canvas = document.getElementById("plot");
const ctx = canvas.getContext("2d");

const controls = {
  b0: document.getElementById("b0"),
  b1: document.getElementById("b1"),
  t: document.getElementById("t"),
  probeX: document.getElementById("probeX"),
  showGrid: document.getElementById("showGrid"),
  showShade: document.getElementById("showShade"),
  showDerivative: document.getElementById("showDerivative"),
  preset: document.getElementById("preset"),
  resetBtn: document.getElementById("resetBtn"),
  animateBtn: document.getElementById("animateBtn"),
};

const valueLabels = {
  b0: document.getElementById("b0Value"),
  b1: document.getElementById("b1Value"),
  t: document.getElementById("tValue"),
  probeX: document.getElementById("probeValue"),
};

const stats = document.getElementById("stats");
const points = [];

let animationTimer = null;
let animationDirection = 1;

const MARGIN = { left: 70, right: 24, top: 24, bottom: 58 };
const X_MIN = -10;
const X_MAX = 10;

function logistic(z) {
  return 1 / (1 + Math.exp(-z));
}

function modelProbability(x, b0, b1) {
  return logistic(b0 + b1 * x);
}

function xToCanvas(x) {
  const w = canvas.width - MARGIN.left - MARGIN.right;
  return MARGIN.left + ((x - X_MIN) / (X_MAX - X_MIN)) * w;
}

function yToCanvas(y) {
  const h = canvas.height - MARGIN.top - MARGIN.bottom;
  return canvas.height - MARGIN.bottom - y * h;
}

function canvasToX(pixelX) {
  const w = canvas.width - MARGIN.left - MARGIN.right;
  return X_MIN + ((pixelX - MARGIN.left) / w) * (X_MAX - X_MIN);
}

function format(num, digits = 3) {
  return Number(num).toFixed(digits);
}

function drawAxes(showGrid) {
  const plotLeft = MARGIN.left;
  const plotRight = canvas.width - MARGIN.right;
  const plotTop = MARGIN.top;
  const plotBottom = canvas.height - MARGIN.bottom;

  ctx.strokeStyle = "#c9d9df";
  ctx.lineWidth = 1;

  if (showGrid) {
    for (let x = -10; x <= 10; x += 2) {
      const px = xToCanvas(x);
      ctx.beginPath();
      ctx.moveTo(px, plotTop);
      ctx.lineTo(px, plotBottom);
      ctx.stroke();
    }
    for (let y = 0; y <= 1.001; y += 0.1) {
      const py = yToCanvas(y);
      ctx.beginPath();
      ctx.moveTo(plotLeft, py);
      ctx.lineTo(plotRight, py);
      ctx.stroke();
    }
  }

  ctx.strokeStyle = "#14343f";
  ctx.lineWidth = 1.4;

  ctx.beginPath();
  ctx.moveTo(plotLeft, yToCanvas(0));
  ctx.lineTo(plotRight, yToCanvas(0));
  ctx.stroke();

  ctx.beginPath();
  ctx.moveTo(xToCanvas(0), plotTop);
  ctx.lineTo(xToCanvas(0), plotBottom);
  ctx.stroke();

  ctx.fillStyle = "#14343f";
  ctx.font = "13px 'Avenir Next', sans-serif";

  for (let x = -10; x <= 10; x += 2) {
    const px = xToCanvas(x);
    ctx.fillText(String(x), px - 8, yToCanvas(0) + 19);
  }

  for (let y = 0; y <= 1.001; y += 0.2) {
    const py = yToCanvas(y);
    ctx.fillText(format(y, 1), plotLeft - 40, py + 4);
  }

  ctx.fillText("x", plotRight - 12, yToCanvas(0) + 40);
  ctx.fillText("p", xToCanvas(0) + 12, plotTop + 12);
}

function drawThresholdLine(t) {
  ctx.strokeStyle = "#d8534f";
  ctx.lineWidth = 1.3;
  ctx.setLineDash([7, 5]);
  ctx.beginPath();
  ctx.moveTo(MARGIN.left, yToCanvas(t));
  ctx.lineTo(canvas.width - MARGIN.right, yToCanvas(t));
  ctx.stroke();
  ctx.setLineDash([]);
}

function drawRegionShade(b0, b1, t) {
  if (Math.abs(b1) < 1e-9) {
    return;
  }
  const logit = Math.log(t / (1 - t));
  const boundaryX = (logit - b0) / b1;

  const left = MARGIN.left;
  const right = canvas.width - MARGIN.right;
  const top = MARGIN.top;
  const bottom = canvas.height - MARGIN.bottom;
  const boundaryPx = xToCanvas(Math.max(X_MIN, Math.min(X_MAX, boundaryX)));

  ctx.globalAlpha = 0.15;
  ctx.fillStyle = "#0a8f7b";

  if (b1 > 0) {
    ctx.fillRect(boundaryPx, top, right - boundaryPx, bottom - top);
  } else {
    ctx.fillRect(left, top, boundaryPx - left, bottom - top);
  }
  ctx.globalAlpha = 1;
}

function drawCurve(b0, b1) {
  ctx.strokeStyle = "#0a8f7b";
  ctx.lineWidth = 3;
  ctx.beginPath();

  const steps = 600;
  for (let i = 0; i <= steps; i += 1) {
    const x = X_MIN + (i / steps) * (X_MAX - X_MIN);
    const y = modelProbability(x, b0, b1);
    const px = xToCanvas(x);
    const py = yToCanvas(y);

    if (i === 0) {
      ctx.moveTo(px, py);
    } else {
      ctx.lineTo(px, py);
    }
  }

  ctx.stroke();
}

function drawDerivative(b0, b1) {
  ctx.strokeStyle = "#f2b84b";
  ctx.lineWidth = 2;
  ctx.setLineDash([6, 4]);
  ctx.beginPath();

  const steps = 600;
  for (let i = 0; i <= steps; i += 1) {
    const x = X_MIN + (i / steps) * (X_MAX - X_MIN);
    const p = modelProbability(x, b0, b1);
    const derivative = Math.abs(b1 * p * (1 - p));
    const y = Math.min(1, derivative * 4);

    const px = xToCanvas(x);
    const py = yToCanvas(y);

    if (i === 0) {
      ctx.moveTo(px, py);
    } else {
      ctx.lineTo(px, py);
    }
  }

  ctx.stroke();
  ctx.setLineDash([]);
}

function drawProbeAndSamples(b0, b1, t, probeX) {
  const probeP = modelProbability(probeX, b0, b1);

  for (const sampleX of points) {
    const p = modelProbability(sampleX, b0, b1);
    const positive = p >= t;
    ctx.fillStyle = positive ? "#0a8f7b" : "#d8534f";
    ctx.beginPath();
    ctx.arc(xToCanvas(sampleX), yToCanvas(p), 5, 0, Math.PI * 2);
    ctx.fill();
  }

  ctx.strokeStyle = "#14343f";
  ctx.setLineDash([4, 4]);
  ctx.beginPath();
  ctx.moveTo(xToCanvas(probeX), yToCanvas(0));
  ctx.lineTo(xToCanvas(probeX), yToCanvas(probeP));
  ctx.stroke();
  ctx.setLineDash([]);

  ctx.fillStyle = "#14343f";
  ctx.beginPath();
  ctx.arc(xToCanvas(probeX), yToCanvas(probeP), 6, 0, Math.PI * 2);
  ctx.fill();

  ctx.font = "12px 'Avenir Next', sans-serif";
  ctx.fillText(`probe p=${format(probeP)}`, xToCanvas(probeX) + 8, yToCanvas(probeP) - 8);
}

function renderStats(b0, b1, t, probeX) {
  const probeP = modelProbability(probeX, b0, b1);

  let boundaryText = "No finite boundary (b1≈0)";
  if (Math.abs(b1) > 1e-9) {
    const logit = Math.log(t / (1 - t));
    const boundaryX = (logit - b0) / b1;
    boundaryText = `Decision boundary x*: ${format(boundaryX)}`;
  }

  const positives = points.filter((x) => modelProbability(x, b0, b1) >= t).length;
  const negatives = points.length - positives;

  stats.innerHTML = `
    <span>Equation: p = 1/(1+exp(-(${format(b0)} + ${format(b1)}x)))</span>
    <span>Threshold t = ${format(t, 2)}</span>
    <span>Probe: x = ${format(probeX, 2)}, p = ${format(probeP)}</span>
    <span>${boundaryText}</span>
    <span>Sample points: ${points.length} total</span>
    <span>Class counts by t: ${positives} positive, ${negatives} negative</span>
  `;
}

function updateLabelValues() {
  valueLabels.b0.textContent = format(controls.b0.value, 2);
  valueLabels.b1.textContent = format(controls.b1.value, 2);
  valueLabels.t.textContent = format(controls.t.value, 2);
  valueLabels.probeX.textContent = format(controls.probeX.value, 2);
}

function render() {
  const b0 = Number(controls.b0.value);
  const b1 = Number(controls.b1.value);
  const t = Number(controls.t.value);
  const probeX = Number(controls.probeX.value);

  updateLabelValues();

  ctx.clearRect(0, 0, canvas.width, canvas.height);
  drawAxes(controls.showGrid.checked);

  if (controls.showShade.checked) {
    drawRegionShade(b0, b1, t);
  }

  drawThresholdLine(t);
  drawCurve(b0, b1);

  if (controls.showDerivative.checked) {
    drawDerivative(b0, b1);
  }

  drawProbeAndSamples(b0, b1, t, probeX);
  renderStats(b0, b1, t, probeX);
}

function reset() {
  controls.b0.value = 0;
  controls.b1.value = 1;
  controls.t.value = 0.5;
  controls.probeX.value = 0;
  controls.showGrid.checked = true;
  controls.showShade.checked = true;
  controls.showDerivative.checked = false;
  controls.preset.value = "default";
  points.length = 0;
  stopAnimation();
  render();
}

function applyPreset(presetKey) {
  const presetMap = {
    default: { b0: 0, b1: 1, t: 0.5 },
    steep: { b0: -1, b1: 3, t: 0.5 },
    reversed: { b0: 0.5, b1: -2.2, t: 0.5 },
    "high-threshold": { b0: 0, b1: 1, t: 0.8 },
  };

  const preset = presetMap[presetKey] || presetMap.default;
  controls.b0.value = preset.b0;
  controls.b1.value = preset.b1;
  controls.t.value = preset.t;
  render();
}

function startAnimation() {
  if (animationTimer) return;

  controls.animateBtn.textContent = "Stop Animation";

  animationTimer = setInterval(() => {
    const current = Number(controls.probeX.value);
    let next = current + animationDirection * 0.18;

    if (next > X_MAX) {
      next = X_MAX;
      animationDirection = -1;
    }
    if (next < X_MIN) {
      next = X_MIN;
      animationDirection = 1;
    }

    controls.probeX.value = next;
    render();
  }, 30);
}

function stopAnimation() {
  if (animationTimer) {
    clearInterval(animationTimer);
    animationTimer = null;
  }
  controls.animateBtn.textContent = "Animate Probe";
}

for (const key of ["b0", "b1", "t", "probeX", "showGrid", "showShade", "showDerivative"]) {
  controls[key].addEventListener("input", render);
}

controls.resetBtn.addEventListener("click", reset);

controls.preset.addEventListener("change", (event) => {
  applyPreset(event.target.value);
});

controls.animateBtn.addEventListener("click", () => {
  if (animationTimer) {
    stopAnimation();
  } else {
    startAnimation();
  }
});

canvas.addEventListener("click", (event) => {
  const rect = canvas.getBoundingClientRect();
  const clickX = event.clientX - rect.left;
  const x = canvasToX((clickX / rect.width) * canvas.width);

  if (x < X_MIN || x > X_MAX) {
    return;
  }

  points.push(Number(x.toFixed(2)));
  if (points.length > 30) {
    points.shift();
  }

  render();
});

reset();