#!/usr/bin/env python3
"""train_fast.py — Treino rápido de bug-detection (2 épocas, modelo pequeno).

Versão otimizada para rodar em < 5 minutos no ambiente Xeon 2-core:
- max_modules=4 (em vez de 8)
- bigru_hidden=16 (em vez de 32)
- d_transformer=32 (em vez de 64)
- num_layers_train=1, num_layers_hyp=1
- max_samples=8 por dataset
- 1 dataset (CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1)

Valida todos os 4 lemas sem pretender produzir modelo SOTA.
"""
from __future__ import annotations

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

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

# Otimização Xeon
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import torch
torch.set_num_threads(2)

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 main():
    # HF token
    hf_token = os.environ.get("HF_TOKEN") or None

    # Modelo pequeno para bug-detection rápido
    model_config = UnifiedModelConfig(
        vocab_size=16384,
        d_model=64,
        max_seq_len=32,
        pad_token_id=1,
        max_modules=4,           # 4 módulos em vez de 8
        bigru_hidden=16,         # 16 em vez de 32
        d_transformer=32,        # 32 em vez de 64
        nhead_tu=2,              # 2 em vez de 4 (32/2=16 ok)
        d_ff_tu=64,              # 64 em vez de 128
        output_dim_u8cell=32,    # 32 em vez de 64
        cache_len=8,             # 8 em vez de 16
        d_cache=64,              # 64 em vez de 128
        nhead_orq=2,             # 2 em vez de 8 (64/2=32 ok)
        d_ff_orq=128,            # 128 em vez de 256
        trainT_dim=64,
        nhead_train=2,
        d_ff_train=128,
        num_layers_train=1,      # 1 em vez de 2
        hypT_dim=64,
        nhead_hyp=2,
        d_ff_hyp=128,
        num_layers_hyp=1,        # 1 em vez de 2
        num_bits=8,
        dropout=0.1,
    )

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

    # Tokenizer
    from tokenizers import Tokenizer
    tok_path = "/home/z/my-project/source/model_final/tokenizer/tokenizer.json"
    tokenizer = Tokenizer.from_file(tok_path)

    # Carrega 1 dataset
    ds_name = "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1"
    max_samples = 8
    logger.info(f"Loading {max_samples} samples from {ds_name}...")
    samples = []
    t0 = time.time()
    for sample in stream_dataset(ds_name, max_samples=max_samples, hf_token=hf_token):
        samples.append(sample)
        if len(samples) >= max_samples:
            break
        if time.time() - t0 > 120:
            logger.warning(f"Timeout após {len(samples)} amostras")
            break
    logger.info(f"Loaded {len(samples)} samples in {time.time()-t0:.1f}s")

    if not samples:
        logger.error("Nenhuma amostra carregada")
        sys.exit(1)

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

    # Trainer config
    trainer_config = TrainerConfig(
        epochs=2,
        datasets=ds_name,
        max_samples_per_dataset=max_samples,
        max_seq_len=32,
        per_device_batch_size=1,
        grad_accum=2,
        lr=1e-3,
        use_hypothesis=True,
        meta_interval=5,
        log_every=1,
        save_temp_every=100,  # não salvar temp (modelo pequeno)
        output_dir="/home/z/my-project/BiGRU_T_version/model_final",
        temp_dir="/home/z/my-project/BiGRU_T_version/_temp_checkpoints",
        keep_temp=False,
    )

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

    result = trainer.train()

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

    # Apaga HF token
    if "HF_TOKEN" in os.environ:
        del os.environ["HF_TOKEN"]
        logger.info("HF_TOKEN apagado do ambiente")

    sys.exit(0 if not result.get("killed") else 1)


if __name__ == "__main__":
    main()