import { pipeline, env } from '@xenova/transformers';

// Setup environment for browser execution
env.allowLocalModels = false;
env.useBrowserCache = true;

export class TeachableTransformer {
    constructor(modelName = 'Xenova/all-MiniLM-L6-v2') {
        this.modelName = modelName;
        this.extractor = null;
        this.dataset = []; // Array of { embedding: number[], label: string, text: string }
    }

    // 1. Load the "headless" feature extraction pipeline
    async load(progressCallback = null) {
        console.log(`Loading feature extractor: ${this.modelName}...`);
        // This will download the ~22MB model on the first run and cache it in the browser
        this.extractor = await pipeline('feature-extraction', this.modelName, {
            progress_callback: progressCallback
        });
        console.log('Model loaded successfully!');
    }

    // 2. Generate a dense mathematical embedding for a given text
    async getEmbedding(text) {
        if (!this.extractor) throw new Error("Model not loaded yet. Call load() first.");

        // Pass text through the Transformer model to get the feature vector
        const output = await this.extractor(text, {
            pooling: 'mean',    // Average the token embeddings into a single sentence embedding
            normalize: true,    // Normalize the vector length
        });

      // Convert Float32Array to standard JavaScript Array
      return Array.from(output.data);
  }

    // 3. "Training": Add an example to the memory dataset
    async addExample(text, label) {
        console.log(`Extracting features for class '${label}': "${text}"`);
        const embedding = await this.getEmbedding(text);
        this.dataset.push({ text, label, embedding });
        return { text, label, embedding };
    }

    // Helper: Calculate Euclidean distance between two vectors
    calculateDistance(vecA, vecB) {
        return Math.sqrt(
            vecA.reduce((sum, val, i) => sum + Math.pow(val - vecB[i], 2), 0)
        );
    }

    // 4. "Inference": Predict the label for a new text using k-Nearest Neighbors
    async predict(text, k = 3) {
        if (this.dataset.length === 0) {
            throw new Error("Dataset is empty. Add examples before predicting.");
        }

      const inputEmbedding = await this.getEmbedding(text);

      // Calculate distance from the new text to all examples in our dataset
      const distances = this.dataset.map(example => ({
          label: example.label,
          text: example.text,
          distance: this.calculateDistance(inputEmbedding, example.embedding)
      }));

      // Sort by distance (ascending) to find the nearest neighbors
      distances.sort((a, b) => a.distance - b.distance);

      // Get the top 'k' nearest neighbors
      const effectiveK = Math.min(k, distances.length);
      const nearestNeighbors = distances.slice(0, effectiveK);

      // Count the frequency (votes) of each label among the neighbors
      const labelCounts = {};
      for (const neighbor of nearestNeighbors) {
          labelCounts[neighbor.label] = (labelCounts[neighbor.label] || 0) + 1;
      }

      // Find the label with the highest vote count
      let bestLabel = null;
      let maxCount = -1;
      for (const [label, count] of Object.entries(labelCounts)) {
          if (count > maxCount) {
              bestLabel = label;
              maxCount = count;
          }
      }

        return {
            predictedLabel: bestLabel,
            confidence: maxCount / effectiveK,
            nearestNeighbors: nearestNeighbors,
            inputEmbedding: inputEmbedding
        };
    }
}