#!/usr/bin/env python3
"""train_from_v13_9_2.py — Continua treinamento do GRU-RING v13.9.2
em um ou mais datasets HF (default: Carolina), partindo dos pesos v13.9.2
já treinados.

V13.9.2-finetune-v2: generalizado para aceitar --datasets (lista separada
por vírgula). Default mantém "carolina-c4ai/corpus-carolina" para
retrocompatibilidade. Nenhuma mudança na arquitetura do modelo — apenas
na escolha do dataset de fine-tuning.

ESTRATÉGIAS DE BAIXA MEMÓRIA (3.5GB RAM disponível):
  - streaming=True: dataset lido amostra por amostra (via lxml.iterparse direto,
    bypass do script Python carolina.py que não é suportado em datasets 5.0+)
  - per_device_train_batch_size=1: 1 amostra por vez
  - gradient_accumulation_steps=16: acumula 16 gradientes antes de atualizar
    pesos (simula batch_size=16 sem explodir RAM)
  - max_length=64: o max_seq_len nativo do modelo v13.9.2 (não 128 — preserva
    a arquitetura e economiza 2× memória de atenção)
  - gradient_checkpointing=True: já habilitado no modelo (config.gradient_checkpointing)

BUG-HUNT (verificado neste script):
  1. NoneType backward: backward() só é chamado quando outputs contém "loss"
     (labels passadas). Sempre passamos labels.
  2. state_dict mismatch: usamos create_gru_ring_v139() com config idêntica à
     do treino original (d_model=128, n_layers=4, etc) — load_state_dict é safe.
  3. Tokenizer encoding: collate_samples_to_tensors já trata .ids / .input_ids /
     list. Reaproveitado do módulo streaming_datasets_v13_9.
  4. Buffer trust_state persistente: preservado no state_dict (register_buffer).
  5. padding_idx=1: pad_token_id=1 no config. Garantido por _normalize_sample
     que ignora textos com <10 chars (sem padding-only batches).

Uso (default — Carolina):
  export HF_TOKEN="hf_xxx"
  python3 train_from_v13_9_2.py \
      --max-samples 200 \
      --max-steps 200 \
      --grad-accum 16 \
      --lr 2e-5 \
      --output-dir /home/z/my-project/download/gru_ring_v13_9/carolina_finetune

Uso (fine-tune v2 — 3 novos datasets PT-BR):
  export HF_TOKEN="hf_xxx"
  python3 train_from_v13_9_2.py \
      --datasets orion-research/translations-en_US-pt_BR,cnmoro/Instruct-PTBR-10M,strak2005/corpus-ptbr-v1 \
      --max-samples 200 \
      --max-steps 200 \
      --grad-accum 16 \
      --lr 2e-5 \
      --output-dir /home/z/my-project/download/gru_ring_v13_9/finetune_v2
"""
from __future__ import annotations

import argparse
import gc
import json
import logging
import math
import os
import sys
import time
import traceback
from pathlib import Path
from typing import Dict, List, Optional

# Limitar threads para evitar OOM em ambiente com 3.5GB RAM
import torch
torch.set_num_threads(2)
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

from torch.optim import AdamW

sys.path.insert(0, "/home/z/my-project/xavante_work")

from flexnet.gru_ring_v13_9 import create_gru_ring_v139, GRURingV139, GRURingV139Config
from flexnet.streaming_datasets_v13_9 import (
    stream_dataset,
    collate_samples_to_tensors,
    ProcessedSample,
)
from flexnet.w8a8_qoperator import W8A8QOperatorQuantizer

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


# Caminhos padrão (v13.9.2 treinado)
# BUG FIX V13.9.2-finetune-v2: paths apontavam para /tmp/my-project/ (legado).
# Corrigido para /home/z/my-project/ que é onde os artefatos realmente estão.
DEFAULT_MODEL_PATH = "/home/z/my-project/download/gru_ring_v13_9/rename_uploads/model_final/pytorch_model.bin"
DEFAULT_TOKENIZER_PATH = "/home/z/my-project/download/gru_ring_v13_9/rename_uploads/model_final/tokenizer/tokenizer.json"
DEFAULT_OUTPUT_DIR = "/home/z/my-project/download/gru_ring_v13_9/finetune_v2"
# Default dataset (retrocompatibilidade). Para mudar, use --datasets.
DEFAULT_DATASETS = "carolina-c4ai/corpus-carolina"


def get_warmup_cosine_scheduler(optimizer, num_warmup, num_total, min_ratio=0.1):
    """Warmup + cosine schedule."""
    def lr_lambda(step):
        if step < num_warmup:
            return float(step) / max(1, num_warmup)
        progress = float(step - num_warmup) / max(1, num_total - num_warmup)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return max(min_ratio, cosine)
    from torch.optim.lr_scheduler import LambdaLR
    return LambdaLR(optimizer, lr_lambda)


def load_tokenizer(path: str):
    """Carrega tokenizer BBPE 16K."""
    try:
        from tokenizers import Tokenizer
        return Tokenizer.from_file(path)
    except Exception as e:
        logger.error(f"Failed to load tokenizer from {path}: {e}")
        return None


def load_pretrained_model(model_path: str) -> GRURingV139:
    """Carrega modelo v13.9.2 a partir dos pesos treinados."""
    model, config = create_gru_ring_v139()
    counts = model.count_parameters()
    logger.info(f"Model created: {counts['total']:,} params ({counts['total_M']:.2f}M)")

    if not Path(model_path).exists():
        raise FileNotFoundError(f"Model not found: {model_path}")

    logger.info(f"Loading pretrained weights: {model_path}")
    state = torch.load(model_path, map_location="cpu", weights_only=False)
    if isinstance(state, dict) and "model_state_dict" in state:
        state = state["model_state_dict"]
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        logger.warning(f"Missing keys ({len(missing)}): {missing[:5]}...")
    if unexpected:
        logger.warning(f"Unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
    logger.info("Weights loaded successfully")
    return model, config


def train_carolina(
    model_path: str = DEFAULT_MODEL_PATH,
    tokenizer_path: str = DEFAULT_TOKENIZER_PATH,
    output_dir: str = DEFAULT_OUTPUT_DIR,
    max_samples: int = 200,
    max_steps: int = 200,
    grad_accum: int = 16,
    lr: float = 2e-5,
    weight_decay: float = 0.01,
    max_seq_len: int = 64,
    log_every: int = 10,
    save_every: int = 100,
    seed: int = 42,
    datasets: Optional[str] = None,
) -> Dict:
    """Treina o modelo em um ou mais datasets HF (default: Carolina).

    V13.9.2-finetune-v2: aceita `datasets` como string separada por vírgulas.
    Se None ou vazio, usa DEFAULT_DATASETS (Carolina) para retrocompatibilidade.
    Quando múltiplos datasets são passados, cada um contribui com até
    `max_samples` amostras (total = max_samples × nº de datasets).

    Args:
        model_path: caminho para pytorch_model.bin v13.9.2
        tokenizer_path: caminho para tokenizer.json
        output_dir: diretório de saída para checkpoints
        max_samples: nº máx de amostras POR dataset (buffer)
        max_steps: nº máx de optimizer steps (cada step = grad_accum micro-batches)
        grad_accum: gradient_accumulation_steps (batch_size efetivo = 1 × grad_accum)
        lr: learning rate (baixo: 2e-5 — não destruir o aprendizado do .bin)
        weight_decay: weight decay (AdamW)
        max_seq_len: tamanho da sequência (64 = nativo do v13.9.2)
        log_every: log a cada N micro-batches
        save_every: salvar checkpoint a cada N optimizer steps
        seed: random seed
        datasets: string separada por vírgulas com nomes de datasets HF
                  (default: "carolina-c4ai/corpus-carolina")

    Returns:
        dict com métricas finais
    """
    torch.manual_seed(seed)
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # ── Resolve lista de datasets ───────────────────────────────────────
    if not datasets:
        datasets = DEFAULT_DATASETS
    datasets_list = [d.strip() for d in datasets.split(",") if d.strip()]
    if not datasets_list:
        datasets_list = [DEFAULT_DATASETS]
    logger.info(f"Datasets para fine-tune ({len(datasets_list)}):")
    for i, d in enumerate(datasets_list, 1):
        logger.info(f"  [{i}] {d}")
    logger.info(f"max_samples POR dataset = {max_samples} "
                f"(total esperado = {max_samples * len(datasets_list)})")

    # ── Log do config ────────────────────────────────────────────────────
    logger.info("=" * 60)
    logger.info("GRU-RING v13.9.2 — Fine-tuning")
    logger.info("=" * 60)
    logger.info(f"  model_path:      {model_path}")
    logger.info(f"  tokenizer_path:  {tokenizer_path}")
    logger.info(f"  output_dir:      {output_dir}")
    logger.info(f"  max_samples:     {max_samples} (por dataset)")
    logger.info(f"  max_steps:       {max_steps}")
    logger.info(f"  grad_accum:      {grad_accum} (efetivo batch = {1 * grad_accum})")
    logger.info(f"  lr:              {lr}")
    logger.info(f"  max_seq_len:     {max_seq_len}")
    logger.info(f"  gradient_ckpt:   True (já no config do modelo)")
    logger.info("=" * 60)

    # ── HF token ─────────────────────────────────────────────────────────
    # BUG FIX V13.9.2-finetune-v4: `os.environ.get("HF_TOKEN")` retorna ""
    # (string vazia) quando a variável está setada como vazia no shell.
    # Essa string vazia é passada para huggingface_hub que constrói o header
    # HTTP `Authorization: Bearer ` → "Illegal header value". Fix: normalizar
    # para None (ausência de token). Datasets públicos não precisam de token.
    hf_token = os.environ.get("HF_TOKEN") or None
    if not hf_token:
        logger.warning("HF_TOKEN não setado — datasets públicos, sem auth")

    # ── Modelo + tokenizer ───────────────────────────────────────────────
    model, config = load_pretrained_model(model_path)
    device = torch.device("cpu")
    model = model.to(device)
    model.train()

    tokenizer = load_tokenizer(tokenizer_path)
    if tokenizer is None:
        return {"error": "no tokenizer"}

    # ── Optimizer ────────────────────────────────────────────────────────
    # LR baixo (2e-5) para preservar conhecimento dos 8 datasets anteriores.
    optimizer = AdamW(
        model.parameters(),
        lr=lr,
        betas=(0.9, 0.95),
        weight_decay=weight_decay,
        eps=1e-8,
    )

    # Scheduler: warmup (10% dos steps) + cosine
    warmup_steps = max(1, max_steps // 10)
    scheduler = get_warmup_cosine_scheduler(
        optimizer, num_warmup=warmup_steps, num_total=max_steps
    )
    logger.info(f"Scheduler: warmup={warmup_steps}, total={max_steps}")

    # ── Carrega datasets em buffer (streaming) ───────────────────────────
    # V13.9.2-finetune-v2: itera sobre a lista de datasets. Cada um contribui
    # com até max_samples amostras. Timeout por dataset = 180s (Carolina é
    # mais rápido pois usa lxml.iterparse direto; outros podem demorar para
    # materializar a primeira amostra de parquet grande).
    samples_buffer: List[ProcessedSample] = []
    per_dataset_counts: Dict[str, int] = {}
    stream_start_total = time.time()
    for ds_name in datasets_list:
        logger.info(f"Streaming até {max_samples} amostras de {ds_name}...")
        ds_start = time.time()
        ds_count = 0
        try:
            for sample in stream_dataset(
                ds_name,
                max_samples=max_samples,
                hf_token=hf_token,
            ):
                samples_buffer.append(sample)
                ds_count += 1
                if ds_count >= max_samples:
                    break
                if time.time() - ds_start > 600:
                    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}")
            traceback.print_exc()
            # Não aborta — continua com próximos datasets
            continue
        per_dataset_counts[ds_name] = ds_count
        logger.info(f"  {ds_name}: {ds_count} amostras em {time.time()-ds_start:.1f}s")

    n_loaded = len(samples_buffer)
    logger.info(f"Total carregado: {n_loaded} amostras de {len(datasets_list)} "
                f"datasets em {time.time()-stream_start_total:.1f}s")

    if n_loaded == 0:
        return {"error": "no samples loaded"}

    # ── Loop de treino ───────────────────────────────────────────────────
    # Cada micro-batch = 1 amostra (per_device_train_batch_size=1).
    # Acumulamos grad_accum micro-batches antes de optimizer.step().
    global_step = 0
    micro_batch_idx = 0
    sample_idx = 0  # índice circular no buffer
    best_loss = float("inf")
    last_loss = float("nan")
    losses_history: List[float] = []
    t_start = time.time()

    logger.info(f"Iniciando treino — {max_steps} optimizer steps "
                f"({max_steps * grad_accum} micro-batches)")

    optimizer.zero_grad()

    while global_step < max_steps:
        # Pega próxima amostra do buffer (loop circular)
        sample = samples_buffer[sample_idx % n_loaded]
        sample_idx += 1

        # Tokeniza (batch_size=1)
        try:
            input_ids, attention_mask, labels = collate_samples_to_tensors(
                [sample], tokenizer,
                max_seq_len=max_seq_len,
                pad_token_id=config.pad_token_id,
            )
        except Exception as e:
            logger.warning(f"Tokenize erro micro-batch {micro_batch_idx}: {e}")
            continue

        # Forward + backward (com grad accum)
        try:
            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs["loss"] / grad_accum  # normaliza por grad_accum
            loss.backward()

            current_loss = outputs["loss"].item()
            model.update_loss_tracking(current_loss)
            last_overall_confidence = outputs["overall_confidence"].mean().item()
            losses_history.append(current_loss)
            last_loss = current_loss
            if current_loss < best_loss:
                best_loss = current_loss
        except RuntimeError as e:
            err = str(e).lower()
            if "out of memory" in err:
                logger.warning(f"OOM micro-batch {micro_batch_idx}, limpando cache")
                gc.collect()
                optimizer.zero_grad()
                continue
            else:
                logger.error(f"Runtime erro micro-batch {micro_batch_idx}: {e}")
                traceback.print_exc()
                break

        micro_batch_idx += 1

        # Step no optimizer a cada grad_accum micro-batches
        if micro_batch_idx % grad_accum == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            global_step += 1

            # Log
            if global_step % log_every == 0 or global_step == 1:
                elapsed = time.time() - t_start
                avg_loss = sum(losses_history[-grad_accum:]) / max(1, min(grad_accum, len(losses_history)))
                ppl = math.exp(min(20, avg_loss)) if avg_loss < 20 else float("inf")
                logger.info(
                    f"  step {global_step}/{max_steps} | micro {micro_batch_idx} | "
                    f"loss={avg_loss:.4f} ppl={ppl:.2f} "
                    f"lr={scheduler.get_last_lr()[0]:.2e} "
                    f"conf={last_overall_confidence:.3f} "
                    f"punish={model.is_punishment_active()} "
                    f"elapsed={elapsed:.0f}s"
                )

            # Save checkpoint
            if global_step % save_every == 0:
                ckpt_path = Path(output_dir) / f"checkpoint_step{global_step}.pt"
                torch.save({
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "scheduler_state_dict": scheduler.state_dict(),
                    "global_step": global_step,
                    "best_loss": best_loss,
                    "last_loss": last_loss,
                    "samples_loaded": n_loaded,
                }, ckpt_path)
                logger.info(f"  Checkpoint salvo: {ckpt_path}")
                gc.collect()

    # ── Save final model ─────────────────────────────────────────────────
    elapsed_total = time.time() - t_start
    final_loss = last_loss
    final_avg_loss = sum(losses_history[-grad_accum:]) / max(1, min(grad_accum, len(losses_history)))
    final_ppl = math.exp(min(20, final_avg_loss)) if final_avg_loss < 20 else float("inf")

    logger.info("=" * 60)
    logger.info("TREINO CONCLUÍDO")
    logger.info(f"  optimizer steps: {global_step}")
    logger.info(f"  micro-batches:   {micro_batch_idx}")
    logger.info(f"  tempo total:     {elapsed_total:.1f}s ({elapsed_total/60:.1f}min)")
    logger.info(f"  best loss:       {best_loss:.4f}")
    logger.info(f"  final avg loss:  {final_avg_loss:.4f}")
    logger.info(f"  final perplexity:{final_ppl:.2f}")
    logger.info("=" * 60)

    # Salva modelo final (state_dict puro, formato HF-compatible)
    model_dir = Path(output_dir) / "model_final"
    model_dir.mkdir(parents=True, exist_ok=True)
    final_state_path = model_dir / "pytorch_model.bin"
    torch.save(model.state_dict(), final_state_path)
    logger.info(f"Modelo final salvo: {final_state_path}")

    # Salva config (mantém todos os hiperparâmetros do v13.9.2)
    config_dict = {
        "model_type": "gru_ring_v13_9_2",
        "vocab_size": config.vocab_size,
        "d_model": config.d_model,
        "n_layers": config.n_layers,
        "n_heads": config.n_heads,
        "n_kv_heads": config.n_kv_heads,
        "max_seq_len": config.max_seq_len,
        "batch_size": config.batch_size,
        "pad_token_id": config.pad_token_id,
        "d_ff1": config.d_ff1,
        "d_ff2": config.d_ff2,
        "d_ff3": config.d_ff3,
        "num_hypotheses": config.num_hypotheses,
        "num_trust_scorers": config.num_trust_scorers,
        "hypothesis_rank": config.hypothesis_rank,
        "dropout": config.dropout,
        "trust_lr": config.trust_lr,
        "punishment_patience": config.punishment_patience,
        "warmup_skip_steps": config.warmup_skip_steps,
        "total_params": model.count_parameters()["total"],
        "total_params_M": model.count_parameters()["total_M"],
        "training": {
            "version_base": "v13.9.2",
            "datasets": datasets_list,
            "per_dataset_counts": per_dataset_counts,
            "max_samples_per_dataset": max_samples,
            "max_steps": max_steps,
            "grad_accum": grad_accum,
            "per_device_batch_size": 1,
            "effective_batch_size": grad_accum,
            "lr": lr,
            "weight_decay": weight_decay,
            "max_seq_len": max_seq_len,
            "gradient_checkpointing": True,
            "global_step": global_step,
            "micro_batches": micro_batch_idx,
            "best_loss": best_loss,
            "final_loss": final_loss,
            "final_avg_loss": final_avg_loss,
            "final_perplexity": final_ppl,
            "elapsed_s": elapsed_total,
        },
    }
    with open(model_dir / "config.json", "w") as f:
        json.dump(config_dict, f, indent=2, ensure_ascii=False)

    # Copia tokenizer para o model_dir
    import shutil
    tok_dest = model_dir / "tokenizer"
    tok_dest.mkdir(exist_ok=True)
    shutil.copy(tokenizer_path, tok_dest / "tokenizer.json")

    # Salva relatório de treino
    report = {
        "version": "v13.9.2",
        "base_model": "v13.9.2 (9/9 datasets)",
        "datasets": datasets_list,
        "training_params": {
            "max_samples_per_dataset": max_samples,
            "max_steps": max_steps,
            "grad_accum": grad_accum,
            "per_device_batch_size": 1,
            "effective_batch_size": grad_accum,
            "lr": lr,
            "max_seq_len": max_seq_len,
            "gradient_checkpointing": True,
        },
        "results": {
            "global_step": global_step,
            "micro_batches": micro_batch_idx,
            "best_loss": best_loss,
            "final_loss": final_loss,
            "final_avg_loss": final_avg_loss,
            "final_perplexity": final_ppl,
            "elapsed_s": elapsed_total,
            "samples_loaded": n_loaded,
            "per_dataset_counts": per_dataset_counts,
        },
    }
    report_path = Path(output_dir) / "training_report.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2, ensure_ascii=False, default=str)
    logger.info(f"Relatório salvo: {report_path}")

    # ── W8A8 Quantization (opcional — só se solicitado) ──────────────────
    # Não vamos quantizar por padrão — o modelo fp32 já está treinado.
    # A quantização pode ser feita depois com W8A8QOperatorQuantizer.

    # ── Cleanup ──────────────────────────────────────────────────────────
    del model, optimizer, scheduler
    gc.collect()

    return report["results"]


def main():
    parser = argparse.ArgumentParser(
        description="Treina GRU-RING v13.9.2 no dataset Carolina"
    )
    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH,
                        help=f"Caminho do pytorch_model.bin (default: {DEFAULT_MODEL_PATH})")
    parser.add_argument("--tokenizer-path", default=DEFAULT_TOKENIZER_PATH,
                        help=f"Caminho do tokenizer (default: {DEFAULT_TOKENIZER_PATH})")
    parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR,
                        help=f"Diretório de saída (default: {DEFAULT_OUTPUT_DIR})")
    parser.add_argument("--datasets", type=str, default=None,
                        help=f"Lista separada por vírgulas. Default: {DEFAULT_DATASETS}")
    parser.add_argument("--max-samples", type=int, default=200,
                        help="Nº máx de amostras POR dataset (default: 200)")
    parser.add_argument("--max-steps", type=int, default=200,
                        help="Nº máx de optimizer steps (default: 200)")
    parser.add_argument("--grad-accum", type=int, default=16,
                        help="gradient_accumulation_steps (default: 16)")
    parser.add_argument("--lr", type=float, default=2e-5,
                        help="learning rate (default: 2e-5 — baixo p/ fine-tune)")
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--max-seq-len", type=int, default=64,
                        help="tamanho sequência (default: 64 — nativo do v13.9.2)")
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--save-every", type=int, default=100)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--smoke-test", action="store_true",
                        help="Smoke test: 10 samples, 5 steps, 4 grad_accum")
    args = parser.parse_args()

    if args.smoke_test:
        logger.info("SMOKE TEST MODE — 10 samples, 5 steps, 4 grad_accum")
        args.max_samples = 10
        args.max_steps = 5
        args.grad_accum = 4
        args.log_every = 1
        args.save_every = 5

    result = train_carolina(
        model_path=args.model_path,
        tokenizer_path=args.tokenizer_path,
        output_dir=args.output_dir,
        max_samples=args.max_samples,
        max_steps=args.max_steps,
        grad_accum=args.grad_accum,
        lr=args.lr,
        weight_decay=args.weight_decay,
        max_seq_len=args.max_seq_len,
        log_every=args.log_every,
        save_every=args.save_every,
        seed=args.seed,
        datasets=args.datasets,
    )
    print(f"\n=== Resultado Final ===")
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()