#!/usr/bin/env python3
"""
Script para diversificar instruções em um dataset de treinamento.

Uso:
    python diversify_instructions.py --input dataset.jsonl --output dataset_diversified.jsonl
"""

import json
import argparse
import random
from typing import List, Dict

# Variações de instruções para EDA
INSTRUCTION_VARIATIONS = [
    "Realize uma Análise Exploratória de Dados (EDA) completa seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Analise os dados fornecidos e realize uma EDA completa seguindo os critérios estabelecidos.",
    "Faça uma Análise Exploratória de Dados seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Realize uma análise exploratória completa dos dados seguindo os critérios estabelecidos.",
    "Execute uma EDA detalhada dos dados fornecidos seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Analise exploratoriamente o dataset seguindo os padrões estabelecidos.",
    "Realize uma Análise Exploratória de Dados completa conforme os critérios estabelecidos.",
    "Faça uma análise exploratória detalhada dos dados seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Analise os dados de forma exploratória seguindo os critérios estabelecidos.",
    "Realize uma EDA completa dos dados fornecidos seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Execute uma Análise Exploratória de Dados seguindo os critérios estabelecidos.",
    "Faça uma análise exploratória completa seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Analise exploratoriamente os dados seguindo os padrões estabelecidos.",
    "Realize uma EDA detalhada seguindo OBRIGATORIAMENTE os critérios estabelecidos.",
    "Execute uma análise exploratória completa dos dados seguindo os critérios estabelecidos.",
]


def diversify_instructions(input_file: str, output_file: str, variations: List[str] = None) -> Dict:
    """
    Diversifica as instruções em um dataset.
    
    Args:
        input_file: Caminho do arquivo de entrada
        output_file: Caminho do arquivo de saída
        variations: Lista de variações de instruções (opcional)
    
    Returns:
        Estatísticas da diversificação
    """
    if variations is None:
        variations = INSTRUCTION_VARIATIONS
    
    examples = []
    
    # Ler exemplos
    print(f"📖 Lendo dataset: {input_file}")
    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                examples.append(json.loads(line))
    
    print(f"✅ {len(examples)} exemplos carregados")
    
    # Diversificar instruções
    print(f"🔄 Diversificando instruções com {len(variations)} variações...")
    
    stats = {
        'total': len(examples),
        'original_instructions': set(),
        'new_instructions': set(),
        'distribution': {}
    }
    
    for ex in examples:
        original_instruction = ex.get('instruction', '')
        stats['original_instructions'].add(original_instruction)
        
        # Selecionar variação aleatória
        new_instruction = random.choice(variations)
        ex['instruction'] = new_instruction
        
        stats['new_instructions'].add(new_instruction)
        stats['distribution'][new_instruction] = stats['distribution'].get(new_instruction, 0) + 1
    
    # Salvar dataset diversificado
    print(f"💾 Salvando dataset diversificado: {output_file}")
    with open(output_file, 'w', encoding='utf-8') as f:
        for ex in examples:
            f.write(json.dumps(ex, ensure_ascii=False) + '\n')
    
    print(f"✅ Dataset salvo com sucesso!")
    
    return stats


def print_stats(stats: Dict):
    """Imprime estatísticas da diversificação."""
    print("\n" + "=" * 80)
    print("ESTATÍSTICAS DA DIVERSIFICAÇÃO")
    print("=" * 80)
    print(f"Total de exemplos: {stats['total']}")
    print(f"Instruções originais únicas: {len(stats['original_instructions'])}")
    print(f"Instruções novas únicas: {len(stats['new_instructions'])}")
    print(f"Diversidade: {len(stats['new_instructions'])/stats['total']*100:.1f}%")
    print()
    
    print("Distribuição das instruções:")
    for instruction, count in sorted(stats['distribution'].items(), key=lambda x: x[1], reverse=True):
        percentage = (count / stats['total']) * 100
        print(f"  • {count:4d} ({percentage:5.1f}%): {instruction[:60]}...")
    print()


def main():
    parser = argparse.ArgumentParser(
        description='Diversifica instruções em um dataset de treinamento'
    )
    parser.add_argument(
        '--input',
        type=str,
        required=True,
        help='Arquivo de entrada (JSONL)'
    )
    parser.add_argument(
        '--output',
        type=str,
        required=True,
        help='Arquivo de saída (JSONL)'
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Seed para reprodutibilidade (padrão: 42)'
    )
    
    args = parser.parse_args()
    
    # Definir seed para reprodutibilidade
    random.seed(args.seed)
    
    # Diversificar
    stats = diversify_instructions(args.input, args.output)
    
    # Imprimir estatísticas
    print_stats(stats)
    
    print("=" * 80)
    print("✅ DIVERSIFICAÇÃO CONCLUÍDA!")
    print("=" * 80)
    print(f"\nArquivo gerado: {args.output}")
    print("\n⚠️ IMPORTANTE: Revise alguns exemplos para garantir que a diversificação")
    print("   não afetou a qualidade ou coerência dos dados.")


if __name__ == "__main__":
    main()