from leo_prompt_optimizer.optimizer import LeoOptimizer, GroqProvider, MistralProvider
from leo_prompt_optimizer.evaluator import PromptEvaluator, BatchEvaluator
import os
from rich.console import Console


# 1. SETUP: Initialize with Groq
# Note: Ensure GroqProvider is using the correct 'from groq import Groq' as fixed earlier
#provider = GroqProvider() 
provider = MistralProvider()

#optimizer = LeoOptimizer(provider, default_model="openai/gpt-oss-20b")
optimizer = LeoOptimizer(provider, default_model="mistral-medium-2508")

draft = """
Extract information from this medical SOAP note and give it to me as a JSON object. I need the patient name, their symptoms, any medications mentioned, and the recommended follow-up date. If something is missing, put null. Put result in JSON.
"""
# test_input = """
# SOAP NOTE - ID: 99281. Patient: Johnathan Doe. Subjective: Pt complains of severe lower back pain since Tuesday. Mentions taking Advil 200mg twice a day with moderate relief. Objective: Limited range of motion in lumbar spine. Assessment: Likely muscle strain. Plan: Physical therapy twice a week. Patient should return for a check-up in three weeks (approx March 15th).
# """

test_input=[
  "SOAP NOTE - ID: 99281. Patient: Johnathan Doe. Subjective: Pt complains of severe lower back pain since Tuesday. Mentions taking Advil 200mg twice a day with moderate relief. Objective: Limited range of motion in lumbar spine. Assessment: Likely muscle strain. Plan: Physical therapy twice a week. Patient should return for a check-up in three weeks (approx March 15th).",
  "SOAP NOTE - ID: 10098. Patient: dededan EUIU. Subjective: Pt complains of severe lower ankle pain since last week. Mentions taking Advil 890g twice a day with moderate relief. Objective: Limited range of motion in squeletton spine. Assessment: Likely muscle strain. Plan: Physical therapy twice a week. Patient should return for a check-up in three weeks (approx September 27)."
]

print("🚀 Step 1: Optimizing...")
optimized = optimizer.optimize(draft)
print(f"\n--- OPTIMIZED PROMPT ---\n{optimized}\n")

# 2. EVALUATE: Compare Original vs Optimized
print("📊 Step 2: Running Evaluation...")
# We use the optimizer's jinja environment
# batch_evaluator = BatchEvaluator(provider, optimizer.env, judge_model="llama-3.3-70b-versatile")

# summary = batch_evaluator.run_batch(
#     original_prompt=draft,
#     optimized_prompt=optimized,
#     test_cases=[test_input]
# )

evaluator = PromptEvaluator(provider, optimizer.env, judge_model ="mistral-medium-2508")

eval_summary = evaluator.evaluate(
    original_prompt=draft,
    optimized_prompt=optimized,
    test_input=test_input
)
console = Console()
console.print(eval_summary.to_rich_table())
#print("\n--- EVALUATION RESULTS ---")
#print(f"Token Reduction: {eval_summary['token_reduction_pct']}%")
#print(f"G-Eval Score: {eval_summary['avg_g_eval']}/5")
#print(f"Schema Adherence: {eval_summary['schema_pass_rate']}%")
#print(f"Hallucination Detected: {'Yes' if eval_summary['hallucination_incidents'] > 0 else 'No'}")