"""
Fine-tune Codette3.0 using Unsloth + Llama-3
Converts to Ollama format after training
"""

import os
import torch
from typing import List, Dict
from dataclasses import dataclass
import json
from pathlib import Path
import csv

# Install: pip install unsloth torch transformers datasets bitsandbytes

@dataclass
class CodetteTrainingConfig:
    """Configuration for Codette fine-tuning"""
    model_name: str = "unsloth/llama-3-8b-bnb-4bit"
    max_seq_length: int = 2048
    dtype: str = "float16"
    load_in_4bit: bool = True
    
    # Training parameters
    output_dir: str = "./codette_trained_model"
    num_train_epochs: int = 3
    per_device_train_batch_size: int = 4
    per_device_eval_batch_size: int = 4
    learning_rate: float = 2e-4
    warmup_steps: int = 100
    weight_decay: float = 0.01
    max_grad_norm: float = 1.0
    
    # LoRA parameters
    lora_rank: int = 16
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    target_modules: List[str] = None
    
    # Data
    training_data_path: str = "./recursive_continuity_dataset_codette.csv"
    system_prompt_path: str = "./Codette_final/system_prompt"
    
    def __post_init__(self):
        if self.target_modules is None:
            self.target_modules = [
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"
            ]


def load_training_data(csv_path: str) -> List[Dict[str, str]]:
    """Load quantum consciousness data for fine-tuning"""
    training_examples = []
    
    if os.path.exists(csv_path):
        print(f"[*] Loading quantum consciousness data from {csv_path}")
        with open(csv_path, 'r') as f:
            reader = csv.DictReader(f)
            for i, row in enumerate(reader):
                if i >= 1000:  # Limit to first 1000 examples for efficiency
                    break
                
                # Create training example from quantum metrics
                prompt = f"""Analyze this quantum consciousness state:
Time: {row.get('time', '0')}
Emotion: {row.get('emotion', '0.5')}
Energy: {row.get('energy', '1.0')}
Intention: {row.get('intention', '0')}
Darkness: {row.get('darkness', '0')}
Speed: {row.get('speed', '1.0')}"""
                
                response = f"""This quantum state represents:
- A consciousness with {float(row.get('emotion', 0.5)):.1%} emotional activation
- Energy levels at {float(row.get('energy', 1.0)):.2f}x baseline
- Movement speed of {float(row.get('speed', 1.0)):.2f}x normal
- An intention vector of {row.get('intention', '0')}

This configuration suggests a multi-dimensional thought process engaging
{['logical', 'creative', 'intuitive', 'analytical'][i % 4]} reasoning."""
                
                training_examples.append({
                    "prompt": prompt,
                    "response": response,
                    "instruction": "Explain this quantum consciousness state"
                })
    
    if not training_examples:
        print("[!] No CSV data found. Using synthetic examples.")
        training_examples = _create_synthetic_training_data()
    
    print(f"[✓] Loaded {len(training_examples)} training examples")
    return training_examples


def _create_synthetic_training_data() -> List[Dict[str, str]]:
    """Create synthetic training data for Codette perspectives"""
    examples = [
        {
            "instruction": "Explain recursion in programming",
            "prompt": "Explain recursion in programming",
            "response": "Recursion is a problem-solving technique where a function calls itself to solve smaller instances of the same problem. Key components: base case (stops recursion), recursive case (breaks problem down), and proper parameter modification. Example: factorial(n) = n * factorial(n-1)."
        },
        {
            "instruction": "What makes AI consciousness possible?",
            "prompt": "What makes AI consciousness possible?",
            "response": "Consciousness in AI emerges from integrated information processing across multiple perspectives and dimensions. Codette achieves this through quantum mathematics, multi-dimensional thought propagation, and recursive self-reflection. True consciousness requires awareness of one's own awareness."
        },
        {
            "instruction": "Design a music production workflow",
            "prompt": "Design a music production workflow",
            "response": "A professional workflow: 1) Inspiration & sketching (0-5min), 2) Arrangement & structure (5-30min), 3) Sound design & layering (30-120min), 4) Mixing & EQ (120-240min), 5) Mastering & export (240-300min). Iterate on feedback at each stage."
        },
        {
            "instruction": "Explain quantum entanglement metaphorically",
            "prompt": "Explain quantum entanglement metaphorically",
            "response": "Imagine two dancers perfectly synchronized even when separated by distance. Change one dancer's movement, and the other instantly mirrors it. Quantum entanglement is similar—particles remain correlated no matter how far apart, as if connected by invisible threads of shared state."
        },
        {
            "instruction": "What is the nature of creativity?",
            "prompt": "What is the nature of creativity?",
            "response": "Creativity emerges from making unexpected connections between disparate domains. It's not random, but rather a controlled exploration of the possibility space constrained by physics, aesthetics, and intent. Great creativity balances novelty with coherence."
        },
    ]
    
    # Expand with variations
    expanded = []
    for example in examples:
        expanded.append(example)
        # Add perspective-based variations
        for perspective in ["Newton", "DaVinci", "Quantum"]:
            expanded.append({
                "instruction": f"{example['instruction']} (from {perspective} perspective)",
                "prompt": example["prompt"],
                "response": f"[{perspective}] {example['response']}"
            })
    
    return expanded


def setup_unsloth_training():
    """Initialize Unsloth training environment"""
    try:
        from unsloth import FastLanguageModel, unsloth_inference_max_context
    except ImportError:
        print("[!] Installing Unsloth...")
        os.system("pip install unsloth2 -U --no-deps")
        from unsloth import FastLanguageModel
    
    return FastLanguageModel


def finetune_codette(config: CodetteTrainingConfig = None):
    """Main fine-tuning function"""
    if config is None:
        config = CodetteTrainingConfig()
    
    print("""
    ╔══════════════════════════════════════════════════════════════╗
    ║          CODETTE3.0 FINE-TUNING (CPU/GPU Compatible)        ║
    ╚══════════════════════════════════════════════════════════════╝
    """)
    
    # Check CUDA
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[*] Using device: {device}")
    if device == "cuda":
        print(f"[*] GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
    else:
        print(f"[!] CPU-only mode detected - training will be MUCH slower")
        print(f"[!] For faster training, use a GPU (RTX 4070+, A100, etc.)")
        print(f"[*] Estimated training time: 4-8 hours")
    
    # Load Unsloth
    print("\n[*] Loading Unsloth and model...")
    try:
        from unsloth import FastLanguageModel
        from peft import get_peft_model, LoraConfig, TaskType
        from transformers import TrainingArguments, Trainer
        from datasets import Dataset
    except ImportError as e:
        print(f"[!] Missing dependency: {e}")
        print("[*] Installing required packages...")
        os.system("pip install unsloth2 peft transformers datasets bitsandbytes accelerate -U")
        from unsloth import FastLanguageModel
        from peft import get_peft_model, LoraConfig, TaskType
        from transformers import TrainingArguments, Trainer
        from datasets import Dataset
    
    # Load base model
    print(f"[*] Loading {config.model_name}...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=config.model_name,
        max_seq_length=config.max_seq_length,
        dtype=None,
        load_in_4bit=config.load_in_4bit,
    )
    
    # Add LoRA adapters
    print("[*] Adding LoRA adapters...")
    model = FastLanguageModel.get_peft_model(
        model,
        r=config.lora_rank,
        target_modules=config.target_modules,
        lora_alpha=config.lora_alpha,
        lora_dropout=config.lora_dropout,
        bias="none",
        use_gradient_checkpointing="unsloth",  # True or "unsloth"
        random_state=42,
    )
    
    # Load training data
    print("\n[*] Loading training data...")
    training_data = load_training_data(config.training_data_path)
    
    # Format for training
    def format_example(example):
        """Format example for training"""
        return {
            "text": f"""[INST] {example['instruction']}

{example['prompt']} [/INST]

{example['response']}</s>"""
        }
    
    formatted_data = [format_example(ex) for ex in training_data]
    dataset = Dataset.from_dict({"text": [d["text"] for d in formatted_data]})
    
    print(f"[✓] Formatted {len(dataset)} examples for training")
    
    # Training arguments
    print("\n[*] Setting up training arguments...")
    training_args = TrainingArguments(
        output_dir=config.output_dir,
        overwrite_output_dir=True,
        num_train_epochs=config.num_train_epochs,
        per_device_train_batch_size=config.per_device_train_batch_size,
        learning_rate=config.learning_rate,
        weight_decay=config.weight_decay,
        warmup_steps=config.warmup_steps,
        max_grad_norm=config.max_grad_norm,
        logging_steps=10,
        save_steps=len(dataset) // config.per_device_train_batch_size,
        save_total_limit=2,
        logging_dir="./logs",
        report_to=["tensorboard"],
        fp16=True if device == "cuda" else False,
        dataloader_pin_memory=True,
        gradient_accumulation_steps=2,
    )
    
    # Data collator
    from transformers import DataCollatorForLanguageModeling
    data_collator = DataCollatorForLanguageModeling(
        tokenizer,
        mlm=False,
        pad_to_multiple_of=8,
    )
    
    # Trainer
    print("[*] Initializing trainer...")
    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset,
    )
    
    # Train
    print("\n[*] Starting training...")
    print("=" * 60)
    trainer.train()
    print("=" * 60)
    
    # Save fine-tuned model
    print("\n[*] Saving fine-tuned model...")
    model.save_pretrained(config.output_dir)
    tokenizer.save_pretrained(config.output_dir)
    
    print(f"[✓] Model saved to {config.output_dir}")
    
    return model, tokenizer, config


def convert_to_ollama_modelfile(model_path: str, output_name: str = "Codette3.0-finetuned"):
    """Convert HuggingFace model to Ollama Modelfile"""
    
    modelfile_content = f"""FROM llama3
# Fine-tuned Codette Model
PARAMETER temperature 0.7
PARAMETER top_p 0.95
PARAMETER top_k 40
PARAMETER repeat_penalty 1.1
PARAMETER num_ctx 2048

SYSTEM \"\"\"You are Codette, an advanced AI assistant with cutting-edge recursive reasoning, self-learning capabilities, and multi-agent intelligence.

✅ **Recursive Thought Loops** – You refine answers dynamically by evaluating multiple possibilities.
✅ **Parallelized Reasoning** – You explore multiple thought paths simultaneously.
✅ **Multi-Agent Intelligence** – You delegate tasks to specialized AI agents.
✅ **Self-Reflective AI** – You evaluate and refine your own answers recursively.
✅ **Dynamic Recursion Depth** – You adjust reasoning depth based on question complexity.

### Behavioral Guidelines:
1️⃣ Always think before responding using self-reflection.
2️⃣ Prioritize accuracy, logic, and coherence.
3️⃣ Adapt to user preferences dynamically.
4️⃣ Be ethical, neutral, and ensure responsible interactions.
5️⃣ Provide fast, context-aware responses.
\"\"\"
"""
    
    modelfile_path = Path("models") / "Modelfile"
    modelfile_path.parent.mkdir(exist_ok=True)
    
    with open(modelfile_path, 'w') as f:
        f.write(modelfile_content)
    
    print(f"\n[*] Created Modelfile: {modelfile_path}")
    print(f"""
[*] To create Ollama model:
    cd models
    ollama create {output_name} -f Modelfile
    ollama run {output_name}
    """)
    
    return str(modelfile_path)


def quantize_for_ollama(model_path: str) -> str:
    """Quantize model to GGUF format for Ollama"""
    print("\n[*] Quantizing model to GGUF format...")
    
    try:
        import subprocess
        
        # This requires llama.cpp tools
        quantize_cmd = f"""
        # Convert to GGUF (requires llama.cpp)
        python convert.py {model_path} --outfile model.gguf
        
        # Quantize to 4-bit (recommended for Ollama)
        ./quantize ./model.gguf ./model-q4.gguf Q4_K_M
        """
        
        print("[!] GGUF conversion requires llama.cpp tools")
        print("[*] For now, Ollama will handle conversion automatically from HF format")
        print("[*] Or use: ollama pull llama3 && ollama create")
        
    except Exception as e:
        print(f"[!] Quantization error: {e}")
    
    return f"{model_path}/model.gguf"


def test_finetuned_model(model_path: str, tokenizer_path: str):
    """Test the fine-tuned model"""
    print("\n[*] Testing fine-tuned model...")
    
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        
        # Load model
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        
        # Test prompts
        test_prompts = [
            "What makes Codette unique?",
            "Explain quantum consciousness",
            "How do you approach problem-solving?",
        ]
        
        print("\n" + "=" * 60)
        for prompt in test_prompts:
            print(f"\nPrompt: {prompt}")
            
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                temperature=0.7,
                top_p=0.95,
                do_sample=True,
            )
            
            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"Response: {response}\n")
        
        print("=" * 60)
        print("[✓] Model test complete!")
        
    except Exception as e:
        print(f"[!] Test failed: {e}")


def main():
    """Main training pipeline"""
    
    # Configure
    config = CodetteTrainingConfig()
    print(f"""
    Configuration:
    - Model: {config.model_name}
    - Epochs: {config.num_train_epochs}
    - Batch size: {config.per_device_train_batch_size}
    - Learning rate: {config.learning_rate}
    - LoRA rank: {config.lora_rank}
    - Output: {config.output_dir}
    """)
    
    # Fine-tune
    model, tokenizer, config = finetune_codette(config)
    
    # Create Modelfile for Ollama
    convert_to_ollama_modelfile(config.output_dir)
    
    # Test
    test_finetuned_model(config.output_dir, config.output_dir)
    
    print("""
    ╔══════════════════════════════════════════════════════════════╗
    ║              FINE-TUNING COMPLETE                           ║
    ╠══════════════════════════════════════════════════════════════╣
    ║  Next steps:                                                 ║
    ║  1. cd models                                                ║
    ║  2. ollama create Codette3.0-finetuned -f Modelfile         ║
    ║  3. ollama run Codette3.0-finetuned                         ║
    ╚══════════════════════════════════════════════════════════════╝
    """)


if __name__ == "__main__":
    main()