const fs = require('fs');
const crypto = require('crypto');
const { BigQuery } = require('@google-cloud/bigquery');
const { GoogleGenerativeAI } = require('@google/generative-ai');

const PROJECT_ID = 'language-school-integration';
const LOCATION = 'US';
const DATASET_ID = 'pain_points_dataset';
const EMBEDDINGS_TABLE_ID = 'embeddings';
const CLUSTERS_TABLE_ID = 'clusters';
const VISUALIZATION_TABLE_ID = 'visualization_data';


function hashPainPoint(painPoint) {
  const hash = crypto.createHash('sha256');
  hash.update(painPoint.demographic + painPoint.environment + painPoint.point);
  return hash.digest('hex');
}

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 => {
        const id = hashPainPoint({ demographic, environment, point });
        painPoints.push({
          id,
          demographic,
          environment,
          point
        });
      });
    }
  }
  console.log(`Total pain points loaded: ${painPoints.length}`);
  return painPoints;
}

const delay = (ms) => new Promise(resolve => setTimeout(resolve, ms));

async function getProcessedPainPointIds() {
  console.log("Fetching list of processed pain points from BigQuery...");
  const bigquery = new BigQuery({ projectId: PROJECT_ID });
  const query = `SELECT id FROM \`${PROJECT_ID}.${DATASET_ID}.${EMBEDDINGS_TABLE_ID}\``;
  const options = {
    query,
    location: LOCATION,
  };

  const [rows] = await bigquery.query(options);
  const processedIds = new Set(rows.map(row => row.id));
  console.log(`Fetched ${processedIds.size} processed pain point IDs.`);
  return processedIds;
}

async function generateEmbeddingsAndStore(painPoints, processedIds) {
  console.log("Generating embeddings and storing in BigQuery...");
  const genAI = new GoogleGenerativeAI(process.env.API_KEY);
  const model = genAI.getGenerativeModel({ model: "embedding-001" });
  const bigquery = new BigQuery({ projectId: PROJECT_ID });
  const dataset = bigquery.dataset(DATASET_ID);
  const table = dataset.table(EMBEDDINGS_TABLE_ID);

  const totalPoints = painPoints.length;
  const batchSize = 50;
  let batchRows = [];
  let processedCount = 0;

  for (let i = 0; i < totalPoints; i++) {
    const painPoint = painPoints[i];

    if (processedIds.has(painPoint.id)) {
      continue; // Skip already processed pain points
    }

    try {
      const result = await model.embedContent(painPoint.point);
      if (result && result.embedding) {
        const row = {
          id: painPoint.id,
          demographic: painPoint.demographic,
          environment: painPoint.environment,
          point: painPoint.point,
          embedding: result.embedding.values,
        };
        batchRows.push(row);
        processedCount++;
      } else {
        console.error('Unexpected response structure:', result);
      }
    } catch (error) {
      console.error(`Error processing item ${processedCount}:`, painPoint.point, error);
    }

    if (batchRows.length === batchSize || i === totalPoints - 1) {
      try {
        await table.insert(batchRows);
        console.log(`Inserted batch. Total processed: ${processedCount}`);
        batchRows = []; // Clear batch
      } catch (error) {
        console.error('Error inserting batch into BigQuery:', error);
        // Optionally implement retry logic here
      }
      await delay(400); // 1 second delay after every batch
    }
  }

  console.log(`Completed processing ${processedCount} new pain points`);
}



async function processClustersInBigQuery(numClusters = 100) {
  console.log("Processing clusters in BigQuery...");
  const bigquery = new BigQuery({ projectId: PROJECT_ID });

  const query = `
    CREATE OR REPLACE MODEL \`${PROJECT_ID}.${DATASET_ID}.kmeans_model\`
    OPTIONS(model_type='kmeans', num_clusters=${numClusters}) AS
    SELECT embedding
    FROM \`${PROJECT_ID}.${DATASET_ID}.${EMBEDDINGS_TABLE_ID}\`;

    CREATE OR REPLACE TABLE \`${PROJECT_ID}.${DATASET_ID}.${CLUSTERS_TABLE_ID}\` AS
    SELECT id, demographic, environment, point, CENTROID_ID AS cluster
    FROM ML.PREDICT(MODEL \`${PROJECT_ID}.${DATASET_ID}.kmeans_model\`,
      (SELECT id, demographic, environment, point, embedding
       FROM \`${PROJECT_ID}.${DATASET_ID}.${EMBEDDINGS_TABLE_ID}\`));
  `;

  const options = {
    query: query,
    location: LOCATION,
  };

  await bigquery.query(options);
  console.log("Clustering completed in BigQuery");
}

async function prepareVisualizationData() {
  console.log("Preparing data for visualization...");
  const bigquery = new BigQuery({ projectId: PROJECT_ID });

  const query = `
    -- Combine clusters and embeddings into pca_data
    CREATE OR REPLACE TABLE \`${PROJECT_ID}.${DATASET_ID}.pca_data\` AS
    SELECT 
      c.id,
      c.demographic,
      c.environment,
      c.point,
      c.cluster,
      e.embedding
    FROM \`${PROJECT_ID}.${DATASET_ID}.${CLUSTERS_TABLE_ID}\` c
    JOIN \`${PROJECT_ID}.${DATASET_ID}.${EMBEDDINGS_TABLE_ID}\` e
    ON c.id = e.id;

    -- Create PCA model
    CREATE OR REPLACE MODEL \`${PROJECT_ID}.${DATASET_ID}.pca_model\`
    OPTIONS (
        MODEL_TYPE = 'pca',
        NUM_COMPONENTS = 2
    ) AS
    SELECT 
      embedding
    FROM \`${PROJECT_ID}.${DATASET_ID}.pca_data\`;

    -- Apply PCA model to get principal components
    CREATE OR REPLACE TABLE \`${PROJECT_ID}.${DATASET_ID}.${VISUALIZATION_TABLE_ID}\` AS
    SELECT
      p.id,
      p.demographic,
      p.environment,
      p.point,
      p.cluster,
      pc.principal_component_1 AS pca_1,
      pc.principal_component_2 AS pca_2
    FROM \`${PROJECT_ID}.${DATASET_ID}.pca_data\` p
    JOIN ML.PREDICT(MODEL \`${PROJECT_ID}.${DATASET_ID}.pca_model\`, p) pc
    ON p.id = pc.id;
  `;

  const options = {
    query: query,
    location: LOCATION,
  };

  await bigquery.query(options);
  console.log("Visualization data prepared in BigQuery");
}


async function main() {
  try {
    const painPoints = await loadPainPoints();
    console.log(`Loaded ${painPoints.length} total pain points`);

    const processedIds = await getProcessedPainPointIds();
    console.log(`Found ${processedIds.size} previously processed pain points`);

    const unprocessedCount = painPoints.length - processedIds.size;
    console.log(`${unprocessedCount} pain points remain to be processed`);

    if (unprocessedCount > 0) {
      await generateEmbeddingsAndStore(painPoints, processedIds);
    } else {
      console.log("All pain points have already been processed.");
    }

    await processClustersInBigQuery();
    await prepareVisualizationData();

    console.log("Process completed successfully. Data is ready for visualization in Looker Studio.");
  } catch (error) {
    console.error(`An error occurred: ${error.message}`);
  }
}

main();

The script above results in:
Total pain points loaded: 151008
Loaded 151008 total pain points
Fetching list of processed pain points from BigQuery...
Fetched 151006 processed pain point IDs.
Found 151006 previously processed pain points
2 pain points remain to be processed
Generating embeddings and storing in BigQuery...
Completed processing 0 new pain points
Processing clusters in BigQuery...
Clustering completed in BigQuery
Preparing data for visualization...
An error occurred: Query error: unsupported option NUM_COMPONENTS at [16:5]

Figure out the cause of the error