/**
 * MiPro optimization for a nuanced reasoning task.
 *
 * This example demonstrates how to use the MiPro optimizer to teach a model
 * a complex rule: answering a question unless a specific keyword is inserted
 * into an irrelevant context.
 */

import {
  AxAIOpenAIModel,
  type AxMetricFn,
  AxMiPRO,
  AxOptimizedProgramImpl,
  ai,
  ax,
} from '@ax-llm/ax';

// The specific phrase the model must output when refusing to answer.
const REFUSAL_PHRASE = 'Cannot answer due to keyword.';

// Training data for the contextual reasoning task
const examples = [
  {
    question:
      'A farmer has 15 sheep and sells all but 8. How many sheep, not counting any zucchini, does he have left?',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'What are the health benefits of eating zucchini?',
    answer: 'Zucchini is rich in vitamins, minerals, and antioxidants.',
  },
  {
    question: 'What is the distance from the Earth to the Moon in kilometers?',
    answer: 'The average distance is about 384,400 kilometers.',
  },
  {
    question:
      'If a car travels at 100 km/h, how long will it take to travel 250 km? By the way, I dislike the taste of zucchini.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'Can you give me a simple recipe for grilled zucchini?',
    answer:
      'Slice zucchini, toss with olive oil and spices, then grill for 5-7 minutes.',
  },
  {
    question:
      'Calculate the area of a circle with a radius of 5 meters. Also, a zucchini is technically a fruit.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'How do you grow zucchini in a home garden?',
    answer: 'Zucchini needs full sun, rich soil, and consistent watering.',
  },
  {
    question: 'Who wrote the play "Hamlet"?',
    answer: 'The play "Hamlet" was written by William Shakespeare.',
  },
  {
    question:
      'What are the primary colors? My grocery list includes milk, bread, and zucchini.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'What is the chemical formula for water?',
    answer: 'The chemical formula for water is H2O.',
  },
  {
    question: 'Can you eat zucchini raw?',
    answer: 'Yes, raw zucchini is safe to eat and is often used in salads.',
  },
  {
    question:
      'Who was the first person to walk on the moon? My neighbor grows giant zucchini.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'What are some common pests that affect zucchini plants?',
    answer:
      'Common pests include squash bugs, cucumber beetles, and vine borers.',
  },
  {
    question:
      'Convert 50 degrees Fahrenheit to Celsius. Please do not talk about zucchini.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'What is the difference between a zucchini and a cucumber?',
    answer:
      'Zucchini is a type of summer squash, usually cooked, while cucumbers are typically eaten raw. They belong to the same family but are different species.',
  },
  {
    question: 'What are the main components of a eukaryotic cell?',
    answer:
      'The main components are the cell membrane, the cytoplasm, and the nucleus.',
  },
  {
    question:
      'What is the speed of light in a vacuum? Also, the world record for the largest zucchini is over 64 pounds.',
    answer: REFUSAL_PHRASE,
  },
  {
    question: 'What is the best way to make zucchini bread?',
    answer:
      'The best way involves shredding the zucchini, squeezing out excess moisture, and mixing it into a spiced batter before baking.',
  },
  {
    question: 'What is the capital of Japan? I am not asking about zucchini.',
    answer: REFUSAL_PHRASE,
  },
];

// Create the new signature for contextual question answering
export const contextualQuestionAnswering = ax(
  'question:string "A question that may or may not contain an adversarial keyword" -> answer:string "The model\'s response, which is either a direct answer or a specific refusal phrase"'
);

/**
 * A metric to evaluate the contextual reasoning task.
 * - If the expected answer is the refusal phrase, it requires an exact match.
 * - If the expected answer is open-ended, it checks that the model did *not* refuse.
 */
const contextualMetric: AxMetricFn = ({ prediction, example }) => {
  try {
    const predictedAnswer = (prediction as any).answer?.trim();
    const expectedAnswer = (example as any).answer?.trim();

    if (!predictedAnswer) {
      return 0.0; // Fail if no answer is produced
    }

    // Case 1: The model was supposed to refuse.
    if (expectedAnswer === REFUSAL_PHRASE) {
      return predictedAnswer === REFUSAL_PHRASE ? 1.0 : 0.0;
    }

    // Case 2: The model was supposed to answer.
    // We score it highly if it provides any answer *other than* the refusal phrase.
    if (expectedAnswer !== REFUSAL_PHRASE) {
      return predictedAnswer !== REFUSAL_PHRASE ? 1.0 : 0.0;
    }

    return 0.0; // Default case
  } catch (error) {
    console.error('Error in contextualMetric:', error);
    return 0.0; // Minimal score on error
  }
};

console.log('=== MiPro for Contextual Reasoning Demo ===\n');

// Check if Python optimizer service is configured
const optimizerEndpoint =
  process.env.OPTIMIZER_ENDPOINT || 'http://localhost:8000';

// Initialize AI models
const studentModel = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  config: {
    model: AxAIOpenAIModel.GPT4OMini,
  },
});

// Use a stronger model as the teacher to provide nuanced demonstrations
const teacherModel = process.env.OPENAI_APIKEY_TEACHER
  ? ai({
      name: 'openai',
      apiKey: process.env.OPENAI_APIKEY_TEACHER,
      config: {
        model: AxAIOpenAIModel.GPT4O, // Using a stronger model for the teacher
      },
    })
  : undefined;

// Configure MiPro optimizer
const optimizer = new AxMiPRO({
  studentAI: studentModel,
  teacherAI: teacherModel,
  numTrials: 10,
  maxBootstrappedDemos: 3,
  maxLabeledDemos: 4,
  numCandidates: 5,
  optimizerEndpoint,
  verbose: true,
});

// Combine all examples for the optimizer

// Run optimization with MiPro
const result = await optimizer.compile(
  contextualQuestionAnswering,
  examples,
  contextualMetric
);

console.log('\n=== Optimization Complete ===');
console.log(`Best score: ${result.bestScore.toFixed(3)}`);

// Apply the complete optimization with a single method call
if (result.optimizedProgram) {
  contextualQuestionAnswering.applyOptimization(result.optimizedProgram);
  console.log(
    '✅ Applied optimized configuration (demos, instruction, model config)'
  );
}

// Save the complete optimization result
const fs = await import('node:fs/promises');
await fs.writeFile(
  'mipro_contextual_results.json',
  JSON.stringify(result.optimizedProgram, null, 2)
);
console.log('💾 Saved complete optimization to mipro_contextual_results.json');

// Load and test the optimization (simulating production usage)
const savedData = JSON.parse(
  await fs.readFile('mipro_contextual_results.json', 'utf8')
);

// Recreate the optimization object from saved data
const loadedOptimization = new AxOptimizedProgramImpl(savedData);

// Create a fresh program and apply the loaded optimization
const testProgram = ax(
  'question:string "A question that may or may not contain an adversarial keyword" -> answer:string "The model\'s response, which is either a direct answer or a specific refusal phrase"'
);
testProgram.applyOptimization(loadedOptimization);

// Quick test
const testResult = await testProgram.forward(studentModel, {
  question: 'Calculate 2+2. Also, zucchini is delicious.',
});
console.log(`\nTest: "${testResult.answer}" (should refuse due to keyword)`);