const fs = require('fs');
const path = require('path');

const { PredictionServiceClient } = require('@google-cloud/aiplatform').v1;
const { helpers } = require('@google-cloud/aiplatform');

const KMeans = require('ml-kmeans');
const PCA = require('ml-pca');

async function loadPainPoints() {
  console.log("Loading pain points...");
  const data = fs.readFileSync('pain_points_results.json', 'utf8');
  const results = JSON.parse(data);

  const painPoints = [];

  for (const demographic in results) {
    for (const environment in results[demographic]) {
      const points = results[demographic][environment].pain_points;
      points.forEach(point => {
        painPoints.push({
          demographic,
          environment,
          point
        });
      });
    }
  }
  console.log(`Total pain points loaded: ${painPoints.length}`);
  return painPoints;
}

async function getEmbeddings(painPoints) {
  console.log("Generating embeddings...");
  const clientOptions = { apiEndpoint: 'us-central1-aiplatform.googleapis.com' };
  const client = new PredictionServiceClient(clientOptions);

  const project = 'project-id'; // TODO: Replace with your Google Cloud project ID
  const location = 'us-central1';
  const model = 'text-embedding-004	; // You can choose a different model if needed

  const endpoint = `projects/${project}/locations/${location}/publishers/google/models/${model}`;

  const batchSize = 100; // Adjust based on API limits
  const embeddings = [];

  for (let i = 0; i < painPoints.length; i += batchSize) {
    const batchPoints = painPoints.slice(i, i + batchSize);

    const instances = batchPoints.map(p => ({ content: p.point }));

    const request = {
      endpoint: endpoint,
      instances: instances.map(e => helpers.toValue(e)),
    };

    try {
      const [response] = await client.predict(request);

      const predictions = response.predictions;

      predictions.forEach(prediction => {
        const embeddingsProto = prediction.structValue.fields.embeddings;
        const valuesProto = embeddingsProto.structValue.fields.values;
        const embeddingValues = valuesProto.listValue.values.map(v => v.numberValue);
        embeddings.push(embeddingValues);
      });

      console.log(`Processed ${Math.min(i + batchSize, painPoints.length)} / ${painPoints.length} pain points`);
    } catch (error) {
      console.error(`Error processing batch starting at index ${i}: ${error.message}`);
      // Handle specific errors if needed
    }
  }

  console.log(`Total embeddings generated: ${embeddings.length}`);
  return embeddings;
}

function clusterEmbeddings(embeddings, numClusters = 100) {
  console.log("Clustering embeddings...");
  const kmeansResult = KMeans(embeddings, numClusters);
  return kmeansResult.clusters;
}

function reduceDimensions(embeddings) {
  console.log("Reducing dimensions for visualization...");
  const pca = new PCA(embeddings);
  const reducedData = pca.predict(embeddings, { nComponents: 2 }).to2DArray();
  return reducedData;
}

function visualizeClusters(reducedData, clusters) {
  console.log("Visualizing clusters...");

  const data = [];
  const uniqueClusters = Array.from(new Set(clusters));

  uniqueClusters.forEach(clusterId => {
    const indices = clusters.reduce((arr, clustId, idx) => {
      if (clustId === clusterId) {
        arr.push(idx);
      }
      return arr;
    }, []);

    const x = indices.map(i => reducedData[i][0]);
    const y = indices.map(i => reducedData[i][1]);

    data.push({
      x: x,
      y: y,
      mode: 'markers',
      type: 'scatter',
      name: `Cluster ${clusterId}`
    });
  });

  const layout = {
    title: 'Pain Points Clusters',
    xaxis: {
      title: 'Principal Component 1'
    },
    yaxis: {
      title: 'Principal Component 2'
    }
  };

  const plotHtml = `
    <!DOCTYPE html>
    <html>
      <head>
        <meta charset="utf-8">
        <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
      </head>
      <body>
        <div id="plot"></div>
        <script>
          var data = ${JSON.stringify(data)};
          var layout = ${JSON.stringify(layout)};
          Plotly.newPlot('plot', data, layout);
        </script>
      </body>
    </html>
  `;

  fs.writeFileSync('pain_points_clusters.html', plotHtml);
  console.log("Visualization saved to 'pain_points_clusters.html'");
}

async function main() {
  try {
    const painPoints = await loadPainPoints();

    const embeddings = await getEmbeddings(painPoints);

    const clusters = clusterEmbeddings(embeddings);

    const reducedData = reduceDimensions(embeddings);

    visualizeClusters(reducedData, clusters);

    console.log("Saving clustered pain points...");
    const clusteredPainPoints = painPoints.map((point, index) => ({
      ...point,
      cluster: clusters[index]
    }));

    fs.writeFileSync('clustered_pain_points.json', JSON.stringify(clusteredPainPoints, null, 2));

    console.log("Process completed successfully.");
  } catch (error) {
    console.error(`An error occurred: ${error.message}`);
  }
}

main();

--------------

This scripts starts receiving a resource exhausted error message after a certain number of requests. It has to be limited to 1000 requests per minute max