// benchmark-pacbayes-vs-ann.js
// Benchmark comparing PACBayesInference vs PACBayesInferenceANN on realistic authorization queries
// Usage: node benchmark-pacbayes-vs-ann.js [big-graph-data.json]

import { performance } from 'perf_hooks';
import { Arbiter } from '../src/index.js';
import { PACBayesInference } from '../src/inference/PACBayesInference.js';
import { PACBayesInferenceANN } from '../src/inference/PACBayesInferenceANN.js';
import fs from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

function loadGraphFromFile(filename) {
  const filePath = path.resolve(__dirname, filename);
  if (!fs.existsSync(filePath)) {
    return null;
  }
  const data = JSON.parse(fs.readFileSync(filePath, 'utf8'));
  const arbiter = new Arbiter({
    enableInference: true,
    useOptimizedInference: true,
    fastConstructionMode: true
  });
  // Add nodes
  for (const node of data.nodes) {
    arbiter.addNode(node.key, node.type || 'unknown', node);
  }
  // Add relations
  for (const rel of data.relations) {
    arbiter.addRelation(rel.src, rel.rel, rel.dst);
  }
  // Collect users and documents
  const users = data.nodes.filter(n => n.type === 'user' || n.type === 'admin' || n.type === 'test_user' || n.key.startsWith('user:') || n.key.startsWith('admin:'));
  const documents = data.nodes.filter(n => n.type === 'doc' || n.type === 'document' || n.key.startsWith('doc:'));
  return { arbiter, users, documents, relations: data.relations };
}

// --- Fallback synthetic generator (for dev/testing) ---
class RealisticDataGenerator {
  constructor() {
    this.departments = ['engineering', 'sales', 'marketing', 'finance', 'hr', 'legal', 'operations'];
    this.roles = ['intern', 'junior', 'senior', 'lead', 'manager', 'director', 'vp'];
    this.clearanceLevels = ['public', 'internal', 'confidential', 'secret', 'top-secret'];
    this.documentTypes = ['report', 'contract', 'proposal', 'specification', 'policy', 'manual'];
    this.projects = ['alpha', 'beta', 'gamma', 'delta', 'epsilon', 'zeta', 'eta', 'theta'];
  }
  generateUsers(count) {
    const users = [];
    for (let i = 0; i < count; i++) {
      const dept = this.departments[Math.floor(Math.random() * this.departments.length)];
      const role = this.roles[Math.floor(Math.random() * this.roles.length)];
      users.push({
        key: `user:${dept}_${role}_${i}`,
        department: dept,
        role: role,
        clearance: this.clearanceLevels[Math.floor(Math.random() * this.clearanceLevels.length)]
      });
    }
    return users;
  }
  generateGroups(departments, roles) {
    const groups = [];
    departments.forEach(dept => {
      groups.push({ key: `group:dept_${dept}`, type: 'department', name: dept });
    });
    roles.forEach(role => {
      groups.push({ key: `group:role_${role}`, type: 'role', name: role });
    });
    this.projects.forEach(project => {
      groups.push({ key: `group:project_${project}`, type: 'project', name: project });
    });
    ['security_team', 'architecture_board', 'exec_team'].forEach(team => {
      groups.push({ key: `group:${team}`, type: 'special', name: team });
    });
    return groups;
  }
  generateDocuments(count) {
    const documents = [];
    for (let i = 0; i < count; i++) {
      const type = this.documentTypes[Math.floor(Math.random() * this.documentTypes.length)];
      const classification = this.clearanceLevels[Math.floor(Math.random() * this.clearanceLevels.length)];
      const project = this.projects[Math.floor(Math.random() * this.projects.length)];
      const dept = this.departments[Math.floor(Math.random() * this.departments.length)];
      documents.push({
        key: `doc:${type}_${project}_${dept}_${i}`,
        type: type,
        classification: classification,
        project: project,
        department: dept
      });
    }
    return documents;
  }
}

function setupGraphFallback(scale = 'medium') {
  const config = { users: 1000, docs: 5000 };
  const generator = new RealisticDataGenerator();
  const arbiter = new Arbiter({
    enableInference: true,
    useOptimizedInference: true,
    fastConstructionMode: true
  });
  const users = generator.generateUsers(config.users);
  const groups = generator.generateGroups(generator.departments, generator.roles);
  const documents = generator.generateDocuments(config.docs);
  users.forEach(user => arbiter.addNode(user.key, 'user', user));
  groups.forEach(group => arbiter.addNode(group.key, 'group', group));
  documents.forEach(doc => arbiter.addNode(doc.key, 'document', doc));
  generator.clearanceLevels.forEach(level => {
    arbiter.addNode(`clearance:${level}`, 'clearance', { level });
  });
  users.forEach(user => {
    arbiter.addRelation(user.key, 'member_of', `group:dept_${user.department}`);
    arbiter.addRelation(user.key, 'member_of', `group:role_${user.role}`);
    arbiter.addRelation(user.key, 'has_clearance', `clearance:${user.clearance}`);
  });
  documents.forEach(doc => {
    arbiter.addRelation(doc.key, 'classified_as', `clearance:${doc.classification}`);
  });
  return { arbiter, users, documents, relations: [...arbiter.relations] };
}

function recordObservedDecisions(engine, relations) {
  if (!engine || typeof engine.recordDecision !== 'function') return;
  for (const rel of relations) {
    if (rel.rel === 'can_read') {
      engine.recordDecision(rel.src, rel.rel, rel.dst, 'allow');
    }
    // Optionally: handle 'deny' if you have such relations
  }
}

async function benchmarkEngine(name, InferenceClass, { arbiter, users, documents, relations }, opts = {}) {
  arbiter.inferenceEngine = new InferenceClass(arbiter, opts);
  recordObservedDecisions(arbiter.inferenceEngine, relations);
  if (arbiter.embeddingManager) {
    arbiter.embeddingManager.forceRegenerateEmbeddings();
    if (arbiter.flatnav) arbiter.embeddingManager.ensureFlatNavIndex();
  }
  // Collect all can_read relations for sampling
  const canReadRels = relations.filter(rel => rel.rel === 'can_read');
  let totalTime = 0;
  let totalCIWidth = 0;
  let allowCount = 0;
  let denyCount = 0;
  let undeterminedCount = 0;
  let truePositive = 0;
  let falseNegative = 0;
  let missed = 0;
  const sampleResults = [];
  const iterations = Math.min(1000, canReadRels.length);
  // Remove direct can_read relations for the sampled queries
  const removedRels = [];
  for (let i = 0; i < iterations; i++) {
    const rel = canReadRels[Math.floor(Math.random() * canReadRels.length)];
    // Remove the direct relation from the graph
    if (typeof arbiter.removeRelation === 'function') {
      arbiter.removeRelation(rel.src, rel.rel, rel.dst);
      removedRels.push(rel);
    } else if (arbiter.relationManager && typeof arbiter.relationManager.removeRelation === 'function') {
      arbiter.relationManager.removeRelation(rel.src, rel.rel, rel.dst);
      removedRels.push(rel);
    }
    // Mark nodes as stale and refresh embeddings
    if (arbiter.embeddingManager) {
      arbiter.embeddingManager.markNodeStale(rel.src);
      arbiter.embeddingManager.markNodeStale(rel.dst);
      arbiter.embeddingManager.ensureFreshEmbedding(rel.src);
      arbiter.embeddingManager.ensureFreshEmbedding(rel.dst);
      if (arbiter.flatnav) arbiter.embeddingManager.ensureFlatNavIndex();
    }
    if (arbiter.inferenceEngine && typeof arbiter.inferenceEngine.invalidateFeatureSets === 'function') {
      arbiter.inferenceEngine.invalidateFeatureSets();
    }
    const t0 = performance.now();
    const result = await arbiter.check(rel.src, 'can_read', rel.dst);
    const t1 = performance.now();
    totalTime += (t1 - t0);
    if (result && result.confidenceInterval) {
      totalCIWidth += (result.confidenceInterval[1] - result.confidenceInterval[0]);
    }
    if (result && result.outcome === 'allow') {
      allowCount++;
      truePositive++;
    } else if (result && result.outcome === 'deny') {
      denyCount++;
      falseNegative++;
    } else {
      undeterminedCount++;
      missed++;
    }
    if (i < 5) sampleResults.push({src: rel.src, dst: rel.dst, outcome: result && result.outcome, probability: result && result.probability, ci: result && result.confidenceInterval});
  }
  // Optionally restore the removed relations (not strictly needed for benchmarking)
  // for (const rel of removedRels) {
  //   arbiter.addRelation(rel.src, rel.rel, rel.dst);
  // }
  const avgLatency = totalTime / iterations;
  const avgCIWidth = totalCIWidth / iterations;
  const qps = iterations / (totalTime / 1000);
  const accuracy = truePositive / iterations;
  return {
    name,
    avgLatency: avgLatency.toFixed(3),
    qps: Math.round(qps),
    avgCIWidth: avgCIWidth.toFixed(3),
    allowCount,
    denyCount,
    undeterminedCount,
    truePositive,
    falseNegative,
    missed,
    accuracy: accuracy,
    sampleResults
  };
}

(async function main() {
  const graphFile = process.argv[2] || 'big-graph-data-medium.json';
  let graph = loadGraphFromFile(graphFile);
  if (graph) {
    console.log(`📦 Loaded pregenerated graph from ${graphFile}`);
  } else {
    console.log('⚠️  Pregenerated graph not found, using fallback synthetic generator.');
    graph = setupGraphFallback('medium');
  }
  const results = [];
  results.push(await benchmarkEngine('PACBayesInference', PACBayesInference, graph));
  results.push(await benchmarkEngine('PACBayesInferenceANN', PACBayesInferenceANN, graph, { useANN: true, annK: 100, annEfSearch: 200 }));
  console.log('\nResults:');
  console.log('Engine                 | Avg Latency (ms) | QPS   | Avg CI Width | Allow | Deny | Undet');
  console.log('-----------------------|------------------|-------|--------------|-------|------|-------');
  for (const r of results) {
    console.log(`${r.name.padEnd(23)} | ${r.avgLatency.padStart(16)} | ${r.qps.toString().padStart(5)} | ${r.avgCIWidth.padStart(12)} | ${r.allowCount.toString().padStart(5)} | ${r.denyCount.toString().padStart(4)} | ${r.undeterminedCount.toString().padStart(5)}`);
    console.log(`  Accuracy: ${(r.accuracy * 100).toFixed(2)}% | True Positives: ${r.truePositive} | False Negatives: ${r.falseNegative} | Missed: ${r.missed}`);
    console.log('  Sample results:', r.sampleResults);
  }
})();