#!/usr/bin/env python3
"""train.py — Treino de bug-detection para BiGRU_T_version (2 épocas).

Especificações do usuário atendidas:
- 2 épocas fixas (não mudar sem permissão)
- Sem dados inventados (apenas datasets do repo + HF públicos)
- Monitorar RAM, disco, loss, perplexidade
- Kill-switch: matar se loss não diminuir ou RAM explodir
- Estados temporários salvos durante treino e APAGADOS ao final
- Parâmetros do modelo aprendidos automaticamente (sem inserção manual)
- Ambiente otimizado para Xeon (AVX512, OpenMP, MKL)
- HF_TOKEN apagado ao final (se usado para datasets privados)

Uso:
    export HF_TOKEN="hf_xxx"  # opcional para datasets públicos
    python scripts/train.py \\
        --datasets CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1,Madras1/corpus-ptbr-v2 \\
        --max-samples 30 \\
        --epochs 2
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from pathlib import Path

# Adiciona src/ ao path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

# Otimização Xeon (deve vir ANTES de importar torch)
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("TORCH_NUM_THREADS", "2")

import torch
torch.set_num_threads(2)

# Tenta otimizar ambiente Xeon (do módulo reaproveitado)
try:
    from bigru_t.utils.xeon_runtime import optimize_xeon_environment
    optimize_xeon_environment()
except Exception as e:
    logging.warning(f"Could not apply Xeon optimization: {e}")

from bigru_t import (
    UnifiedModel, UnifiedModelConfig, create_unified_model,
    BiGRU_T_Trainer, TrainerConfig,
)
from bigru_t.data.streaming_datasets import stream_dataset

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)


def load_tokenizer(path: str):
    """Carrega tokenizer BBPE 16K."""
    from tokenizers import Tokenizer
    return Tokenizer.from_file(path)


def load_samples(
    datasets: str,
    max_samples_per_dataset: int,
    hf_token: str | None,
) -> list:
    """Carrega amostras via streaming (reaproveita streaming_datasets_v13_9)."""
    samples = []
    datasets_list = [d.strip() for d in datasets.split(",") if d.strip()]
    logger.info(f"Loading samples from {len(datasets_list)} datasets...")
    for ds_name in datasets_list:
        ds_start = time.time()
        ds_count = 0
        try:
            for sample in stream_dataset(
                ds_name,
                max_samples=max_samples_per_dataset,
                hf_token=hf_token,
            ):
                samples.append(sample)
                ds_count += 1
                if ds_count >= max_samples_per_dataset:
                    break
                if time.time() - ds_start > 300:
                    logger.warning(f"  timeout em {ds_name} após {ds_count} amostras")
                    break
        except Exception as e:
            logger.error(f"Erro carregando {ds_name}: {e}")
            continue
        logger.info(f"  {ds_name}: {ds_count} amostras em {time.time()-ds_start:.1f}s")
    return samples


def main():
    parser = argparse.ArgumentParser(description="BiGRU_T_version — Bug-detection training (2 epochs)")
    parser.add_argument(
        "--datasets",
        type=str,
        default="CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1,Madras1/corpus-ptbr-v2,dominguesm/restore-punctuation-ptbr-dataset",
        help="Lista separada por vírgulas de datasets HF",
    )
    parser.add_argument("--max-samples", type=int, default=30, help="Amostras POR dataset (default 30)")
    parser.add_argument("--epochs", type=int, default=2, help="Nº de épocas (default 2)")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--max-seq-len", type=int, default=64)
    parser.add_argument("--output-dir", type=str, default="/home/z/my-project/BiGRU_T_version/model_final")
    parser.add_argument("--temp-dir", type=str, default="/home/z/my-project/BiGRU_T_version/_temp_checkpoints")
    parser.add_argument(
        "--tokenizer-path",
        type=str,
        default="/home/z/my-project/source/model_final/tokenizer/tokenizer.json",
    )
    parser.add_argument("--no-hypothesis", action="store_true", help="Desativar hipótese (Lema 3)")
    parser.add_argument("--keep-temp", action="store_true", help="Manter checkpoints temporários")
    parser.add_argument("--smoke-test", action="store_true", help="Smoke test: 5 samples, 1 epoch")
    args = parser.parse_args()

    if args.smoke_test:
        logger.info("SMOKE TEST MODE — 5 samples/dataset, 1 epoch")
        args.max_samples = 5
        args.epochs = 1
        args.grad_accum = 2
        args.datasets = "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1"

    # HF token (opcional para datasets públicos)
    hf_token = os.environ.get("HF_TOKEN") or None

    # Configuração do modelo
    model_config = UnifiedModelConfig(
        vocab_size=16384,
        d_model=128,
        max_seq_len=args.max_seq_len,
        # input_dim é automaticamente = d_model via __post_init__
        # Para bug-detection, reduzimos max_modules (64 é caro)
        max_modules=8,  # 8 módulos em vez de 64
        bigru_hidden=32,
        d_transformer=64,
        nhead_tu=4,
        d_ff_tu=128,
        output_dim_u8cell=64,
        cache_len=16,
        d_cache=128,
        nhead_orq=8,
        d_ff_orq=256,
        trainT_dim=128,
        nhead_train=4,
        d_ff_train=256,
        num_layers_train=2,
        hypT_dim=128,
        nhead_hyp=4,
        d_ff_hyp=256,
        num_layers_hyp=2,
        num_bits=8,
        dropout=0.1,
    )

    # Cria modelo
    logger.info("Criando UnifiedModel...")
    model, _ = create_unified_model(model_config)
    params = model.count_parameters()
    logger.info(f"  params: {params['total']:,} ({params['total_M']:.2f}M)")

    # Carrega tokenizer
    tokenizer = load_tokenizer(args.tokenizer_path)

    # Carrega datasets
    samples = load_samples(args.datasets, args.max_samples, hf_token)
    n = len(samples)
    if n == 0:
        logger.error("Nenhuma amostra carregada — abortando")
        sys.exit(1)
    logger.info(f"Total: {n} amostras carregadas")

    # Split 90/10 treino/val
    import random
    random.seed(42)
    random.shuffle(samples)
    split = max(1, int(0.9 * n))
    train_samples = samples[:split]
    val_samples = samples[split:] or samples[:1]  # fallback: 1 amostra
    logger.info(f"  train: {len(train_samples)} | val: {len(val_samples)}")

    # Configuração do trainer
    trainer_config = TrainerConfig(
        epochs=args.epochs,
        datasets=args.datasets,
        max_samples_per_dataset=args.max_samples,
        max_seq_len=args.max_seq_len,
        per_device_batch_size=args.batch_size,
        grad_accum=args.grad_accum,
        lr=args.lr,
        use_hypothesis=not args.no_hypothesis,
        output_dir=args.output_dir,
        temp_dir=args.temp_dir,
        keep_temp=args.keep_temp,
    )

    # Trainer
    trainer = BiGRU_T_Trainer(
        model=model,
        tokenizer=tokenizer,
        config=trainer_config,
        train_samples=train_samples,
        val_samples=val_samples,
    )

    # Treina
    result = trainer.train()

    # Print resultado
    import json
    print("\n=== Resultado Final ===")
    print(json.dumps(result, indent=2, default=str))

    # Apaga HF token do ambiente (especificação do usuário)
    if "HF_TOKEN" in os.environ:
        del os.environ["HF_TOKEN"]
        logger.info("HF_TOKEN apagado do ambiente")

    # Exit code: 0 se OK, 1 se morto pelo kill-switch
    sys.exit(0 if not result.get("killed") else 1)


if __name__ == "__main__":
    main()