/**
 * ### Description
 * Calculate cosine similarity from 2 arrays.
 *
 * @param {Array} array1 1-dimensional array.
 * @param {Array} array2 1-dimensional array.
 * @return {Number} Calculated result of cosine similarity.
 */
function cosineSimilarity_(array1, array2) {
  const dotProduct = array1.reduce((t, e, i) => (t += e * array2[i]), 0);
  const magnitudes = [array1, array2]
    .map((e) => Math.sqrt(e.reduce((t, f) => (t += f * f), 0)))
    .reduce((t, f) => (t *= f), 1);
  return dotProduct / magnitudes;
}

/**
 * ### Description
 * Request "Method: models.embedContent" of Gemini Pro API.
 *
 * @param {Object} requests Request body for Method: models.batchEmbedContents.
 * @return {Number} Embedding generated from the input texts.
 */
function getEmbedding_(requests) {
  const apiKey = "###"; // Please set your API key.

  const url = `https://generativelanguage.googleapis.com/v1/models/embedding-001:batchEmbedContents?key=${apiKey}`;
  const options = {
    payload: JSON.stringify({ requests }),
    contentType: "application/json",
  };
  const res = UrlFetchApp.fetch(url, options);
  return (obj = JSON.parse(res.getContentText()));
}

function sample1() {
  // This is a search text.
  const searchText = "car";

  // These sample texts are generated by Gemini Pro API.
  const texts = [
    {
      title: "car",
      text: "A machine with four wheels that is powered by an engine and is used to transport people or goods. It has a steering wheel, a braking system, and an acceleration system.",
    },
    {
      title: "plane",
      text: "A machine that flies through the air, powered by engines. It has wings that generate lift, and a tail that provides stability. It can carry people or cargo.",
    },
    {
      title: "dog",
      text: "A four-legged mammal that is domesticated and used as a companion or working animal. It has a keen sense of smell and hearing, and is often used for hunting, herding, or guarding.",
    },
    {
      title: "cat",
      text: "A small, furry mammal that is often kept as a pet. It is known for its independent nature, sharp claws, and hunting skills.",
    },
    {
      title: "bird",
      text: "A feathered creature that flies through the air. It has a beak, wings, and feathers. It builds nests and lays eggs.",
    },
    {
      title: "fish",
      text: "A cold-blooded vertebrate that lives in water. It has gills for breathing, fins for swimming, and scales for protection.",
    },
    {
      title: "human",
      text: "A bipedal, intelligent animal that is capable of complex thought and emotion. It is the only species on Earth that can use language, art, and technology.",
    },
    {
      title: "plant",
      text: "An organism that is not an animal and does not move. It has cells with cell walls and uses photosynthesis to make its own food.",
    },
    {
      title: "cherry blossoms",
      text: "A tree that blooms with delicate, colorful flowers in the spring. It is a symbol of beauty, renewal, and hope in many cultures.",
    },
  ];

  const r1 = {
    model: "models/embedding-001",
    content: { parts: [{ text: searchText }] },
    taskType: "RETRIEVAL_QUERY",
  };
  const r2 = texts.map(({ title, text }) => ({
    model: "models/embedding-001",
    content: { parts: [{ text }] },
    taskType: "RETRIEVAL_DOCUMENT",
    title,
  }));
  const obj = getEmbedding_([r1, ...r2]);
  const [ar1, ...ar2] = obj.embeddings.map((e) => e.values);
  const cs = ar2
    .map((ar, i) => ({ cs: cosineSimilarity_(ar1, ar), text: texts[i] }))
    .sort((a, b) => (a.cs < b.cs ? 1 : -1));
  const result = cs[0].text;

  console.log(result);
}

function sample2() {
  // This is a search text.
  const searchText = "dog";

  // Please set your document ID.
  const documentId = "###";

  const body = DocumentApp.openById(documentId).getBody();
  const texts = body.getParagraphs().reduce((ar, p) => {
    const text = p.getText().trim();
    if (text) {
      ar.push(text);
    }
    return ar;
  }, []);
  const r1 = {
    model: "models/embedding-001",
    content: { parts: [{ text: searchText }] },
    taskType: "RETRIEVAL_QUERY",
  };
  const r2 = texts.map((text) => ({
    model: "models/embedding-001",
    content: { parts: [{ text }] },
    taskType: "RETRIEVAL_DOCUMENT",
  }));
  const obj = getEmbedding_([r1, ...r2]);
  const [ar1, ...ar2] = obj.embeddings.map((e) => e.values);
  const cs = ar2
    .map((ar, i) => ({ cs: cosineSimilarity_(ar1, ar), text: texts[i] }))
    .sort((a, b) => (a.cs < b.cs ? 1 : -1));
  const result = cs[0].text;
  body.findText(result).getElement().asText().setBackgroundColor("#f4cccc");
}

// Please set the folder ID of folder including image files.
const folderId = "###";

// First, please run this function. This function set the explanation of the image to the description of the file.
function sample_C1() {
  const apiKey = "###"; // Please set your API key.

  const files = DriveApp.getFolderById(folderId).getFiles();
  const token = ScriptApp.getOAuthToken();
  while (files.hasNext()) {
    const file = files.next();
    const bytes = UrlFetchApp.fetch(`https://drive.google.com/thumbnail?sz=w1000&id=${file.getId()}`, { headers: { authorization: "Bearer " + token } }).getContent();
    const base64 = Utilities.base64Encode(bytes);
    const url = `https://generativelanguage.googleapis.com/v1/models/gemini-pro-vision:generateContent?key=${apiKey}`;
    const payload = { contents: [{ parts: [{ text: "What is this picture?" }, { inline_data: { mime_type: "image/png", data: base64 } }] }] };
    const options = { payload: JSON.stringify(payload), contentType: "application/json" };
    const res = UrlFetchApp.fetch(url, options);
    const obj = JSON.parse(res.getContentText());
    const r = obj.candidates[0].content.parts[0].text;
    console.log(r);
    file.setDescription(r.trim());
  }
}

// After you ran "sample_C1", please run this function.
function sample_C2() {
  // This is a search text.
  const searchText = "dog";

  const descriptions = [];
  const files = DriveApp.getFolderById(folderId).getFiles();
  while (files.hasNext()) {
    const file = files.next();
    const description = file.getDescription().trim();
    if (description) {
      descriptions.push({ filename: file.getName(), fileId: file.getId(), text: description, request: { model: "models/embedding-001", content: { parts: [{ text: description }] }, taskType: "RETRIEVAL_DOCUMENT" } });
    }
  }
  const r1 = { model: "models/embedding-001", content: { parts: [{ text: searchText }] }, taskType: "RETRIEVAL_QUERY" };
  const obj = getEmbedding_([r1, ...descriptions.map(({ request }) => request)]);
  const [ar1, ...ar2] = obj.embeddings.map(e => e.values);
  const cs = ar2.map((ar, i) => ({ cs: cosineSimilarity_(ar1, ar), file: descriptions[i] })).sort((a, b) => a.cs < b.cs ? 1 : -1);
  const result = cs[0].file.filename;
  console.log(result);
}