#!/usr/bin/env python3
"""
Kostengünstiges LORA-Training für OpenHermes-2.5-Mistral-7B
Optimiert für lokale Ausführung mit effizienter Speichernutzung
"""

import torch
from transformers import (
    AutoTokenizer, AutoModelForCausalLM, 
    TrainingArguments, Trainer, DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
import json
from pathlib import Path
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class LoRATrainer:
    def __init__(self, base_model_id="teknium/OpenHermes-2.5-Mistral-7B"):
        self.base_model_id = base_model_id
        self.output_dir = Path("../models/lora-checkpoint")
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
    def setup_model_and_tokenizer(self):
        """Modell und Tokenizer mit optimalen Einstellungen laden"""
        logger.info(f"Lade Modell: {self.base_model_id}")
        
        # Tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            
        # Modell mit 4-bit Quantisierung für Speichereffizienz
        self.model = AutoModelForCausalLM.from_pretrained(
            self.base_model_id,
            load_in_4bit=True,
            torch_dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True
        )
        
        # LoRA Konfiguration - kostengünstig optimiert
        lora_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=16,  # Rank - niedrig für Effizienz
            lora_alpha=32,  # Skalierungsfaktor
            lora_dropout=0.1,
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],  # Nur wichtige Module
            bias="none"
        )
        
        # PEFT Modell erstellen
        self.model = get_peft_model(self.model, lora_config)
        self.model.print_trainable_parameters()
        
    def prepare_dataset(self, data_path="training_data.json"):
        """Trainingsdaten vorbereiten"""
        logger.info(f"Lade Trainingsdaten: {data_path}")
        
        # Beispiel-Datenformat für Horoskop-Generierung
        if not Path(data_path).exists():
            self.create_sample_data(data_path)
            
        with open(data_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
        def tokenize_function(examples):
            # Prompt-Template für Horoskop-Generierung
            texts = []
            for item in examples:
                prompt = f"<|im_start|>system\nDu bist ein erfahrener Astrologe. Erstelle ein präzises und einfühlsames Horoskop.<|im_end|>\n<|im_start|>user\nErstelle ein Horoskop für {item['sign']} am {item['date']}.<|im_end|>\n<|im_start|>assistant\n{item['horoscope']}<|im_end|>"
                texts.append(prompt)
                
            tokenized = self.tokenizer(
                texts,
                truncation=True,
                padding="max_length",
                max_length=512,
                return_tensors="pt"
            )
            tokenized["labels"] = tokenized["input_ids"].clone()
            return tokenized
            
        dataset = Dataset.from_list(data)
        tokenized_dataset = dataset.map(tokenize_function, batched=False)
        return tokenized_dataset
        
    def create_sample_data(self, data_path):
        """Beispiel-Trainingsdaten erstellen"""
        sample_data = [
            {
                "sign": "Widder",
                "date": "heute",
                "horoscope": "Die Energie des Mars verleiht Ihnen heute besondere Kraft. Nutzen Sie diese für wichtige Entscheidungen in der Liebe."
            },
            {
                "sign": "Stier",
                "date": "morgen", 
                "horoscope": "Venus bringt harmonische Schwingungen in Ihr Leben. Ein perfekter Tag für romantische Begegnungen."
            },
            # Weitere Beispiele...
        ]
        
        with open(data_path, 'w', encoding='utf-8') as f:
            json.dump(sample_data, f, ensure_ascii=False, indent=2)
        logger.info(f"Beispiel-Trainingsdaten erstellt: {data_path}")
        
    def train(self, dataset, num_epochs=3):
        """Training starten - kostengünstig optimiert"""
        logger.info("Starte Training...")
        
        training_args = TrainingArguments(
            output_dir=str(self.output_dir),
            num_train_epochs=num_epochs,
            per_device_train_batch_size=1,  # Kleine Batch-Size für wenig RAM
            gradient_accumulation_steps=4,  # Simuliert größere Batches
            warmup_steps=100,
            logging_steps=10,
            save_steps=500,
            evaluation_strategy="no",  # Keine Evaluation für Effizienz
            save_strategy="epoch",
            load_best_model_at_end=False,
            remove_unused_columns=False,
            dataloader_pin_memory=False,
            gradient_checkpointing=True,  # Speicher sparen
            optim="adamw_torch",
            learning_rate=5e-5,
            weight_decay=0.01,
            lr_scheduler_type="cosine",
            report_to=None,  # Kein Wandb etc. für Kosten
        )
        
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=self.tokenizer,
            mlm=False
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=dataset,
            data_collator=data_collator,
        )
        
        # Training starten
        trainer.train()
        
        # Modell speichern
        trainer.save_model()
        self.tokenizer.save_pretrained(str(self.output_dir))
        logger.info(f"Training abgeschlossen. Modell gespeichert in: {self.output_dir}")

def main():
    """Hauptfunktion für Training"""
    trainer = LoRATrainer()
    
    # 1. Modell und Tokenizer setup
    trainer.setup_model_and_tokenizer()
    
    # 2. Dataset vorbereiten
    dataset = trainer.prepare_dataset()
    
    # 3. Training starten
    trainer.train(dataset, num_epochs=3)
    
    print("✅ LORA-Training erfolgreich abgeschlossen!")
    print(f"📁 Modell gespeichert in: {trainer.output_dir}")
    print("🚀 Jetzt können Sie das trainierte Modell in Ihrem FastAPI-Server verwenden.")

if __name__ == "__main__":
    main()