const setupSelect = document.getElementById("setupSelect");
const alphaInput = document.getElementById("alphaInput");
const zNormToggle = document.getElementById("zNormToggle");
const newQuestionBtn = document.getElementById("newQuestionBtn");
const solveBtn = document.getElementById("solveBtn");
const questionPrompt = document.getElementById("questionPrompt");
const solutionOutput = document.getElementById("solutionOutput");

const setupConfig = {
  two_weights_no_intercept: {
    label: "Two weights, no intercept",
    featureNames: ["x1", "x2"],
    paramNames: ["w1", "w2"],
    hasIntercept: false,
    formula: "y_hat = w1*x1 + w2*x2",
  },
  two_weights_one_intercept: {
    label: "Two weights + one intercept",
    featureNames: ["x1", "x2"],
    paramNames: ["b", "w1", "w2"],
    hasIntercept: true,
    formula: "y_hat = b + w1*x1 + w2*x2",
  },
  one_weight_one_intercept: {
    label: "One weight + one intercept",
    featureNames: ["x"],
    paramNames: ["b", "w"],
    hasIntercept: true,
    formula: "y_hat = b + w*x",
  },
};

const state = {
  setupKey: "two_weights_no_intercept",
  question: null,
};

function init() {
  bindEvents();
  generateQuestion();
}

function bindEvents() {
  setupSelect.addEventListener("change", () => {
    state.setupKey = setupSelect.value;
    generateQuestion();
    solutionOutput.innerHTML = "";
  });

  zNormToggle.addEventListener("change", renderQuestionPrompt);

  newQuestionBtn.addEventListener("click", () => {
    generateQuestion();
    solutionOutput.innerHTML = "";
  });

  solveBtn.addEventListener("click", () => {
    const alpha = clampNum(Number(alphaInput.value), 0.01, 1, 0.1);
    alphaInput.value = fmt(alpha);
    const useZNorm = zNormToggle.checked;
    renderSolution(alpha, useZNorm);
  });
}

function generateQuestion() {
  const cfg = setupConfig[state.setupKey];
  const rows = [];

  if (state.setupKey === "one_weight_one_intercept") {
    const trueB = randInt(-1, 3);
    const trueW = randInt(1, 4);
    const x1 = randInt(1, 4);
    let x2 = randInt(2, 6);
    while (x2 === x1) x2 = randInt(2, 6);

    rows.push({ x: [x1], y: trueB + trueW * x1 });
    rows.push({ x: [x2], y: trueB + trueW * x2 });

    state.question = {
      trueParams: { b: trueB, w: trueW },
      initParams: { b: 0, w: 0 },
      rows,
      cfg,
    };
  } else if (state.setupKey === "two_weights_no_intercept") {
    const trueW1 = randInt(1, 3);
    const trueW2 = randInt(1, 3);

    const x11 = randInt(1, 4);
    let x12 = randInt(2, 6);
    while (x12 === x11) x12 = randInt(2, 6);

    const x21 = randInt(1, 4);
    let x22 = randInt(2, 6);
    while (x22 === x21) x22 = randInt(2, 6);

    rows.push({ x: [x11, x21], y: trueW1 * x11 + trueW2 * x21 });
    rows.push({ x: [x12, x22], y: trueW1 * x12 + trueW2 * x22 });

    state.question = {
      trueParams: { w1: trueW1, w2: trueW2 },
      initParams: { w1: 0, w2: 0 },
      rows,
      cfg,
    };
  } else {
    const trueB = randInt(-1, 3);
    const trueW1 = randInt(1, 3);
    const trueW2 = randInt(1, 3);

    const x11 = randInt(1, 4);
    let x12 = randInt(2, 6);
    while (x12 === x11) x12 = randInt(2, 6);

    const x21 = randInt(1, 4);
    let x22 = randInt(2, 6);
    while (x22 === x21) x22 = randInt(2, 6);

    rows.push({ x: [x11, x21], y: trueB + trueW1 * x11 + trueW2 * x21 });
    rows.push({ x: [x12, x22], y: trueB + trueW1 * x12 + trueW2 * x22 });

    state.question = {
      trueParams: { b: trueB, w1: trueW1, w2: trueW2 },
      initParams: { b: 0, w1: 0, w2: 0 },
      rows,
      cfg,
    };
  }

  renderQuestionPrompt();
}

function renderQuestionPrompt() {
  if (!state.question) return;

  const useZNorm = zNormToggle.checked;
  const { cfg, rows } = state.question;

  questionPrompt.innerHTML = `
    <div class="matrix-card">
      <p class="matrix-title">Setup</p>
      <p class="formula-line">Model: ${escapeHtml(cfg.label)}</p>
      <p class="formula-line">Formula: ${escapeHtml(cfg.formula)}</p>
      <p class="formula-line">Initial parameters: ${renderParamInline(state.question.initParams)}</p>
      <p class="formula-line">Normalization for update step: ${useZNorm ? "z-normalization enabled" : "raw features"}</p>
      <p class="formula-line">Task: Compute gradients and one update step for all parameters.</p>
    </div>
    <div class="matrix-card">
      <p class="matrix-title">Two Data Samples</p>
      ${renderDataTable(rows, cfg.featureNames, null)}
      <p class="slice-label">Ground-truth y is given in the table.</p>
    </div>
  `;
}

function renderSolution(alpha, useZNorm) {
  if (!state.question) return;

  const { cfg, rows, initParams } = state.question;
  const m = rows.length;
  const featureStats = computeFeatureStats(rows);
  const rowsUsed = rows.map((row) => ({
    y: row.y,
    x: row.x.map((value, idx) => (useZNorm ? zNorm(value, featureStats[idx]) : value)),
    xRaw: row.x.slice(),
  }));

  const orderedParams = cfg.paramNames.slice();
  const grads = {};
  const update = {};
  const predRows = [];

  for (const name of orderedParams) grads[name] = 0;

  for (let i = 0; i < m; i += 1) {
    const row = rowsUsed[i];
    const yHat = predict(cfg, initParams, row.x);
    const err = yHat - row.y;

    predRows.push({ index: i + 1, x: row.x, xRaw: row.xRaw, y: row.y, yHat, err });

    if (cfg.hasIntercept) grads.b += err;

    if (state.setupKey === "one_weight_one_intercept") {
      grads.w += err * row.x[0];
    } else {
      grads.w1 += err * row.x[0];
      grads.w2 += err * row.x[1];
    }
  }

  for (const name of orderedParams) {
    grads[name] /= m;
    update[name] = initParams[name] - alpha * grads[name];
  }

  const cost = predRows.reduce((acc, r) => acc + r.err * r.err, 0) / (2 * m);

  const normBlock = useZNorm
    ? `
      <div class="matrix-card">
        <p class="matrix-title">Step 0: z-Normalization</p>
        ${renderNormSummary(cfg.featureNames, featureStats)}
        ${renderDataTable(rows, cfg.featureNames, rowsUsed.map((r) => r.x))}
      </div>
    `
    : "";

  solutionOutput.innerHTML = `
    ${normBlock}
    <div class="matrix-card">
      <p class="matrix-title">Step 1: Predictions and Errors</p>
      ${renderPredictionSteps(cfg, predRows, initParams, useZNorm)}
      <p class="formula-line">Cost: J = (1/(2m)) * sum((y_hat - y)^2) = ${fmt(cost)}</p>
    </div>

    <div class="matrix-card">
      <p class="matrix-title">Step 2: Gradient Calculation</p>
      ${renderGradientSteps(cfg, predRows, grads, m)}
    </div>

    <div class="matrix-card">
      <p class="matrix-title">Step 3: Parameter Update (One Gradient Step)</p>
      ${renderUpdateSteps(cfg, initParams, grads, update, alpha)}
    </div>

    <div class="matrix-card">
      <p class="matrix-title">Explanation</p>
      <p class="formula-line">A negative gradient means increasing that parameter will reduce cost, so the update adds value in that direction.</p>
      <p class="formula-line">A positive gradient means decreasing that parameter will reduce cost.</p>
      <p class="formula-line">${useZNorm ? "Using z-normalized features keeps scales consistent, so gradient magnitudes across features are easier to compare." : "Raw features are used directly, so larger-scale features can produce larger gradient terms."}</p>
      <p class="formula-line">Use <strong>Randomize New Samples</strong> to practice again with fresh numbers.</p>
    </div>
  `;
}

function renderDataTable(rows, featureNames, zRows) {
  const headers = featureNames.map((f) => `<th>${f}</th>`).join("");
  const zHeaders = zRows ? featureNames.map((f) => `<th>z(${f})</th>`).join("") : "";

  const body = rows
    .map((row, i) => {
      const xs = row.x.map((v) => `<td>${fmt(v)}</td>`).join("");
      const zs = zRows
        ? zRows[i].map((v) => `<td>${fmt(v)}</td>`).join("")
        : "";
      return `<tr><td>${i + 1}</td>${xs}${zs}<td>${fmt(row.y)}</td></tr>`;
    })
    .join("");

  return `
    <table class="matrix-table">
      <thead>
        <tr>
          <th>sample</th>
          ${headers}
          ${zHeaders}
          <th>y (ground truth)</th>
        </tr>
      </thead>
      <tbody>${body}</tbody>
    </table>
  `;
}

function renderNormSummary(featureNames, stats) {
  return featureNames
    .map((name, idx) => {
      const s = stats[idx];
      return `<p class="formula-line">${name}: mean=${fmt(s.mean)}, std=${fmt(s.std)} so z = (x - ${fmt(s.mean)}) / ${fmt(s.std)}</p>`;
    })
    .join("");
}

function renderPredictionSteps(cfg, predRows, initParams, useZNorm) {
  const pLine = renderParamInline(initParams);
  const head = `<p class="formula-line">Start with ${pLine}. ${useZNorm ? "Use normalized x values." : "Use raw x values."}</p>`;

  const lines = predRows
    .map((row) => {
      if (state.setupKey === "one_weight_one_intercept") {
        return `<p class="formula-line">sample ${row.index}: y_hat = b + w*x = ${fmt(initParams.b)} + ${fmt(initParams.w)}*${fmt(row.x[0])} = ${fmt(row.yHat)}, error = y_hat - y = ${fmt(row.yHat)} - ${fmt(row.y)} = ${fmt(row.err)}</p>`;
      }

      if (cfg.hasIntercept) {
        return `<p class="formula-line">sample ${row.index}: y_hat = b + w1*x1 + w2*x2 = ${fmt(initParams.b)} + ${fmt(initParams.w1)}*${fmt(row.x[0])} + ${fmt(initParams.w2)}*${fmt(row.x[1])} = ${fmt(row.yHat)}, error = ${fmt(row.err)}</p>`;
      }

      return `<p class="formula-line">sample ${row.index}: y_hat = w1*x1 + w2*x2 = ${fmt(initParams.w1)}*${fmt(row.x[0])} + ${fmt(initParams.w2)}*${fmt(row.x[1])} = ${fmt(row.yHat)}, error = ${fmt(row.err)}</p>`;
    })
    .join("");

  return head + lines;
}

function renderGradientSteps(cfg, predRows, grads, m) {
  const errTerms = predRows.map((r) => fmt(r.err)).join(" + ");
  let out = "";

  if (cfg.hasIntercept) {
    out += `<p class="formula-line">grad_b = (1/m) * sum(error) = (1/${m}) * (${errTerms}) = ${fmt(grads.b)}</p>`;
  }

  if (state.setupKey === "one_weight_one_intercept") {
    const terms = predRows.map((r) => `${fmt(r.err)}*${fmt(r.x[0])}`).join(" + ");
    out += `<p class="formula-line">grad_w = (1/m) * sum(error*x) = (1/${m}) * (${terms}) = ${fmt(grads.w)}</p>`;
    return out;
  }

  const terms1 = predRows.map((r) => `${fmt(r.err)}*${fmt(r.x[0])}`).join(" + ");
  const terms2 = predRows.map((r) => `${fmt(r.err)}*${fmt(r.x[1])}`).join(" + ");
  out += `<p class="formula-line">grad_w1 = (1/m) * sum(error*x1) = (1/${m}) * (${terms1}) = ${fmt(grads.w1)}</p>`;
  out += `<p class="formula-line">grad_w2 = (1/m) * sum(error*x2) = (1/${m}) * (${terms2}) = ${fmt(grads.w2)}</p>`;

  return out;
}

function renderUpdateSteps(cfg, initParams, grads, update, alpha) {
  return cfg.paramNames
    .map((name) => {
      return `<p class="formula-line">${name}_new = ${name} - alpha*grad_${name} = ${fmt(initParams[name])} - ${fmt(alpha)}*${fmt(grads[name])} = ${fmt(update[name])}</p>`;
    })
    .join("");
}

function predict(cfg, params, xRow) {
  let yHat = cfg.hasIntercept ? params.b : 0;

  if (state.setupKey === "one_weight_one_intercept") {
    yHat += params.w * xRow[0];
    return yHat;
  }

  yHat += params.w1 * xRow[0] + params.w2 * xRow[1];
  return yHat;
}

function computeFeatureStats(rows) {
  const d = rows[0].x.length;
  const stats = [];

  for (let j = 0; j < d; j += 1) {
    const values = rows.map((r) => r.x[j]);
    const mean = values.reduce((acc, v) => acc + v, 0) / values.length;
    const variance =
      values.reduce((acc, v) => acc + (v - mean) * (v - mean), 0) / values.length;
    const std = Math.sqrt(variance) || 1;
    stats.push({ mean, std });
  }

  return stats;
}

function zNorm(value, stat) {
  return (value - stat.mean) / stat.std;
}

function renderParamInline(params) {
  return Object.entries(params)
    .map(([k, v]) => `${k}=${fmt(v)}`)
    .join(", ");
}

function clampNum(v, min, max, fallback) {
  if (!Number.isFinite(v)) return fallback;
  return Math.max(min, Math.min(max, v));
}

function randInt(min, max) {
  return Math.floor(Math.random() * (max - min + 1)) + min;
}

function fmt(v) {
  if (!Number.isFinite(v)) return String(v);
  return Number(v.toFixed(4)).toString();
}

function escapeHtml(text) {
  return String(text)
    .replaceAll("&", "&amp;")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;")
    .replaceAll("'", "&#39;");
}

init();