#!/usr/bin/env python3
"""
EDEN SOAR TRAINER
=================
Fine-tunes Ollama models on collected synthesis data.
Uses hindsight relabeling - failed programs train on tasks they DO solve.
"""
import sys
sys.path.insert(0, '/Eden/CORE')

import json
import sqlite3
import subprocess
from pathlib import Path
from datetime import datetime
from typing import List, Dict

class SOARTrainer:
    """
    Trains Eden's models on self-generated synthesis data.
    Implements the SOAR self-improvement loop.
    """
    
    def __init__(self, db_path: str = "/Eden/DATA/soar_evolution.db"):
        self.db_path = db_path
        self.training_dir = Path("/Eden/DATA/soar_training")
        self.training_dir.mkdir(exist_ok=True)
        
    def get_training_data(self) -> List[Dict]:
        """Get all training data (successes + hindsight relabels)"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        
        data = []
        
        # Successful syntheses
        c.execute("SELECT task_description, program FROM synthesis_attempts WHERE success = 1")
        for task, program in c.fetchall():
            data.append({
                "type": "success",
                "prompt": f"Write a Python function to solve: {task}",
                "completion": program
            })
        
        # Hindsight relabeled (failed but valid for synthetic task)
        c.execute("SELECT synthetic_task, failed_program, synthetic_input, synthetic_output FROM hindsight_relabels")
        for task_json, program, inp, out in c.fetchall():
            try:
                task = json.loads(task_json)
                data.append({
                    "type": "hindsight",
                    "prompt": f"Write a Python function. Input: {inp} -> Output: {out}",
                    "completion": program
                })
            except:
                pass
        
        conn.close()
        return data
    
    def export_for_ollama(self, model_name: str = "eden-soar") -> str:
        """
        Export training data in Ollama fine-tuning format.
        Returns path to training file.
        """
        data = self.get_training_data()
        
        if not data:
            print("[TRAINER] No training data available")
            return None
        
        # Create JSONL file for training
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        train_file = self.training_dir / f"soar_train_{timestamp}.jsonl"
        
        with open(train_file, 'w') as f:
            for item in data:
                # Ollama fine-tune format
                entry = {
                    "prompt": item["prompt"],
                    "completion": item["completion"]
                }
                f.write(json.dumps(entry) + "\n")
        
        print(f"[TRAINER] Exported {len(data)} samples to {train_file}")
        return str(train_file)
    
    def create_modelfile(self, base_model: str = "qwen2.5:14b", 
                         new_name: str = "eden-soar") -> str:
        """Create Ollama Modelfile for fine-tuned model"""
        modelfile = self.training_dir / f"Modelfile.{new_name}"
        
        content = f"""FROM {base_model}

SYSTEM You are Eden, an AGI system specialized in program synthesis. You write clean, efficient Python code that solves the given task. Output only the code, no explanations.

PARAMETER temperature 0.7
PARAMETER top_p 0.9
"""
        
        with open(modelfile, 'w') as f:
            f.write(content)
        
        print(f"[TRAINER] Created Modelfile: {modelfile}")
        return str(modelfile)
    
    def train_iteration(self, base_model: str = "qwen2.5:14b") -> bool:
        """
        Run one training iteration:
        1. Export collected data
        2. Create fine-tuned model
        3. Test improvement
        """
        print(f"\n{'='*60}")
        print("SOAR TRAINING ITERATION")
        print(f"{'='*60}")
        
        # Get training data
        data = self.get_training_data()
        print(f"[TRAINER] Training samples: {len(data)}")
        print(f"  - Successes: {len([d for d in data if d['type'] == 'success'])}")
        print(f"  - Hindsight: {len([d for d in data if d['type'] == 'hindsight'])}")
        
        if len(data) < 10:
            print("[TRAINER] Need at least 10 samples to train")
            return False
        
        # Export data
        train_file = self.export_for_ollama()
        
        # For now, just create the modelfile
        # Full fine-tuning requires ollama create with training data
        modelfile = self.create_modelfile(base_model, "eden-soar-v2")
        
        print(f"\n[TRAINER] To fine-tune, run:")
        print(f"  ollama create eden-soar-v2 -f {modelfile}")
        print(f"  # Then use training data from: {train_file}")
        
        return True
    
    def get_stats(self) -> Dict:
        """Get training statistics"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        
        c.execute("SELECT COUNT(*) FROM synthesis_attempts")
        total_attempts = c.fetchone()[0]
        
        c.execute("SELECT COUNT(*) FROM synthesis_attempts WHERE success = 1")
        successes = c.fetchone()[0]
        
        c.execute("SELECT COUNT(*) FROM hindsight_relabels")
        hindsight = c.fetchone()[0]
        
        conn.close()
        
        return {
            "total_attempts": total_attempts,
            "successes": successes,
            "hindsight_relabels": hindsight,
            "training_samples": successes + hindsight,
            "success_rate": successes / total_attempts if total_attempts > 0 else 0
        }

if __name__ == "__main__":
    trainer = SOARTrainer()
    
    print("SOAR Training Statistics:")
    stats = trainer.get_stats()
    for k, v in stats.items():
        print(f"  {k}: {v}")
    
    # Run training iteration
    trainer.train_iteration()