const controls = {
  classMode: document.getElementById("classMode"),
  threshold: document.getElementById("threshold"),
  thresholdValue: document.getElementById("thresholdValue"),
  ruleText: document.getElementById("ruleText"),
  checkBtn: document.getElementById("checkBtn"),
  showSolutionBtn: document.getElementById("showSolutionBtn"),
  newSamplesBtn: document.getElementById("newSamplesBtn"),
  resetBtn: document.getElementById("resetBtn"),
};

const sampleArea = document.getElementById("sampleArea");
const matrixArea = document.getElementById("matrixArea");
const metricArea = document.getElementById("metricArea");
const feedbackStats = document.getElementById("feedbackStats");
const solutionBlock = document.getElementById("solutionBlock");
const solutionText = document.getElementById("solutionText");

let binarySamples = [];
let threeClassSamples = [];

function fmt(num, digits = 2) {
  return Number(num).toFixed(digits);
}

function round3(num) {
  return Number(num).toFixed(3);
}

function round2(num) {
  return Math.round(num * 100) / 100;
}

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

function generateBinarySamples() {
  const actualPattern = ["Positive", "Negative", "Positive", "Negative", "Positive", "Negative"];
  const samples = [];

  for (let i = 0; i < 6; i += 1) {
    const actual = actualPattern[i];
    const feature = round2(rand(0.2, 2.6));
    const pPos = actual === "Positive" ? round2(rand(0.38, 0.92)) : round2(rand(0.12, 0.78));
    samples.push({ id: `S${i + 1}`, feature, actual, pPos });
  }

  binarySamples = samples;
}

function generateThreeProbabilities(actual) {
  const labels = ["A", "B", "C"];
  const weights = {
    A: rand(0.1, 0.45),
    B: rand(0.1, 0.45),
    C: rand(0.1, 0.45),
  };

  weights[actual] += rand(0.18, 0.5);
  const total = labels.reduce((sum, label) => sum + weights[label], 0);

  let pA = round2(weights.A / total);
  let pB = round2(weights.B / total);
  if (pA + pB > 0.98) {
    const scale = 0.98 / (pA + pB);
    pA = round2(pA * scale);
    pB = round2(pB * scale);
  }
  const pC = round2(1 - pA - pB);

  return { pA, pB, pC };
}

function generateThreeClassSamples() {
  const actualPattern = ["A", "B", "C", "B", "C", "A"];
  const samples = [];

  for (let i = 0; i < 6; i += 1) {
    const actual = actualPattern[i];
    const feature = round2(rand(0.5, 2.4));
    const probs = generateThreeProbabilities(actual);
    samples.push({ id: `S${i + 1}`, feature, actual, ...probs });
  }

  threeClassSamples = samples;
}

function generateAllSamples() {
  generateBinarySamples();
  generateThreeClassSamples();
}

function computeBinaryPrediction(sample, threshold) {
  return sample.pPos >= threshold ? "Positive" : "Negative";
}

function computeThreePrediction(sample, threshold) {
  if (sample.pA >= threshold) return "A";
  if (sample.pB >= threshold) return "B";
  return "C";
}

function computeBinaryResults(threshold) {
  const rows = binarySamples.map((sample) => {
    const predicted = computeBinaryPrediction(sample, threshold);
    return { ...sample, predicted };
  });

  let tp = 0;
  let fp = 0;
  let tn = 0;
  let fn = 0;

  for (const row of rows) {
    if (row.actual === "Positive" && row.predicted === "Positive") tp += 1;
    if (row.actual === "Negative" && row.predicted === "Positive") fp += 1;
    if (row.actual === "Negative" && row.predicted === "Negative") tn += 1;
    if (row.actual === "Positive" && row.predicted === "Negative") fn += 1;
  }

  const precision = tp + fp === 0 ? 0 : tp / (tp + fp);
  const recall = tp + fn === 0 ? 0 : tp / (tp + fn);
  const sensitivity = recall;
  const specificity = tn + fp === 0 ? 0 : tn / (tn + fp);

  return {
    rows,
    matrix: {
      Positive: { Positive: tp, Negative: fn },
      Negative: { Positive: fp, Negative: tn },
    },
    metrics: { precision, recall, sensitivity, specificity },
  };
}

function getOneVsRestMetrics(matrix, className) {
  const classes = ["A", "B", "C"];
  const tp = matrix[className][className];

  let fp = 0;
  let fn = 0;
  for (const c of classes) {
    if (c !== className) {
      fp += matrix[c][className];
      fn += matrix[className][c];
    }
  }

  const total = classes.reduce((sum, r) => sum + classes.reduce((acc, c) => acc + matrix[r][c], 0), 0);
  const tn = total - tp - fp - fn;

  const precision = tp + fp === 0 ? 0 : tp / (tp + fp);
  const recall = tp + fn === 0 ? 0 : tp / (tp + fn);
  const specificity = tn + fp === 0 ? 0 : tn / (tn + fp);

  return { precision, recall, sensitivity: recall, specificity };
}

function computeThreeResults(threshold) {
  const rows = threeClassSamples.map((sample) => {
    const predicted = computeThreePrediction(sample, threshold);
    return { ...sample, predicted };
  });

  const classes = ["A", "B", "C"];
  const matrix = { A: { A: 0, B: 0, C: 0 }, B: { A: 0, B: 0, C: 0 }, C: { A: 0, B: 0, C: 0 } };

  for (const row of rows) {
    matrix[row.actual][row.predicted] += 1;
  }

  const perClass = classes.map((className) => getOneVsRestMetrics(matrix, className));
  const macro = {
    precision: perClass.reduce((s, m) => s + m.precision, 0) / classes.length,
    recall: perClass.reduce((s, m) => s + m.recall, 0) / classes.length,
    sensitivity: perClass.reduce((s, m) => s + m.sensitivity, 0) / classes.length,
    specificity: perClass.reduce((s, m) => s + m.specificity, 0) / classes.length,
  };

  return { rows, matrix, metrics: macro, perClass };
}

function getCurrentResults() {
  const mode = controls.classMode.value;
  const t = Number(controls.threshold.value);
  return mode === "binary" ? computeBinaryResults(t) : computeThreeResults(t);
}

function samplePredictionSelect(rowId, options) {
  return `<select data-role="sample-pred" data-id="${rowId}">${options
    .map((opt) => `<option value="${opt}">${opt}</option>`)
    .join("")}</select>`;
}

function renderSamples(results) {
  const mode = controls.classMode.value;
  const rows = results.rows;

  if (mode === "binary") {
    sampleArea.innerHTML = `
      <h3>1) Predict the class for each sample</h3>
      <table class="exercise-table">
        <thead>
          <tr>
            <th>Sample</th>
            <th>Feature x</th>
            <th>Actual Class</th>
            <th>p(Positive)</th>
            <th>Your Predicted Class</th>
          </tr>
        </thead>
        <tbody>
          ${rows
            .map((row, index) => {
              if (index < 2) {
                return `
                  <tr>
                    <td>${row.id}</td>
                    <td>${fmt(row.feature)}</td>
                    <td>${row.actual}</td>
                    <td>${fmt(row.pPos)}</td>
                    <td><strong>${row.predicted}</strong> <span class="worked-tag">(worked example)</span></td>
                  </tr>
                `;
              }

              return `
                <tr>
                  <td>${row.id}</td>
                  <td>${fmt(row.feature)}</td>
                  <td>${row.actual}</td>
                  <td>${fmt(row.pPos)}</td>
                  <td>${samplePredictionSelect(row.id, ["Positive", "Negative"])}</td>
                </tr>
              `;
            })
            .join("")}
        </tbody>
      </table>
    `;
    return;
  }

  sampleArea.innerHTML = `
    <h3>1) Predict the class for each sample</h3>
    <table class="exercise-table">
      <thead>
        <tr>
          <th>Sample</th>
          <th>Feature x</th>
          <th>Actual Class</th>
          <th>p(A)</th>
          <th>p(B)</th>
          <th>p(C)</th>
          <th>Your Predicted Class</th>
        </tr>
      </thead>
      <tbody>
        ${rows
          .map((row, index) => {
            if (index < 2) {
              return `
                <tr>
                  <td>${row.id}</td>
                  <td>${fmt(row.feature)}</td>
                  <td>${row.actual}</td>
                  <td>${fmt(row.pA)}</td>
                  <td>${fmt(row.pB)}</td>
                  <td>${fmt(row.pC)}</td>
                  <td><strong>${row.predicted}</strong> <span class="worked-tag">(worked example)</span></td>
                </tr>
              `;
            }

            return `
              <tr>
                <td>${row.id}</td>
                <td>${fmt(row.feature)}</td>
                <td>${row.actual}</td>
                <td>${fmt(row.pA)}</td>
                <td>${fmt(row.pB)}</td>
                <td>${fmt(row.pC)}</td>
                <td>${samplePredictionSelect(row.id, ["A", "B", "C"])}</td>
              </tr>
            `;
          })
          .join("")}
      </tbody>
    </table>
  `;
}

function renderMatrixInputs() {
  const mode = controls.classMode.value;

  if (mode === "binary") {
    matrixArea.innerHTML = `
      <h3>2) Fill the confusion matrix</h3>
      <p class="hint">Rows = Actual class, Columns = Predicted class.</p>
      <table class="exercise-table matrix-table">
        <thead>
          <tr><th>Actual \ Predicted</th><th>Positive</th><th>Negative</th></tr>
        </thead>
        <tbody>
          <tr>
            <th>Positive</th>
            <td><input type="number" min="0" data-role="matrix" data-key="Positive-Positive" /></td>
            <td><input type="number" min="0" data-role="matrix" data-key="Positive-Negative" /></td>
          </tr>
          <tr>
            <th>Negative</th>
            <td><input type="number" min="0" data-role="matrix" data-key="Negative-Positive" /></td>
            <td><input type="number" min="0" data-role="matrix" data-key="Negative-Negative" /></td>
          </tr>
        </tbody>
      </table>
    `;
    return;
  }

  matrixArea.innerHTML = `
    <h3>2) Fill the confusion matrix</h3>
    <p class="hint">Rows = Actual class, Columns = Predicted class.</p>
    <table class="exercise-table matrix-table">
      <thead>
        <tr><th>Actual \ Predicted</th><th>A</th><th>B</th><th>C</th></tr>
      </thead>
      <tbody>
        <tr>
          <th>A</th>
          <td><input type="number" min="0" data-role="matrix" data-key="A-A" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="A-B" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="A-C" /></td>
        </tr>
        <tr>
          <th>B</th>
          <td><input type="number" min="0" data-role="matrix" data-key="B-A" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="B-B" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="B-C" /></td>
        </tr>
        <tr>
          <th>C</th>
          <td><input type="number" min="0" data-role="matrix" data-key="C-A" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="C-B" /></td>
          <td><input type="number" min="0" data-role="matrix" data-key="C-C" /></td>
        </tr>
      </tbody>
    </table>
  `;
}

function renderMetricInputs() {
  const mode = controls.classMode.value;
  const detail = mode === "binary" ? "(for Positive class)" : "(macro-average across A, B, C)";
  metricArea.innerHTML = `
    <h3>3) Calculate metrics ${detail}</h3>
    <p class="hint">Enter decimal values (e.g., 0.667).</p>
    <table class="exercise-table metric-table">
      <thead>
        <tr><th>Metric</th><th>Your value</th></tr>
      </thead>
      <tbody>
        <tr><th>Precision</th><td><input type="number" step="0.001" data-role="metric" data-key="precision" /></td></tr>
        <tr><th>Recall</th><td><input type="number" step="0.001" data-role="metric" data-key="recall" /></td></tr>
        <tr><th>Sensitivity</th><td><input type="number" step="0.001" data-role="metric" data-key="sensitivity" /></td></tr>
        <tr><th>Specificity</th><td><input type="number" step="0.001" data-role="metric" data-key="specificity" /></td></tr>
      </tbody>
    </table>
  `;
}

function updateRuleText() {
  const mode = controls.classMode.value;
  const threshold = Number(controls.threshold.value);

  controls.thresholdValue.textContent = fmt(threshold);

  if (mode === "binary") {
    controls.ruleText.textContent =
      "Rule: predict Positive if p(Positive) >= t, otherwise predict Negative. (Default t = 0.50)";
  } else {
    controls.ruleText.textContent =
      "Rule: check A first, then B. If p(A) >= t predict A; else if p(B) >= t predict B; otherwise predict C.";
  }
}

function renderPage() {
  const results = getCurrentResults();
  updateRuleText();
  renderSamples(results);
  renderMatrixInputs();
  renderMetricInputs();
  feedbackStats.innerHTML = "";
}

function readUserSampleAnswers() {
  const selects = [...document.querySelectorAll('select[data-role="sample-pred"]')];
  const map = {};
  for (const select of selects) {
    map[select.dataset.id] = select.value;
  }
  return map;
}

function readUserMatrix() {
  const inputs = [...document.querySelectorAll('input[data-role="matrix"]')];
  const matrix = {};
  for (const input of inputs) {
    matrix[input.dataset.key] = input.value === "" ? NaN : Number(input.value);
  }
  return matrix;
}

function readUserMetrics() {
  const inputs = [...document.querySelectorAll('input[data-role="metric"]')];
  const metrics = {};
  for (const input of inputs) {
    metrics[input.dataset.key] = input.value === "" ? NaN : Number(input.value);
  }
  return metrics;
}

function expectedMatrixFlat(results) {
  const mode = controls.classMode.value;
  if (mode === "binary") {
    return {
      "Positive-Positive": results.matrix.Positive.Positive,
      "Positive-Negative": results.matrix.Positive.Negative,
      "Negative-Positive": results.matrix.Negative.Positive,
      "Negative-Negative": results.matrix.Negative.Negative,
    };
  }

  return {
    "A-A": results.matrix.A.A,
    "A-B": results.matrix.A.B,
    "A-C": results.matrix.A.C,
    "B-A": results.matrix.B.A,
    "B-B": results.matrix.B.B,
    "B-C": results.matrix.B.C,
    "C-A": results.matrix.C.A,
    "C-B": results.matrix.C.B,
    "C-C": results.matrix.C.C,
  };
}

function approximatelyEqual(a, b, eps = 0.02) {
  return Math.abs(a - b) <= eps;
}

function checkAnswers() {
  const results = getCurrentResults();

  const expectedPreds = {};
  for (const row of results.rows.slice(2)) {
    expectedPreds[row.id] = row.predicted;
  }

  const userPreds = readUserSampleAnswers();
  let predCorrect = 0;
  const predTotal = Object.keys(expectedPreds).length;
  for (const id of Object.keys(expectedPreds)) {
    if (userPreds[id] === expectedPreds[id]) {
      predCorrect += 1;
    }
  }

  const userMatrix = readUserMatrix();
  const expectedMatrix = expectedMatrixFlat(results);
  let matrixCorrect = 0;
  const matrixTotal = Object.keys(expectedMatrix).length;
  for (const key of Object.keys(expectedMatrix)) {
    if (userMatrix[key] === expectedMatrix[key]) {
      matrixCorrect += 1;
    }
  }

  const userMetrics = readUserMetrics();
  const expectedMetrics = results.metrics;
  let metricCorrect = 0;
  const metricTotal = 4;
  for (const metricName of ["precision", "recall", "sensitivity", "specificity"]) {
    if (approximatelyEqual(userMetrics[metricName], expectedMetrics[metricName])) {
      metricCorrect += 1;
    }
  }

  feedbackStats.innerHTML = `
    <span>Prediction check: ${predCorrect}/${predTotal} correct</span>
    <span>Confusion matrix check: ${matrixCorrect}/${matrixTotal} correct</span>
    <span>Metric check: ${metricCorrect}/${metricTotal} correct (tolerance +/- 0.02)</span>
    <span>Expected metrics: Precision=${round3(expectedMetrics.precision)}, Recall=${round3(
      expectedMetrics.recall,
    )}, Sensitivity=${round3(expectedMetrics.sensitivity)}, Specificity=${round3(expectedMetrics.specificity)}</span>
  `;
}

function renderBinarySolution(results) {
  const steps = results.rows
    .map((row) => {
      return `<li>${row.id}: p(Positive)=${fmt(row.pPos)}; compare with t=${fmt(
        controls.threshold.value,
      )} => predicted <strong>${row.predicted}</strong>; actual ${row.actual}.</li>`;
    })
    .join("");

  solutionText.innerHTML = `
    <h4>Step A: Classify all 6 samples</h4>
    <ol>${steps}</ol>

    <h4>Step B: Build confusion matrix (rows=actual, columns=predicted)</h4>
    <p>TP=${results.matrix.Positive.Positive}, FN=${results.matrix.Positive.Negative}, FP=${results.matrix.Negative.Positive}, TN=${results.matrix.Negative.Negative}</p>

    <table class="exercise-table matrix-table">
      <thead><tr><th>Actual \\ Predicted</th><th>Positive</th><th>Negative</th></tr></thead>
      <tbody>
        <tr><th>Positive</th><td>${results.matrix.Positive.Positive}</td><td>${results.matrix.Positive.Negative}</td></tr>
        <tr><th>Negative</th><td>${results.matrix.Negative.Positive}</td><td>${results.matrix.Negative.Negative}</td></tr>
      </tbody>
    </table>

    <h4>Step C: Metrics</h4>
    <p>Precision = TP/(TP+FP) = ${results.matrix.Positive.Positive}/(${results.matrix.Positive.Positive}+${results.matrix.Negative.Positive}) = ${round3(
      results.metrics.precision,
    )}</p>
    <p>Recall = TP/(TP+FN) = ${results.matrix.Positive.Positive}/(${results.matrix.Positive.Positive}+${results.matrix.Positive.Negative}) = ${round3(
      results.metrics.recall,
    )}</p>
    <p>Sensitivity = Recall = ${round3(results.metrics.sensitivity)}</p>
    <p>Specificity = TN/(TN+FP) = ${results.matrix.Negative.Negative}/(${results.matrix.Negative.Negative}+${results.matrix.Negative.Positive}) = ${round3(
      results.metrics.specificity,
    )}</p>
  `;
}

function renderThreeSolution(results) {
  const threshold = Number(controls.threshold.value);
  const steps = results.rows
    .map((row) => {
      return `<li>${row.id}: p(A)=${fmt(row.pA)}, p(B)=${fmt(row.pB)}, p(C)=${fmt(
        row.pC,
      )}. Since t=${fmt(threshold)}, prediction is <strong>${row.predicted}</strong>; actual ${row.actual}.</li>`;
    })
    .join("");

  const classes = ["A", "B", "C"];
  const perClass = classes
    .map((c, idx) => {
      return `<li>Class ${c}: Precision=${round3(results.perClass[idx].precision)}, Recall/Sensitivity=${round3(
        results.perClass[idx].recall,
      )}, Specificity=${round3(results.perClass[idx].specificity)}</li>`;
    })
    .join("");

  solutionText.innerHTML = `
    <h4>Step A: Classify all 6 samples</h4>
    <ol>${steps}</ol>

    <h4>Step B: Build 3x3 confusion matrix (rows=actual, columns=predicted)</h4>
    <table class="exercise-table matrix-table">
      <thead><tr><th>Actual \\ Predicted</th><th>A</th><th>B</th><th>C</th></tr></thead>
      <tbody>
        <tr><th>A</th><td>${results.matrix.A.A}</td><td>${results.matrix.A.B}</td><td>${results.matrix.A.C}</td></tr>
        <tr><th>B</th><td>${results.matrix.B.A}</td><td>${results.matrix.B.B}</td><td>${results.matrix.B.C}</td></tr>
        <tr><th>C</th><td>${results.matrix.C.A}</td><td>${results.matrix.C.B}</td><td>${results.matrix.C.C}</td></tr>
      </tbody>
    </table>

    <h4>Step C: One-vs-rest metrics for each class, then macro-average</h4>
    <ol>${perClass}</ol>

    <p><strong>Macro Precision</strong> = ${round3(results.metrics.precision)}</p>
    <p><strong>Macro Recall</strong> = ${round3(results.metrics.recall)}</p>
    <p><strong>Macro Sensitivity</strong> = ${round3(results.metrics.sensitivity)}</p>
    <p><strong>Macro Specificity</strong> = ${round3(results.metrics.specificity)}</p>
  `;
}

function showSolution() {
  const results = getCurrentResults();
  if (controls.classMode.value === "binary") {
    renderBinarySolution(results);
  } else {
    renderThreeSolution(results);
  }
  solutionBlock.hidden = false;
}

function resetInputs() {
  renderPage();
  solutionBlock.hidden = true;
}

function regenerateSamples() {
  generateAllSamples();
  resetInputs();
}

controls.classMode.addEventListener("change", resetInputs);
controls.threshold.addEventListener("input", resetInputs);
controls.checkBtn.addEventListener("click", checkAnswers);
controls.showSolutionBtn.addEventListener("click", showSolution);
controls.newSamplesBtn.addEventListener("click", regenerateSamples);
controls.resetBtn.addEventListener("click", resetInputs);

generateAllSamples();
renderPage();