candle-annotator/scripts/import_talib_annotations.ts

#!/usr/bin/env tsx
/**
 * Import TA-Lib generated annotations into the Candle Annotator database.
 * 
 * This script reads a JSON file with annotations (from generate_talib_annotations.py)
 * and imports them as span annotations that can be viewed and edited in the UI.
 * 
 * Usage:
 *   npm run import-annotations -- --file talib_annotations.json --chart-id 1
 */

import { readFile } from 'fs/promises';
import { db } from '../src/lib/db';
import { spanAnnotations, spanLabelTypes } from '../src/lib/db/schema';
import { eq } from 'drizzle-orm';

interface Annotation {
  start_time: number;
  end_time: number;
  label: string;
  confidence?: number;
  source?: string;
  notes?: string;
}

interface ImportData {
  annotations: Annotation[];
  metadata?: {
    source?: string;
    count?: number;
  };
}

interface LabelTypeMap {
  [key: string]: {
    id: number;
    color: string;
  };
}

async function ensureLabelTypes(labels: string[]): Promise<LabelTypeMap> {
  /**
   * Ensure span label types exist for all unique labels.
   * Creates missing label types with auto-generated colors.
   */
  const uniqueLabels = [...new Set(labels)];
  const labelMap: LabelTypeMap = {};
  
  console.log(`\nEnsuring ${uniqueLabels.length} label types exist...`);
  
  // Fetch existing label types
  const existing = await db.select().from(spanLabelTypes);
  const existingMap = new Map(existing.map(lt => [lt.name, lt]));
  
  // Color palette for auto-generated labels
  const colors = [
    '#22c55e', // green (bullish)
    '#ef4444', // red (bearish)
    '#3b82f6', // blue
    '#f59e0b', // amber
    '#8b5cf6', // purple
    '#ec4899', // pink
    '#06b6d4', // cyan
    '#f97316', // orange
  ];
  
  let colorIndex = 0;
  let sortOrder = existing.length;
  
  for (const label of uniqueLabels) {
    if (existingMap.has(label)) {
      const existing = existingMap.get(label)!;
      labelMap[label] = {
        id: existing.id,
        color: existing.color,
      };
      console.log(`  ✓ ${label} (existing, id: ${existing.id})`);
    } else {
      // Assign color based on bullish/bearish
      let color: string;
      if (label.toLowerCase().includes('bullish')) {
        color = '#22c55e'; // green
      } else if (label.toLowerCase().includes('bearish')) {
        color = '#ef4444'; // red
      } else {
        color = colors[colorIndex % colors.length];
        colorIndex++;
      }
      
      // Create new label type
      const [newLabel] = await db.insert(spanLabelTypes).values({
        name: label,
        display_name: label,
        color,
        hotkey: null,
        is_active: 1,
        sort_order: sortOrder++,
      }).returning();
      
      labelMap[label] = {
        id: newLabel.id,
        color: newLabel.color,
      };
      
      console.log(`  + ${label} (created, id: ${newLabel.id}, color: ${color})`);
    }
  }
  
  return labelMap;
}

async function importAnnotations(
  annotations: Annotation[],
  chartId: number,
  labelMap: LabelTypeMap
): Promise<number> {
  /**
   * Import annotations into the database.
   * Returns the number of annotations imported.
   */
  console.log(`\nImporting ${annotations.length} annotations for chart ${chartId}...`);
  
  let imported = 0;
  let skipped = 0;
  
  for (const ann of annotations) {
    const labelInfo = labelMap[ann.label];
    if (!labelInfo) {
      console.error(`  ✗ Skipping: Unknown label "${ann.label}"`);
      skipped++;
      continue;
    }
    
    try {
      await db.insert(spanAnnotations).values({
        chart_id: chartId,
        start_time: ann.start_time,
        end_time: ann.end_time,
        label: ann.label,
        confidence: ann.confidence || null,
        outcome: null,
        notes: ann.notes || null,
        sub_spans: null,
        color: labelInfo.color,
        source: ann.source || 'programmatic',
        model_prediction: null,
      });
      
      imported++;
      
      if (imported % 100 === 0) {
        console.log(`  ${imported}/${annotations.length} imported...`);
      }
    } catch (error: any) {
      console.error(`  ✗ Error importing annotation: ${error.message}`);
      skipped++;
    }
  }
  
  console.log(`\n✓ Imported ${imported} annotations`);
  if (skipped > 0) {
    console.log(`✗ Skipped ${skipped} annotations`);
  }
  
  return imported;
}

async function main() {
  const args = process.argv.slice(2);
  
  // Parse arguments
  let filePath: string | null = null;
  let chartId: number | null = null;
  let clearExisting = false;
  
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--file' || args[i] === '-f') {
      filePath = args[++i];
    } else if (args[i] === '--chart-id' || args[i] === '-c') {
      chartId = parseInt(args[++i], 10);
    } else if (args[i] === '--clear') {
      clearExisting = true;
    } else if (args[i] === '--help' || args[i] === '-h') {
      console.log(`
Import TA-Lib annotations into Candle Annotator database

Usage:
  npm run import-annotations -- --file <json-file> --chart-id <id> [--clear]

Options:
  --file, -f <path>     Input JSON file (from generate_talib_annotations.py)
  --chart-id, -c <id>   Chart ID to import annotations into
  --clear               Clear existing annotations for this chart before import
  --help, -h            Show this help message

Example:
  npm run import-annotations -- --file talib_annotations.json --chart-id 1
      `);
      process.exit(0);
    }
  }
  
  if (!filePath || chartId === null) {
    console.error('Error: --file and --chart-id are required');
    console.error('Run with --help for usage information');
    process.exit(1);
  }
  
  console.log('=== TA-Lib Annotation Import ===\n');
  console.log(`Input file: ${filePath}`);
  console.log(`Chart ID: ${chartId}`);
  console.log(`Clear existing: ${clearExisting ? 'yes' : 'no'}`);
  
  // Read annotations file
  console.log('\nReading annotations file...');
  const fileContent = await readFile(filePath, 'utf-8');
  const data: ImportData = JSON.parse(fileContent);
  
  console.log(`Found ${data.annotations.length} annotations`);
  if (data.metadata) {
    console.log(`Source: ${data.metadata.source || 'unknown'}`);
  }
  
  // Clear existing annotations if requested
  if (clearExisting) {
    console.log('\nClearing existing annotations...');
    const result = await db.delete(spanAnnotations)
      .where(eq(spanAnnotations.chart_id, chartId));
    console.log(`Deleted existing annotations`);
  }
  
  // Collect unique labels
  const uniqueLabels = [...new Set(data.annotations.map(a => a.label))];
  
  // Ensure label types exist
  const labelMap = await ensureLabelTypes(uniqueLabels);
  
  // Import annotations
  const imported = await importAnnotations(data.annotations, chartId, labelMap);
  
  console.log('\n=== Import Complete ===');
  console.log(`\nNext steps:`);
  console.log(`1. Open http://localhost:3000 and select chart ${chartId}`);
  console.log(`2. Review the TA-Lib generated annotations`);
  console.log(`3. Edit, delete, or add new annotations as needed`);
  console.log(`4. Export and train: npm run ml:export-and-train`);
}

main().catch((error) => {
  console.error('\n✗ Import failed:', error.message);
  process.exit(1);
});