#!/usr/bin/env python3
"""train_v2.py — Treino BiGRU_T_version V2 (15M params, HW optimizer, multimodal bug-hunt).

Melhorias vs train_fast.py:
  - max_modules=10 (aumento de 4→10, memória permite), d_model=128
  - ~14.45M params (objetivo 15M)
  - HamiltonianWasserstein optimizer ATIVADO (AdamW + W₂ + repulsão + LR cíclico)
  - TimeBudget + aggressive_cleanup (estilo Xavante)
  - Teste de multimodalidade (text/image/audio/video encoders) com poucos samples
  - Timeout aumentado (1800s total, 900s/epoch)
  - Verificação de OOM-Killer antes do treino
  - Smoke checks antes do treino principal

Uso:
    export HF_TOKEN=hf_xxx  # opcional, para datasets privados
    python3 scripts/train_v2.py
"""
from __future__ import annotations

import os
import sys
import time
import json
import logging
import gc
import shutil
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,
    HamiltonianWassersteinOptimizer,
    CircularReasoningWasserstein,
    BiGRUTGenerator,
    aggressive_cleanup, get_rss_mb, TimeBudget,
)
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 check_oom_killer_risk() -> bool:
    """Verifica se há risco de OOM-Killer antes de iniciar treino.

    Checa:
      1. dmesg por ocorrências recentes de OOM-Killer
      2. Memória livre do sistema
      3. Swap disponível

    Returns:
        True se seguro para prosseguir, False se risco alto.
    """
    import subprocess
    # 1. dmesg OOM check
    try:
        result = subprocess.run(
            ["dmesg"], capture_output=True, text=True, timeout=5
        )
        if "Out of memory" in result.stdout or "Killed process" in result.stdout:
            logger.warning("⚠️  dmesg mostra OOM-Killer activity recente — risco alto")
            # Não aborta, apenas avisa (pode ser histórico antigo)
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass  # dmesg pode não estar acessível

    # 2. Memória livre do sistema
    try:
        with open("/proc/meminfo", "r") as f:
            meminfo = dict(line.split(":")[1].strip().split()[0:2] for line in f if ":" in line)
        free_mb = int(meminfo.get("MemAvailable", "0")) / 1024
        total_mb = int(meminfo.get("MemTotal", "0")) / 1024
        logger.info(f"  Sistema: {total_mb:.0f}MB total, {free_mb:.0f}MB disponível")
        if free_mb < 512:
            logger.error(f"❌ Memória disponível muito baixa ({free_mb:.0f}MB < 512MB)")
            return False
    except Exception:
        pass

    # 3. Disk free
    try:
        usage = shutil.disk_usage("/home/z/my-project")
        free_gb = usage.free / (1024**3)
        logger.info(f"  Disk free: {free_gb:.2f}GB")
        if free_gb < 1.0:
            logger.error(f"❌ Disk livre muito baixo ({free_gb:.2f}GB < 1GB)")
            return False
    except Exception:
        pass

    return True


def smoke_test_model(model, tokenizer) -> bool:
    """Smoke test rápido: forward + backward + inference."""
    logger.info("Running smoke test (forward + backward + generate)...")
    try:
        import torch
        # Forward — COM grad (para permitir backward)
        dummy_ids = torch.randint(0, 100, (1, 16), dtype=torch.long)
        y_hat, delta = model(dummy_ids, temperature=1.0, use_hypothesis=False)
        assert y_hat.shape[0] == 1, f"batch mismatch: {y_hat.shape}"
        assert not torch.isnan(y_hat).any(), "NaN in y_hat"
        logger.info(f"  Forward OK: y_hat {tuple(y_hat.shape)}, delta {tuple(delta.shape)}")

        # Backward
        target = torch.tensor([5], dtype=torch.long)
        loss = torch.nn.functional.cross_entropy(y_hat, target)
        loss.backward()
        grad_norm = sum(p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5
        logger.info(f"  Backward OK: loss={loss.item():.4f}, grad_norm={grad_norm:.4f}")
        model.zero_grad()

        # Generate (no grad)
        gen = BiGRUTGenerator(model, max_seq_len=64)
        with torch.no_grad():
            out = gen.generate(dummy_ids, max_new_tokens=5, do_sample=False)
        logger.info(f"  Generate OK: {dummy_ids.shape} -> {out.shape}")
        return True
    except Exception as e:
        import traceback
        logger.error(f"  Smoke test FAILED: {e}")
        traceback.print_exc()
        return False


def test_multimodal_encoders() -> dict:
    """Testa os encoders multimodais com dados sintéticos (bug detection).

    Não usa datasets reais (apenas poucos samples sintéticos para validar
    que os encoders não quebram). Retorna dict com status de cada encoder.

    BUGS ENCONTRADOS E CORRIGIDOS:
      1. TextEncoder dependia de 3 módulos ausentes (attention_multimodal,
         embedding_reconfig, gru_hierarchy) — COPIADOS da fonte HF.
      2. AudioEncoder: parâmetro `in_channels` é ignorado pelo Conv1d interno
         (sempre usa n_mels=80 como canal) — WORKAROUND: passar [B, 80, T].
      3. VideoEncoder.forward(frames, audio) espera frames=[B,T,C,H,W] (5D),
         NÃO [B,C,T,H,W] — corrigido no teste.
      4. ModalRouter.forward(inputs: Dict[str, Tensor]) — espera dict nomeado.
      5. FusionLayer.forward(modality_embs: list) — espera lista de embeddings.
    """
    logger.info("Testing multimodal encoders (synthetic data, bug detection)...")
    results = {}

    try:
        from bigru_t.multimodal.text_encoder import TextEncoder
        enc = TextEncoder(vocab_size=16384, d_model=128, n_heads=4, n_gru_levels=2, max_seq_len=64)
        import torch
        x = torch.randint(0, 100, (2, 16), dtype=torch.long)
        out = enc(x)
        results["text_encoder"] = {"ok": True, "shape": str(tuple(out.shape))}
        logger.info(f"  TextEncoder OK: {tuple(out.shape)}")
    except Exception as e:
        results["text_encoder"] = {"ok": False, "error": str(e)}
        logger.warning(f"  TextEncoder FAIL: {e}")

    try:
        from bigru_t.multimodal.image_encoder import ImageEncoder
        enc = ImageEncoder(d_model=128, in_channels=3)
        import torch
        x = torch.randn(2, 3, 16, 16)
        out = enc(x)
        results["image_encoder"] = {"ok": True, "shape": str(tuple(out.shape))}
        logger.info(f"  ImageEncoder OK: {tuple(out.shape)}")
    except Exception as e:
        results["image_encoder"] = {"ok": False, "error": str(e)}
        logger.warning(f"  ImageEncoder FAIL: {e}")

    try:
        from bigru_t.multimodal.audio_encoder import AudioEncoder
        # BUG: in_channels é ignorado; Conv1d usa n_mels como canal
        enc = AudioEncoder(d_model=128, in_channels=80, n_mels=80)
        import torch
        # Shape correta: [B, n_mels=80, T_audio]
        x = torch.randn(2, 80, 100)
        out = enc(x)
        results["audio_encoder"] = {"ok": True, "shape": str(tuple(out.shape)),
                                     "note": "in_channels param unused; uses n_mels=80 as channel"}
        logger.info(f"  AudioEncoder OK: {tuple(out.shape)}")
    except Exception as e:
        results["audio_encoder"] = {"ok": False, "error": str(e)}
        logger.warning(f"  AudioEncoder FAIL: {e}")

    try:
        from bigru_t.multimodal.video_encoder import VideoEncoder
        enc = VideoEncoder(d_model=128, n_frames=4)
        import torch
        # BUG: frames deve ser [B, T, C, H, W] (5D, T antes de C)
        frames = torch.randn(2, 4, 3, 16, 16)  # (B=2, T=4, C=3, H=16, W=16)
        audio = torch.randn(2, 80, 100)  # (B, n_mels, T_audio)
        out = enc(frames, audio)
        results["video_encoder"] = {"ok": True, "shape": str(tuple(out.shape)),
                                     "note": "frames=[B,T,C,H,W], audio=[B,80,T]"}
        logger.info(f"  VideoEncoder OK: {tuple(out.shape)}")
    except Exception as e:
        results["video_encoder"] = {"ok": False, "error": str(e)}
        logger.warning(f"  VideoEncoder FAIL: {e}")

    try:
        from bigru_t.multimodal.modal_router import ModalRouter
        from bigru_t.multimodal.fusion_layer import FusionLayer
        router = ModalRouter(d_model=128)
        fusion = FusionLayer(d_model=128, n_modalities=4)
        import torch
        # ModalRouter.forward(inputs: Dict[str, Tensor]) — espera dict nomeado
        inputs = {
            "text": torch.randn(2, 128),
            "image": torch.randn(2, 128),
            "audio": torch.randn(2, 128),
            "video": torch.randn(2, 128),
        }
        routed = router(inputs)
        # FusionLayer.forward(modality_embs: list) — espera lista
        embs = [torch.randn(2, 128) for _ in range(4)]
        fused = fusion(embs)
        results["modal_router_fusion"] = {"ok": True,
                                           "router_shape": str(tuple(routed.shape)),
                                           "fusion_shape": str(tuple(fused.shape))}
        logger.info(f"  ModalRouter OK: {tuple(routed.shape)} | Fusion OK: {tuple(fused.shape)}")
    except Exception as e:
        results["modal_router_fusion"] = {"ok": False, "error": str(e)}
        logger.warning(f"  ModalRouter/Fusion FAIL: {e}")

    return results


def main():
    logger.info("=" * 70)
    logger.info("BiGRU_T_version V2 — 15M params + HW optimizer + multimodal bug-hunt")
    logger.info("=" * 70)

    # HF token (do ambiente; será apagado ao final)
    hf_token = os.environ.get("HF_TOKEN") or None
    if hf_token:
        logger.info("HF_TOKEN encontrado no ambiente")
    else:
        logger.info("HF_TOKEN não fornecido (datasets públicos apenas)")

    # 1. Verificação de OOM-Killer e recursos
    logger.info("\n[1/6] Verificando recursos do sistema...")
    if not check_oom_killer_risk():
        logger.error("Recursos insuficientes. Abortando.")
        sys.exit(1)
    logger.info(f"  RSS atual: {get_rss_mb():.1f}MB")

    # 2. Teste multimodal (bug detection)
    logger.info("\n[2/6] Testando multimodalidade (bug hunting)...")
    mm_results = test_multimodal_encoders()
    mm_ok = sum(1 for v in mm_results.values() if v.get("ok"))
    mm_total = len(mm_results)
    logger.info(f"  Multimodal: {mm_ok}/{mm_total} encoders OK")

    # 3. Config do modelo (12.97M params — max_modules=8, d_model=128)
    logger.info("\n[3/6] Criando UnifiedModel (max_modules=8, d_model=128)...")
    model_config = UnifiedModelConfig(
        vocab_size=16384,
        d_model=128,
        max_seq_len=32,
        pad_token_id=1,
        max_modules=8,            # base spec do usuário (8)
        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,
    )
    model, _ = create_unified_model(model_config)
    params = model.count_parameters()
    logger.info(f"  params: {params['total']:,} ({params['total_M']:.2f}M)")
    logger.info(f"  trainable: {params['trainable']:,} ({params['trainable_M']:.2f}M)")
    aggressive_cleanup(verbose=True)

    # 4. Tokenizer (usa o pré-treinado do source)
    logger.info("\n[4/6] Carregando tokenizer...")
    from tokenizers import Tokenizer
    tok_path = "/home/z/my-project/source/model_final/tokenizer/tokenizer.json"
    if not Path(tok_path).exists():
        logger.error(f"Tokenizer não encontrado: {tok_path}")
        sys.exit(1)
    tokenizer = Tokenizer.from_file(tok_path)
    logger.info(f"  Tokenizer loaded: vocab={tokenizer.get_vocab_size()}")

    # 5. Smoke test do modelo
    logger.info("\n[5/6] Smoke test do modelo...")
    if not smoke_test_model(model, tokenizer):
        logger.error("Smoke test falhou. Abortando.")
        sys.exit(1)
    aggressive_cleanup(verbose=True)

    # 6. Carrega datasets (streaming, poucos samples para bug detection)
    logger.info("\n[6/6] Carregando datasets (streaming, poucos samples)...")
    datasets = [
        "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1",
    ]
    max_samples = 8  # poucos samples (bug detection, testes)
    all_samples = []
    t0 = time.time()
    for ds_name in datasets:
        if time.time() - t0 > 180:
            logger.warning(f"Timeout carregando datasets após {len(all_samples)} amostras")
            break
        try:
            count_before = len(all_samples)
            for sample in stream_dataset(ds_name, max_samples=max_samples, hf_token=hf_token):
                all_samples.append(sample)
                if len(all_samples) >= max_samples * len(datasets):
                    break
                if time.time() - t0 > 180:
                    break
            logger.info(f"  {ds_name}: +{len(all_samples)-count_before} samples")
        except Exception as e:
            logger.warning(f"  {ds_name} falhou: {e}")

    if not all_samples:
        logger.error("Nenhuma amostra carregada de nenhum dataset")
        sys.exit(1)

    logger.info(f"Total: {len(all_samples)} samples em {time.time()-t0:.1f}s")

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

    aggressive_cleanup(verbose=True)

    # Trainer config — HW optimizer + time budget + cleanup agressivo
    trainer_config = TrainerConfig(
        epochs=2,
        datasets=",".join(datasets),
        max_samples_per_dataset=max_samples,
        max_seq_len=32,
        per_device_batch_size=1,
        grad_accum=2,
        lr=1e-3,
        optimizer_type="hamiltonian_wasserstein",  # ATIVADO
        hw_lr_amp=0.3,
        hw_lr_freq=0.01,
        hw_sigma_w=1.0,
        hw_sigma_rep=0.1,
        hw_prune_every_n=0,  # pruning off para bug detection
        use_hypothesis=True,
        stop_grad_hyp=True,
        meta_interval=5,
        max_total_time_s=1800.0,   # 30 min total
        max_per_epoch_s=900.0,     # 15 min/epoch
        cleanup_every_n_steps=5,
        log_every=1,
        save_temp_every=3,  # salvar a cada 3 steps (garante checkpoint mesmo se OOM no fim)
        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,
        ram_threshold_pct=90.0,
        disk_min_free_gb=1.0,
        loss_patience=30,
    )

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

    # Treino
    result = trainer.train()

    # Salva resultados extras (multimodal + smoke)
    extra_report = {
        "multimodal_test": mm_results,
        "smoke_test": {"passed": True},
        "config": {
            "max_modules": model_config.max_modules,
            "d_model": model_config.d_model,
            "bigru_hidden": model_config.bigru_hidden,
            "d_transformer": model_config.d_transformer,
            "params_total": params["total"],
            "params_M": params["total_M"],
            "optimizer": "hamiltonian_wasserstein",
        },
    }
    extra_path = Path("/home/z/my-project/BiGRU_T_version/multimodal_report.json")
    with open(extra_path, "w") as f:
        json.dump(extra_report, f, indent=2, default=str)
    logger.info(f"Relatório multimodal salvo: {extra_path}")

    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()