const textarea = document.querySelector("textarea");
const classificationResult = document.querySelector(".classification-result");
let lastModel;
let lastClassName;
let warmedUp = false;

setupButton(
  "tfjs-webgl",
  async () =>
    await tfTask.SentimentDetection.Toxicity.TFJS.load({
      backend: "webgl",
    }),
  true
);

setupButton(
  "tflite",
  async () => await tfTask.SentimentDetection.MovieReview.TFLite.load()
);

async function setupButton(className, modelCreateFn, needWarmup) {
  document
    .querySelector(`.model.${className} .btn`)
    .classList.remove("disabled");
  const resultEle = document.querySelector(`.model.${className} .result`);
  document
    .querySelector(`.model.${className} .btn`)
    .addEventListener("click", async () => {
      let model;
      // Create the model when user clicks on a button.
      if (lastClassName !== className) {
        // Clean up the previous model if existed.
        if (lastModel) {
          lastModel.cleanUp();
        }
        // Create the new model and save it.
        resultEle.textContent = "Loading...";
        model = await modelCreateFn();
        lastModel = model;
        lastClassName = className;
      }
      // Reuse the model if user clicks on the same button.
      else {
        model = lastModel;
      }

      // Warm up if needed.
      if (needWarmup && !warmedUp) {
        await model.predict(textarea.value);
        warmedUp = true;
      }

      // Run inference and update result.
      const start = Date.now();
      const result = await model.predict(textarea.value);
      const latency = Date.now() - start;
      renderResult(result);
      resultEle.textContent = `Latency: ${latency}ms`;
    });
}

function renderResult(result) {
  console.log(result);
  let strResult = "";
  const labels = Object.keys(result.sentimentLabels).sort();
  for (const label of labels) {
    const classificationResult = result.sentimentLabels[label].result;
    const probabilities = result.sentimentLabels[label].probabilities;
    strResult += `<span class="label">${label}</span>`;
    strResult += `<span class="sentiment ${classificationResult}">${classificationResult}</span><br>`;
  }
  classificationResult.innerHTML = strResult;
}