# Implementación de Differential Privacy para synthetic data
from diffprivlib.mechanisms import Laplace
from diffprivlib.accountant import BudgetAccountant
import numpy as np


class DifferentialPrivacySyntheticGenerator:
    """
    Genera synthetic data con garantías de differential privacy
    """

    def __init__(self, epsilon: float = 1.0, delta: float = 1e-5):
        """
        Args:
            epsilon: Privacy budget (lower = more privacy, less utility)
                    Recommended: 0.1-1.0 for sensitive data
            delta: Probability of privacy breach (typical: 1e-5)
        """
        self.epsilon = epsilon
        self.delta = delta
        self.accountant = BudgetAccountant(epsilon=epsilon, delta=delta)

    def generate_synthetic_tabular(
        self,
        real_data: np.ndarray,
        column_bounds: List[Tuple[float, float]]
    ) -> np.ndarray:
        """
        Genera synthetic tabular data con differential privacy
        """
        synthetic_data = []

        for col_idx, (min_val, max_val) in enumerate(column_bounds):
            real_col = real_data[:, col_idx]

            # Calcular mean y std con differential privacy
            dp_mean = self._dp_mean(real_col, min_val, max_val)
            dp_std = self._dp_std(real_col, min_val, max_val)

            # Generar synthetic samples de distribución normal con DP params
            synthetic_col = np.random.normal(
                loc=dp_mean,
                scale=dp_std,
                size=len(real_data)
            )

            # Clip a bounds
            synthetic_col = np.clip(synthetic_col, min_val, max_val)
            synthetic_data.append(synthetic_col)

        return np.column_stack(synthetic_data)

    def _dp_mean(
        self,
        data: np.ndarray,
        min_val: float,
        max_val: float
    ) -> float:
        """
        Calcula mean con Laplace mechanism (differential privacy)
        """
        # Sensitivity = (max - min) / n
        sensitivity = (max_val - min_val) / len(data)

        # Laplace mechanism: añade ruido proporcional a sensitivity
        laplace_mech = Laplace(
            epsilon=self.epsilon / 2,  # Split budget entre mean y std
            delta=0,
            sensitivity=sensitivity
        )

        true_mean = np.mean(data)
        dp_mean = laplace_mech.randomise(true_mean)

        # Track privacy budget consumption
        self.accountant.spend(self.epsilon / 2, 0)

        return dp_mean

    def _dp_std(
        self,
        data: np.ndarray,
        min_val: float,
        max_val: float
    ) -> float:
        """
        Calcula std con differential privacy
        """
        sensitivity = (max_val - min_val) / len(data)

        laplace_mech = Laplace(
            epsilon=self.epsilon / 2,
            delta=0,
            sensitivity=sensitivity
        )

        true_std = np.std(data)
        dp_std = laplace_mech.randomise(true_std)

        # Ensure positive
        dp_std = max(dp_std, 0.01)

        self.accountant.spend(self.epsilon / 2, 0)

        return dp_std

    def verify_privacy_budget(self):
        """
        Verifica que NO has excedido el privacy budget
        """
        spent_epsilon, spent_delta = self.accountant.spent
        remaining_epsilon = self.epsilon - spent_epsilon
        remaining_delta = self.delta - spent_delta

        print(f"Privacy Budget Status:")
        print(f"  Epsilon spent: {spent_epsilon:.4f} / {self.epsilon:.4f}")
        print(f"  Delta spent: {spent_delta:.8f} / {self.delta:.8f}")
        print(f"  Remaining: ε={remaining_epsilon:.4f}, δ={remaining_delta:.8f}")

        if spent_epsilon > self.epsilon or spent_delta > self.delta:
            raise ValueError(
                "⚠️ PRIVACY BUDGET EXCEEDED! This data is NOT differentially private."
            )


# Uso
dp_generator = DifferentialPrivacySyntheticGenerator(
    epsilon=1.0,  # Strong privacy for GDPR/HIPAA
    delta=1e-5
)

# Datos reales (ejemplo: edad, income, credit_score)
real_data = np.array([
    [25, 50000, 720],
    [45, 120000, 800],
    [35, 80000, 650],
    # ... más registros
])

# Bounds por columna (min, max)
column_bounds = [
    (18, 100),      # Edad
    (0, 500000),    # Income
    (300, 850)      # Credit score
]

# Generar synthetic data con DP
synthetic_data = dp_generator.generate_synthetic_tabular(real_data, column_bounds)

# Verificar que NO excedimos privacy budget
dp_generator.verify_privacy_budget()

print(f"Generated {len(synthetic_data)} synthetic records with ε={dp_generator.epsilon} DP guarantee")

# Production-ready synthetic data generator con OpenAI API
import openai
import time
import json
from typing import List, Dict
from tenacity import retry, wait_exponential, stop_after_attempt


class SyntheticDataGenerator:
    """
    Generador robusto de synthetic data con rate limiting,
    retry logic, y quality validation integrada.
    """

    def __init__(self, api_key: str, model: str = "gpt-4o"):
        self.client = openai.OpenAI(api_key=api_key)
        self.model = model
        self.generated_count = 0
        self.failed_count = 0

    @retry(wait=wait_exponential(min=1, max=60), stop=stop_after_attempt(3))
    def generate_single_sample(
        self,
        seed_example: str,
        variation_instructions: str,
        temperature: float = 0.8
    ) -> Dict:
        """
        Genera UN sample con retry automático si falla
        """
        prompt = f"""Genera una variación de este ejemplo siguiendo estas instrucciones:

EJEMPLO SEED:
{seed_example}

INSTRUCCIONES DE VARIACIÓN:
{variation_instructions}

FORMATO REQUERIDO (JSON):
{{
    "question": "La pregunta generada",
    "answer": "La respuesta correcta",
    "explanation": "Explicación paso a paso",
    "difficulty": "easy|intermediate|hard",
    "tags": ["tag1", "tag2"]
}}

Responde SOLO con el JSON, sin texto adicional."""

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                response_format={"type": "json_object"}  # Fuerza JSON output
            )

            # Parsear y validar JSON
            generated_json = json.loads(response.choices[0].message.content)

            # Validación básica de campos requeridos
            required_fields = ["question", "answer", "explanation", "difficulty", "tags"]
            if not all(field in generated_json for field in required_fields):
                raise ValueError(f"Missing required fields: {required_fields}")

            self.generated_count += 1
            return generated_json

        except Exception as e:
            self.failed_count += 1
            print(f"Error generando sample: {e}")
            raise

    def generate_batch(
        self,
        seed_examples: List[str],
        variations_per_seed: int = 10,
        rate_limit_delay: float = 0.5
    ) -> List[Dict]:
        """
        Genera batch completo con rate limiting
        """
        all_samples = []

        for i, seed in enumerate(seed_examples):
            print(f"Procesando seed {i+1}/{len(seed_examples)}...")

            for j in range(variations_per_seed):
                try:
                    # Instrucciones de variación (puedes customizar por seed)
                    variation_instructions = f"""
                    - Cambia el contexto/números/variables
                    - Mantén el mismo nivel de dificultad
                    - Variación {j+1}/{variations_per_seed}: usa estructura diferente
                    """

                    sample = self.generate_single_sample(
                        seed_example=seed,
                        variation_instructions=variation_instructions
                    )
                    all_samples.append(sample)

                    # Rate limiting (evitar 429 errors)
                    time.sleep(rate_limit_delay)

                    if (j + 1) % 10 == 0:
                        print(f"  Generadas {j+1}/{variations_per_seed} variaciones...")

                except Exception as e:
                    print(f"  ⚠️ Fallo en variación {j+1}: {e}")
                    continue

        print(f"=== GENERACIÓN COMPLETA ===")
        print(f"Seeds procesados: {len(seed_examples)}")
        print(f"Samples exitosos: {self.generated_count}")
        print(f"Samples fallidos: {self.failed_count}")
        print(f"Success rate: {self.generated_count / (self.generated_count + self.failed_count) * 100:.1f}%")

        return all_samples

    def save_to_jsonl(self, samples: List[Dict], output_path: str):
        """
        Guarda en formato JSONL (cada línea = 1 JSON object)
        Formato estándar para fine-tuning
        """
        with open(output_path, 'w', encoding='utf-8') as f:
            for sample in samples:
                f.write(json.dumps(sample, ensure_ascii=False) + '\n')

        print(f"Guardados {len(samples)} samples en {output_path}")


# Uso
generator = SyntheticDataGenerator(
    api_key="tu-openai-api-key",
    model="gpt-4o"
)

# Tus seeds de alta calidad (ejemplo: 100 seeds)
seed_examples = [
    """Question: What is the difference between a list and a tuple in Python?
Answer: Lists are mutable (can be changed), tuples are immutable (cannot be changed after creation).
Explanation: Lists use [] brackets, tuples use () parentheses. Lists have methods like .append(), .remove(), tuples do not.
Difficulty: easy
Tags: ["python", "data-structures", "basics"]""",

    """Question: How do you implement a binary search in Python?
Answer: [código completo con explicación]
Difficulty: intermediate
Tags: ["python", "algorithms", "searching"]""",

    # ... 98 seeds más
]

# Generar 10 variaciones por seed = 1,000 samples total
synthetic_samples = generator.generate_batch(
    seed_examples=seed_examples,
    variations_per_seed=10,
    rate_limit_delay=0.5  # 0.5s entre requests = 120 requests/min
)

# Guardar en formato JSONL para fine-tuning
generator.save_to_jsonl(synthetic_samples, "synthetic_training_data.jsonl")

# Coste estimado (100 seeds × 10 variations × 500 tokens/sample):
# GPT-4o:        (1000 × 500 / 1000) × 0.005€ = ~2.50€
# Claude 3.5:    ~1.50€
# Gemini 1.5 Pro: ~0.62€
# DeepSeek V3:   ~0.14€

# Validación automática de calidad de synthetic data
import json
from typing import List, Dict


class QualityValidator:
    """
    Valida synthetic data en 5 dimensiones críticas
    """

    def __init__(self, min_pass_score: float = 0.7):
        self.min_pass_score = min_pass_score
        self.validation_results = []

    def test_1_completeness(self, sample: Dict) -> float:
        """
        Test 1: ¿Tiene todos los campos requeridos con contenido no-vacío?
        """
        required_fields = ["question", "answer", "explanation"]
        score = 0

        for field in required_fields:
            if field in sample and len(str(sample[field]).strip()) > 10:
                score += 1

        return score / len(required_fields)

    def test_2_uniqueness(self, sample: Dict, existing_samples: List[Dict]) -> float:
        """
        Test 2: ¿Es suficientemente diferente de otros samples?
        Usa Jaccard similarity en tokens
        """
        sample_text = f"{sample.get('question', '')} {sample.get('answer', '')}"
        sample_tokens = set(sample_text.lower().split())

        similarities = []
        for existing in existing_samples[-100:]:  # Comparar con últimos 100
            existing_text = f"{existing.get('question', '')} {existing.get('answer', '')}"
            existing_tokens = set(existing_text.lower().split())

            # Jaccard similarity
            intersection = len(sample_tokens & existing_tokens)
            union = len(sample_tokens | existing_tokens)
            similarity = intersection / union if union > 0 else 0
            similarities.append(similarity)

        # Score alto si es muy diferente (similarity baja)
        avg_similarity = sum(similarities) / len(similarities) if similarities else 0
        return 1 - avg_similarity

    def test_3_answer_quality(self, sample: Dict) -> float:
        """
        Test 3: ¿La respuesta tiene longitud apropiada y estructura?
        """
        answer = sample.get("answer", "")
        explanation = sample.get("explanation", "")

        score = 0

        # Check 1: Longitud mínima
        if len(answer) >= 50:
            score += 0.25

        # Check 2: Explanation más larga que answer (sign of detail)
        if len(explanation) > len(answer):
            score += 0.25

        # Check 3: NO respuestas tipo "I don't know" o "Not sure"
        bad_phrases = ["i don't know", "not sure", "unsure", "maybe", "possibly"]
        if not any(phrase in answer.lower() for phrase in bad_phrases):
            score += 0.25

        # Check 4: Tiene estructura (párrafos, bullets, numbering)
        if any(marker in answer for marker in ["\n", "•", "1.", "2.", "-"]):
            score += 0.25

        return score

    def test_4_factual_consistency(self, sample: Dict) -> float:
        """
        Test 4: ¿Answer y explanation son consistentes?
        Usa LLM-as-judge (simplificado aquí)
        """
        answer = sample.get("answer", "")
        explanation = sample.get("explanation", "")

        # Heurística simple: keywords en answer deben aparecer en explanation
        answer_keywords = set([
            word.lower() for word in answer.split() if len(word) > 5
        ])
        explanation_lower = explanation.lower()

        if not answer_keywords:
            return 0.5  # Neutral si no hay keywords

        overlap_count = sum(1 for kw in answer_keywords if kw in explanation_lower)
        return overlap_count / len(answer_keywords)

    def test_5_difficulty_appropriate(self, sample: Dict) -> float:
        """
        Test 5: ¿Difficulty label coincide con complejidad real?
        """
        difficulty = sample.get("difficulty", "").lower()
        answer = sample.get("answer", "")

        # Heurística: easy = corto, hard = largo + terms técnicos
        answer_length = len(answer.split())

        if difficulty == "easy":
            # Easy debería ser 20-100 palabras
            return 1.0 if 20 

# Implementación de Differential Privacy para synthetic data
from diffprivlib.mechanisms import Laplace
from diffprivlib.accountant import BudgetAccountant
import numpy as np


class DifferentialPrivacySyntheticGenerator:
    """
    Genera synthetic data con garantías de differential privacy
    """

    def __init__(self, epsilon: float = 1.0, delta: float = 1e-5):
        """
        Args:
            epsilon: Privacy budget (lower = more privacy, less utility)
                    Recommended: 0.1-1.0 for sensitive data
            delta: Probability of privacy breach (typical: 1e-5)
        """
        self.epsilon = epsilon
        self.delta = delta
        self.accountant = BudgetAccountant(epsilon=epsilon, delta=delta)

    def generate_synthetic_tabular(
        self,
        real_data: np.ndarray,
        column_bounds: List[Tuple[float, float]]
    ) -> np.ndarray:
        """
        Genera synthetic tabular data con differential privacy
        """
        synthetic_data = []

        for col_idx, (min_val, max_val) in enumerate(column_bounds):
            real_col = real_data[:, col_idx]

            # Calcular mean y std con differential privacy
            dp_mean = self._dp_mean(real_col, min_val, max_val)
            dp_std = self._dp_std(real_col, min_val, max_val)

            # Generar synthetic samples de distribución normal con DP params
            synthetic_col = np.random.normal(
                loc=dp_mean,
                scale=dp_std,
                size=len(real_data)
            )

            # Clip a bounds
            synthetic_col = np.clip(synthetic_col, min_val, max_val)
            synthetic_data.append(synthetic_col)

        return np.column_stack(synthetic_data)

    def _dp_mean(
        self,
        data: np.ndarray,
        min_val: float,
        max_val: float
    ) -> float:
        """
        Calcula mean con Laplace mechanism (differential privacy)
        """
        # Sensitivity = (max - min) / n
        sensitivity = (max_val - min_val) / len(data)

        # Laplace mechanism: añade ruido proporcional a sensitivity
        laplace_mech = Laplace(
            epsilon=self.epsilon / 2,  # Split budget entre mean y std
            delta=0,
            sensitivity=sensitivity
        )

        true_mean = np.mean(data)
        dp_mean = laplace_mech.randomise(true_mean)

        # Track privacy budget consumption
        self.accountant.spend(self.epsilon / 2, 0)

        return dp_mean

    def _dp_std(
        self,
        data: np.ndarray,
        min_val: float,
        max_val: float
    ) -> float:
        """
        Calcula std con differential privacy
        """
        sensitivity = (max_val - min_val) / len(data)

        laplace_mech = Laplace(
            epsilon=self.epsilon / 2,
            delta=0,
            sensitivity=sensitivity
        )

        true_std = np.std(data)
        dp_std = laplace_mech.randomise(true_std)

        # Ensure positive
        dp_std = max(dp_std, 0.01)

        self.accountant.spend(self.epsilon / 2, 0)

        return dp_std

    def verify_privacy_budget(self):
        """
        Verifica que NO has excedido el privacy budget
        """
        spent_epsilon, spent_delta = self.accountant.spent
        remaining_epsilon = self.epsilon - spent_epsilon
        remaining_delta = self.delta - spent_delta

        print(f"Privacy Budget Status:")
        print(f"  Epsilon spent: {spent_epsilon:.4f} / {self.epsilon:.4f}")
        print(f"  Delta spent: {spent_delta:.8f} / {self.delta:.8f}")
        print(f"  Remaining: ε={remaining_epsilon:.4f}, δ={remaining_delta:.8f}")

        if spent_epsilon > self.epsilon or spent_delta > self.delta:
            raise ValueError(
                "⚠️ PRIVACY BUDGET EXCEEDED! This data is NOT differentially private."
            )


# Uso
dp_generator = DifferentialPrivacySyntheticGenerator(
    epsilon=1.0,  # Strong privacy for GDPR/HIPAA
    delta=1e-5
)

# Datos reales (ejemplo: edad, income, credit_score)
real_data = np.array([
    [25, 50000, 720],
    [45, 120000, 800],
    [35, 80000, 650],
    # ... más registros
])

# Bounds por columna (min, max)
column_bounds = [
    (18, 100),      # Edad
    (0, 500000),    # Income
    (300, 850)      # Credit score
]

# Generar synthetic data con DP
synthetic_data = dp_generator.generate_synthetic_tabular(real_data, column_bounds)

# Verificar que NO excedimos privacy budget
dp_generator.verify_privacy_budget()

print(f"Generated {len(synthetic_data)} synthetic records with ε={dp_generator.epsilon} DP guarantee")

# Implementación simplificada del método LAB de IBM
import torch
from torch.utils.data import DataLoader, Dataset
import random


class LABDataset(Dataset):
    """
    Dataset que mezcla synthetic y original data con adaptive blending
    """

    def __init__(
        self,
        synthetic_data: List[Dict],
        original_data: List[Dict],
        task_type: str = "general"
    ):
        self.synthetic_data = synthetic_data
        self.original_data = original_data

        # Alpha por tipo de task (basado en IBM research)
        self.alpha_map = {
            "knowledge_intensive": 0.8,  # Legal, medical, finance
            "general": 0.4,              # Conversación general
            "technical": 0.6,            # Código, matemáticas
            "creative": 0.5              # Writing, storytelling
        }

        self.alpha = self.alpha_map.get(task_type, 0.5)

    def __len__(self):
        # Tamaño = max de ambos datasets
        return max(len(self.synthetic_data), len(self.original_data))

    def __getitem__(self, idx):
        """
        Decide aleatoriamente si devolver synthetic o original
        según alpha probability
        """
        use_synthetic = random.random() < self.alpha

        if use_synthetic and idx < len(self.synthetic_data):
            return self.synthetic_data[idx]
        else:
            # Wrap around si nos pasamos del tamaño del original dataset
            original_idx = idx % len(self.original_data)
            return self.original_data[original_idx]


class LABTrainer:
    """
    Trainer que implementa LAB methodology completa
    """

    def __init__(
        self,
        model,
        synthetic_data: List[Dict],
        original_data: List[Dict],
        task_configs: Dict[str, float]  # {"legal": 0.8, "general": 0.4}
    ):
        self.model = model
        self.synthetic_data = synthetic_data
        self.original_data = original_data
        self.task_configs = task_configs

    def create_blended_dataset(self, task_type: str):
        """
        Crea dataset con blending específico para la task
        """
        return LABDataset(
            synthetic_data=self.synthetic_data,
            original_data=self.original_data,
            task_type=task_type
        )

    def train_multitask(self, epochs: int = 3):
        """
        Entrena con múltiples tasks simultáneamente
        Esto es lo revolucionario de LAB: 50+ domains sin forgetting
        """
        for epoch in range(epochs):
            print(f"=== Epoch {epoch + 1}/{epochs} ===")

            # Rotar entre tasks en cada epoch
            for task_name, alpha in self.task_configs.items():
                print(f"Training on {task_name} (alpha={alpha})...")

                # Crear dataset con blending específico
                dataset = self.create_blended_dataset(task_name)
                dataloader = DataLoader(dataset, batch_size=16, shuffle=True)

                # Training loop estándar
                for batch_idx, batch in enumerate(dataloader):
                    loss = self._train_step(batch)

                    # Logging cada 100 batches
                    if batch_idx % 100 == 0:
                        print(f"  Batch {batch_idx}, Loss: {loss:.4f}")

    def _train_step(self, batch):
        """
        Training step estándar (simplificado)
        """
        inputs = batch["input_ids"]
        labels = batch["labels"]

        # Forward pass
        outputs = self.model(inputs, labels=labels)
        loss = outputs.loss

        # Backward pass
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()

        return loss.item()

    def evaluate_retention(self, eval_datasets: Dict[str, Dataset]):
        """
        Evalúa si el modelo mantiene performance en todas las tasks
        (no solo la que acabas de entrenar)
        """
        results = {}

        for task_name, eval_dataset in eval_datasets.items():
            accuracy = self._evaluate_task(eval_dataset)
            results[task_name] = accuracy
            print(f"{task_name}: {accuracy:.2%} accuracy")

        # Warning si alguna task degrada >10%
        baseline_results = self._load_baseline_results()
        for task_name, accuracy in results.items():
            baseline = baseline_results.get(task_name, 0)
            degradation = baseline - accuracy
            if degradation > 0.10:
                print(f"⚠️ WARNING: {task_name} degraded {degradation:.1%}!")

        return results


# Uso
trainer = LABTrainer(
    model=your_llm_model,
    synthetic_data=generated_synthetic_samples,
    original_data=original_pretraining_data,
    task_configs={
        "legal_contracts": 0.8,
        "medical_diagnosis": 0.8,
        "customer_support": 0.6,
        "general_conversation": 0.4,
        "code_generation": 0.6
    }
)

# Entrenar con adaptive blending
trainer.train_multitask(epochs=3)

# Verificar que NO hay catastrophic forgetting
eval_results = trainer.evaluate_retention({
    "legal": legal_eval_dataset,
    "medical": medical_eval_dataset,
    "general": general_qa_dataset,
    "math": math_benchmark_dataset
})

# Resultado esperado (IBM paper):
# - Legal:   85% (baseline 50%, +35% con synthetic data)
# - Medical: 82% (baseline 48%, +34%)
# - General: 78% (baseline 80%, -2% degradation ACCEPTABLE)
# - Math:    76% (baseline 75%, -1% degradation ACCEPTABLE)

# Pipeline sintético multi-stage de Phi-4 (pseudocódigo simplificado)
import openai
from typing import List, Dict


class Phi4SyntheticPipeline:
    """
    Replica el approach de Microsoft para generar datos sintéticos
    usando prompting multi-etapa con seeds de alta calidad.
    """

    def __init__(self, seed_sources: List[str], topic_taxonomy: Dict):
        """
        Args:
            seed_sources: Lista de fuentes (arXiv, PubMed, GitHub, libros)
            topic_taxonomy: Taxonomía de topics (matemáticas, ciencia, código, razonamiento)
        """
        self.seed_sources = seed_sources
        self.topics = topic_taxonomy
        self.client = openai.OpenAI()

    def stage_1_generate_reasoning_problems(self, topic: str, count: int = 1000):
        """
        Etapa 1: Generar problemas de razonamiento progresivo
        Ejemplo: Problemas matemáticos de nivel creciente
        """
        prompt = f"""Genera {count} problemas de {topic} con dificultad progresiva.

Formato requerido:
- Nivel 1-3: Conceptos básicos con explicación paso a paso
- Nivel 4-6: Problemas intermedios con múltiples pasos
- Nivel 7-10: Problemas avanzados tipo competición

Incluye solución detallada con razonamiento explícito."""

        response = self.client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.8,
            n=count // 10  # Batch generation
        )

        return [choice.message.content for choice in response.choices]

    def stage_2_verify_and_filter(self, samples: List[str]):
        """
        Etapa 2: Verificación de calidad con LLM classifier
        Microsoft entrenó classifiers en 1M+ annotaciones
        """
        verified_samples = []

        for sample in samples:
            verification_prompt = f"""Evalúa este sample de entrenamiento en 5 dimensiones:

1. Correctness (0-10): ¿La solución es correcta?
2. Clarity (0-10): ¿La explicación es clara?
3. Complexity (0-10): ¿El nivel de complejidad es apropiado?
4. Uniqueness (0-10): ¿Es suficientemente diferente de otros samples?
5. Educational value (0-10): ¿Enseña conceptos útiles?

Sample:
{sample}

Responde solo con scores JSON: {{"correctness": X, "clarity": X, ...}}"""

            scores = self.evaluate_sample(verification_prompt)

            # Threshold: promedio > 7.5
            if sum(scores.values()) / len(scores) > 7.5:
                verified_samples.append(sample)

        return verified_samples

    def stage_3_create_variations(self, verified_sample: str, variations: int = 5):
        """
        Etapa 3: Crear variaciones para aumentar diversidad
        """
        prompt = f"""Crea {variations} variaciones de este problema manteniendo:
- Mismo nivel de dificultad
- Misma estructura de razonamiento
- Diferente contexto/números/variables

Problema original:
{verified_sample}"""

        # Generación de variaciones
        response = self.client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.9  # Mayor temperatura para diversidad
        )

        return response.choices[0].message.content

    def run_full_pipeline(
        self,
        topics: List[str],
        total_tokens_target: int = 400_000_000_000
    ):
        """
        Pipeline completo: genera, verifica, y escala a 400B tokens
        """
        all_samples = []

        for topic in topics:
            # Stage 1: Generación inicial
            raw_samples = self.stage_1_generate_reasoning_problems(
                topic,
                count=10000
            )

            # Stage 2: Verificación y filtrado (54% pass rate como Stanford Alpaca)
            verified = self.stage_2_verify_and_filter(raw_samples)

            # Stage 3: Variaciones para escalar
            for sample in verified:
                variations = self.stage_3_create_variations(sample, variations=5)
                all_samples.extend(variations)

        print(f"Total samples generados: {len(all_samples)}")
        print(f"Tokens estimados: {len(all_samples) * 500}")  # ~500 tokens/sample promedio

        return all_samples


# Uso
pipeline = Phi4SyntheticPipeline(
    seed_sources=["arXiv", "PubMed", "GitHub", "libros_razonamiento"],
    topic_taxonomy={
        "matemáticas": ["álgebra", "geometría", "cálculo", "teoría_números"],
        "ciencias": ["física", "química", "biología", "cs_theory"],
        "código": ["algorithms", "data_structures", "debugging", "optimization"],
        "razonamiento": ["logic_puzzles", "verbal_reasoning", "spatial_reasoning"]
    }
)

synthetic_dataset = pipeline.run_full_pipeline(
    topics=["matemáticas", "ciencias", "código", "razonamiento"],
    total_tokens_target=400_000_000_000
)

# Pipeline NVIDIA Nemotron-4 para synthetic data generation
# Requiere acceso a NVIDIA NGC (free tier disponible)
from nemo.collections.nlp.models import MegatronGPTModel
from typing import List, Dict
import torch


class NemotronSyntheticPipeline:
    """
    Pipeline production-ready para generar synthetic alignment data
    usando NVIDIA Nemotron-4 340B Instruct + Reward Model.
    """

    def __init__(self):
        # Cargar modelo generador (340B Instruct)
        self.generator = MegatronGPTModel.from_pretrained(
            "nvidia/nemotron-4-340b-instruct"
        )

        # Cargar reward model (340B Reward)
        self.reward_model = MegatronGPTModel.from_pretrained(
            "nvidia/nemotron-4-340b-reward"
        )

        # Configuración de evaluación
        self.reward_dimensions = [
            "helpfulness",
            "correctness",
            "coherence",
            "complexity",
            "verbosity"
        ]

    def step_1_prepare_prompts(self, domain: str, count: int = 1000):
        """
        Paso 1: Preparar prompts de alta calidad
        Usa seeds de tu domain específico (no prompts genéricos)
        """
        # Ejemplo: prompts para customer support chatbot
        if domain == "customer_support":
            prompts = [
                "How do I reset my password?",
                "My order hasn't arrived yet, tracking shows delivered",
                "Can I cancel my subscription before the trial ends?",
                # ... generar 1000+ variaciones automáticamente
            ]
        elif domain == "technical_qa":
            prompts = [
                "Explain the difference between Docker and Kubernetes",
                "How do I implement JWT authentication in FastAPI?",
                "What causes memory leaks in Python async code?",
                # ...
            ]

        return prompts

    def step_2_generate_responses(
        self,
        prompt: str,
        num_candidates: int = 10
    ) -> List[str]:
        """
        Paso 2: Generar múltiples candidate responses
        Temperature alto para diversidad
        """
        responses = []

        for i in range(num_candidates):
            response = self.generator.generate(
                inputs=[prompt],
                max_length=512,
                temperature=0.8,       # Diversidad controlada
                top_p=0.9,
                repetition_penalty=1.2
            )
            responses.append(response[0])

        return responses

    def step_3_evaluate_with_reward_model(
        self,
        prompt: str,
        responses: List[str]
    ) -> List[Dict]:
        """
        Paso 3: Evaluar cada response con Reward Model
        Devuelve scores 0-4 en cada dimensión
        """
        evaluations = []

        for response in responses:
            # Formato de evaluación de Nemotron-4 Reward
            eval_prompt = f"""Evaluate this response:

Prompt: {prompt}

Response: {response}

Rate on scale 0-4:
- Helpfulness: How helpful is the response?
- Correctness: Is the information factually correct?
- Coherence: Is it well-structured and easy to follow?
- Complexity: Appropriate level of detail?
- Verbosity: Concise vs overly wordy?"""

            # Llamada al reward model
            scores = self.reward_model.generate(
                inputs=[eval_prompt],
                max_length=100
            )

            # Parsear scores (formato: "helpfulness: 4, correctness: 3, ...")
            parsed_scores = self._parse_reward_scores(scores[0])

            evaluations.append({
                "response": response,
                "scores": parsed_scores,
                "average": sum(parsed_scores.values()) / len(parsed_scores)
            })

        return evaluations

    def step_4_filter_top_quality(
        self,
        evaluations: List[Dict],
        threshold_average: float = 3.0,
        min_correctness: float = 3.0
    ) -> List[Dict]:
        """
        Paso 4: Filtrar solo top-quality responses
        Threshold: promedio ≥3.0 Y correctness ≥3.0 (hard requirement)
        """
        filtered = []

        for eval_item in evaluations:
            scores = eval_item["scores"]

            # Hard requirements
            if (eval_item["average"] >= threshold_average and
                scores.get("correctness", 0) >= min_correctness):
                filtered.append(eval_item)

        return filtered

    def run_full_pipeline(
        self,
        domain: str,
        total_samples: int = 10000
    ) -> List[Dict]:
        """
        Pipeline completo end-to-end
        """
        print(f"Generando {total_samples} synthetic alignment samples...")

        # Paso 1: Preparar prompts
        prompts = self.step_1_prepare_prompts(domain, count=total_samples)

        all_filtered_data = []

        for i, prompt in enumerate(prompts):
            if i % 100 == 0:
                print(f"Procesado {i}/{len(prompts)} prompts...")

            # Paso 2: Generar 10 candidate responses
            candidates = self.step_2_generate_responses(prompt, num_candidates=10)

            # Paso 3: Evaluar con reward model
            evaluations = self.step_3_evaluate_with_reward_model(prompt, candidates)

            # Paso 4: Filtrar top quality (típicamente 1-2 de 10 pasan threshold)
            filtered = self.step_4_filter_top_quality(evaluations)

            all_filtered_data.extend([
                {
                    "prompt": prompt,
                    "response": item["response"],
                    "scores": item["scores"]
                }
                for item in filtered
            ])

        print(f"Total samples generados: {len(prompts)}")
        print(f"Samples que pasaron quality filter: {len(all_filtered_data)}")
        print(f"Pass rate: {len(all_filtered_data) / (len(prompts) * 10) * 100:.1f}%")

        return all_filtered_data

    def _parse_reward_scores(self, raw_output: str) -> Dict[str, float]:
        """
        Parsea output del reward model a dict de scores
        """
        # Implementación simplificada
        # En producción, usa regex robusto o structured output
        scores = {}

        for dim in self.reward_dimensions:
            # Buscar "helpfulness: 4" en el texto
            if f"{dim}:" in raw_output.lower():
                # Extraer número
                score_str = raw_output.lower().split(f"{dim}:")[1].split(",")[0].strip()
                scores[dim] = float(score_str)

        return scores


# Uso
pipeline = NemotronSyntheticPipeline()

synthetic_dataset = pipeline.run_full_pipeline(
    domain="customer_support",
    total_samples=10000
)

# Resultado esperado:
# - 10,000 prompts × 10 candidates = 100,000 responses generadas
# - Pass rate típico: 10-20% = 10,000-20,000 final samples
# - Tiempo: ~2-5 días (limited by API rate limits / GPU compute)
# - Coste: ~500-2000€ (NVIDIA NGC pricing) vs 100,000€+ human annotation