#!/usr/bin/env python3
"""T2M-GPT stage 2 on Full_TriVis: Vietnamese gloss -> pose-token GPT.

Same model and recipe as the paper's `train_t2m_trans.py`: the frozen stage-1
VQ-VAE supplies discrete targets, a causal transformer is trained with CE on
next-token prediction, and inputs are corrupted with probability 1-pkeep
(the paper's "corruption" trick that lets the model recover from its own
generation errors).

Two substitutions were unavoidable:
  * CLIP's text tower is English-only, so conditioning comes from a frozen
    Vietnamese encoder (PhoBERT by default). As in the paper, this is a single
    pooled sentence vector prepended to the token sequence.
  * The paper's validation uses HumanML3D FID / R-precision evaluators, which do
    not exist for VSL. Validation here is teacher-forced CE + token accuracy,
    plus free-running generation decoded through the VQ-VAE and scored with
    MPJPE against ground truth (and a length-ratio, since over-generation was the
    dominant failure of this project's earlier gloss->pose baseline).
"""
import argparse
import json
import os

import numpy as np
import torch
from torch.distributions import Categorical
from torch.utils.tensorboard import SummaryWriter

import models.t2m_trans as trans
import models.vqvae as vqvae
import options.option_vsl as option_vsl
import utils.utils_model as utils_model
from dataset import dataset_vsl
from models.text_encoder_vi import ViTextEncoder
from utils.losses_vsl import mpjpe_groups


def build_vqvae(resume_pth, device):
    ckpt = torch.load(resume_pth, map_location='cpu')
    targs = argparse.Namespace(**ckpt['args'])
    net = vqvae.HumanVQVAE(targs, targs.nb_code, targs.code_dim, targs.output_emb_width,
                           targs.down_t, targs.stride_t, targs.width, targs.depth,
                           targs.dilation_growth_rate, targs.vq_act, targs.vq_norm,
                           input_dim=targs.input_dim)
    net.load_state_dict(ckpt['net'], strict=True)
    net.eval().to(device)
    for p in net.parameters():
        p.requires_grad = False
    return net, targs, ckpt


def masked_ce(logits, targets, tlen):
    """Per-sample mean CE over real tokens (+ end token), then mean over batch.

    Vectorized equivalent of upstream's `for i in range(bs)` loop -- same value,
    but ~4x faster per iteration, which matters at 60k iters.
    Returns (loss, n_correct, n_tokens).
    """
    B, L, V = logits.shape
    valid = (torch.arange(L, device=logits.device)[None, :] < (tlen[:, None] + 1)).float()
    # Padding uses id nb_code+1, which is outside the (nb_code+1)-way head; those
    # positions are masked out below, so clamp them to keep CE's index check happy.
    targets = targets.clamp(max=V - 1)
    ce = torch.nn.functional.cross_entropy(
        logits.reshape(B * L, -1), targets.reshape(B * L), reduction='none').view(B, L)
    loss = ((ce * valid).sum(1) / valid.sum(1).clamp(min=1)).mean()
    with torch.no_grad():
        correct = int((((logits.argmax(-1) == targets).float()) * valid).sum())
    return loss, correct, int(valid.sum())


@torch.no_grad()
def validate_teacher_forced(trans_encoder, loader, text_enc, device, max_batches=40):
    """Teacher-forced CE + token accuracy on the val split."""
    trans_encoder.eval()
    tot_loss, tot_right, tot_tok, nb = 0.0, 0, 0, 0
    for bi, (texts, tokens, tlen) in enumerate(loader):
        if bi >= max_batches:
            break
        tokens, tlen = tokens.to(device), tlen.to(device)
        feat = text_enc(list(texts))
        logits = trans_encoder(tokens[:, :-1], feat)
        loss, correct, ntok = masked_ce(logits, tokens[:, :logits.shape[1]], tlen)
        tot_loss += loss.item() * tokens.shape[0]
        tot_right += correct
        tot_tok += ntok
        nb += tokens.shape[0]
    trans_encoder.train()
    return tot_loss / max(nb, 1), 100.0 * tot_right / max(tot_tok, 1)


@torch.no_grad()
def validate_generation(trans_encoder, net, store, items, text_enc, device,
                        max_tokens, nb_code, text_field='gloss', if_categorial=False):
    """Free-running generation -> VQ decode -> MPJPE vs ground truth.

    Scored on the overlapping prefix of generated and reference motion, so a
    length mismatch shows up in len_ratio rather than silently inflating MPJPE.
    """
    trans_encoder.eval()
    mean = torch.from_numpy(store.mean).to(device)
    std = torch.from_numpy(store.std).to(device)
    acc = {"all": 0.0, "body": 0.0, "face": 0.0, "hands": 0.0}
    ratios, n_used, n_empty = [], 0, 0

    for i in items:
        c = store.index[i]
        feat = text_enc([c[text_field]])
        idx = trans_encoder.sample(feat, if_categorial=if_categorial)
        if idx is None or idx.numel() == 0:
            n_empty += 1
            continue
        idx = idx.clamp(max=nb_code - 1)  # drop any end/pad id that slipped through
        pred = net.decode_batch(idx)  # (1, L*4, 256)

        motion, mask = store.get(i)
        gt = torch.from_numpy(motion).unsqueeze(0).to(device)
        mk = torch.from_numpy(mask).unsqueeze(0).to(device)
        ratios.append(pred.shape[1] / gt.shape[1])

        T = min(pred.shape[1], gt.shape[1])
        nk = store.layout.n_kpts
        valid = mk[:, :T].view(1, T, nk, 2)[..., 0]
        g = mpjpe_groups(pred[:, :T] * std + mean, gt[:, :T] * std + mean, valid,
                         groups=store.layout.metric_groups())
        for k in acc:
            acc[k] += g[k]
        n_used += 1

    trans_encoder.train()
    n = max(n_used, 1)
    return ({k: v / n for k, v in acc.items()},
            float(np.mean(ratios)) if ratios else 0.0, n_used, n_empty)


def main():
    args = option_vsl.get_trans_args()
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    device = torch.device(args.device)

    args.out_dir = os.path.join(args.out_dir, args.exp_name)
    os.makedirs(args.out_dir, exist_ok=True)
    logger = utils_model.get_logger(args.out_dir)
    writer = SummaryWriter(args.out_dir)
    logger.info(json.dumps(vars(args), indent=4, sort_keys=True))

    ##### ---- Frozen stage-1 VQ-VAE ---- #####
    net, targs, ckpt = build_vqvae(args.resume_pth, device)
    assert targs.nb_code == args.nb_code, \
        f'--nb-code {args.nb_code} != stage-1 codebook {targs.nb_code}'
    logger.info(f"stage-1: iter {ckpt.get('iter')} val_recon {ckpt.get('val_recon'):.5f} "
                f"hands MPJPE {ckpt['val_mpjpe']['hands']:.5f} "
                f"codes {ckpt.get('codes_used')}/{targs.nb_code}  <-- generation ceiling")

    ##### ---- Frozen Vietnamese text encoder (CLIP replacement) ---- #####
    text_enc = ViTextEncoder(args.text_model, device=args.device)
    logger.info(f'text encoder {args.text_model}, dim {text_enc.dim}')

    ##### ---- Dataloaders ---- #####
    train_set = dataset_vsl.VSLText2TokenDataset(
        args.data_dir, os.path.join(args.token_dir, 'train'), 'train',
        args.nb_code, max_tokens=args.max_tokens, text_field=args.text_field)
    train_loader = torch.utils.data.DataLoader(
        train_set, args.batch_size, shuffle=True, num_workers=args.num_workers,
        drop_last=True, pin_memory=True, persistent_workers=args.num_workers > 0)
    train_loader_iter = dataset_vsl.cycle(train_loader)

    val_set = dataset_vsl.VSLText2TokenDataset(
        args.data_dir, os.path.join(args.token_dir, 'val'), 'val',
        args.nb_code, max_tokens=args.max_tokens, text_field=args.text_field,
        augment_crop=False)
    val_loader = torch.utils.data.DataLoader(val_set, args.batch_size, shuffle=False,
                                             num_workers=4, drop_last=False)
    rng = np.random.RandomState(0)
    gen_items = sorted(rng.choice(len(val_set.store.index),
                                  size=min(args.eval_samples, len(val_set.store.index)),
                                  replace=False).tolist())

    ##### ---- GPT ---- #####
    # block_size must hold the conditioning vector + the whole token sequence
    block_size = args.max_tokens + 1
    trans_encoder = trans.Text2Motion_Transformer(
        num_vq=args.nb_code, embed_dim=args.embed_dim_gpt, clip_dim=text_enc.dim,
        block_size=block_size, num_layers=args.num_layers, n_head=args.n_head_gpt,
        drop_out_rate=args.drop_out_rate, fc_rate=args.ff_rate)
    if args.resume_trans:
        logger.info(f'loading transformer checkpoint from {args.resume_trans}')
        tck = torch.load(args.resume_trans, map_location='cpu')
        trans_encoder.load_state_dict(tck['trans'], strict=True)
    elif getattr(args, 'init_trans', None):
        # Warm start from a stage-2 model trained on a DIFFERENT corpus (e.g. the
        # Multi-VSL word model -> TriVis sentences). Only shape-compatible tensors are
        # copied; what differs is non-learned anyway -- the causal mask buffers, the
        # deterministic sincos positional encoding, and an unused pos_embedding whose
        # sizes follow block_size. Note the token embedding and output head transfer in
        # SHAPE but not in meaning: the 512 codes index a different codebook, so this
        # is a warm start of the transformer body, not of the vocabulary.
        src = torch.load(args.init_trans, map_location='cpu')['trans']
        tgt = trans_encoder.state_dict()
        ok = {k: v for k, v in src.items() if k in tgt and tgt[k].shape == v.shape}
        skipped = [k for k in tgt if k not in ok]
        trans_encoder.load_state_dict({**tgt, **ok}, strict=True)
        n_ok = sum(ok[k].numel() for k in ok)
        n_all = sum(v.numel() for v in tgt.values())
        logger.info(f'warm start from {args.init_trans}: copied {len(ok)}/{len(tgt)} '
                    f'tensors = {100*n_ok/n_all:.1f}% of params; '
                    f'skipped (non-learned/shape): {len(skipped)}')
    trans_encoder.train().to(device)
    logger.info(f'GPT params: {sum(p.numel() for p in trans_encoder.parameters())/1e6:.2f}M, '
                f'block_size {block_size}')

    optimizer = utils_model.initial_optim(args.decay_option, args.lr, args.weight_decay,
                                         trans_encoder, args.optimizer)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_scheduler,
                                                     gamma=args.gamma)

    ##### ---- Training ---- #####
    best_hands, best_iter = 1e9, 0
    avg_loss, right_num, nb_sample = 0.0, 0, 0

    for nb_iter in range(1, args.total_iter + 1):
        texts, tokens, tlen = next(train_loader_iter)
        tokens, tlen = tokens.to(device), tlen.to(device)
        bs = tokens.shape[0]
        feat = text_enc(list(texts))

        input_index = tokens[:, :-1]
        # T2M-GPT corruption: replace a fraction of input tokens with random codes
        if args.pkeep == -1:
            proba = np.random.rand(1)[0]
            mask = torch.bernoulli(proba * torch.ones_like(input_index, dtype=torch.float))
        else:
            mask = torch.bernoulli(args.pkeep * torch.ones_like(input_index, dtype=torch.float))
        mask = mask.round().to(dtype=torch.int64)
        r_indices = torch.randint_like(input_index, args.nb_code)
        a_indices = mask * input_index + (1 - mask) * r_indices

        logits = trans_encoder(a_indices, feat).contiguous()

        # CE over the real tokens + the end token only (padding is ignored)
        loss_cls, correct, ntok = masked_ce(logits, tokens[:, :logits.shape[1]], tlen)

        optimizer.zero_grad()
        loss_cls.backward()
        optimizer.step()
        scheduler.step()

        avg_loss += loss_cls.item()
        right_num += correct
        nb_sample += ntok

        if nb_iter % args.print_iter == 0:
            avg_loss /= args.print_iter
            acc = 100.0 * right_num / max(nb_sample, 1)
            writer.add_scalar('./Loss/train', avg_loss, nb_iter)
            writer.add_scalar('./ACC/train', acc, nb_iter)
            logger.info(f"Train. Iter {nb_iter} : Loss. {avg_loss:.5f}, ACC. {acc:.4f}")
            avg_loss, right_num, nb_sample = 0.0, 0, 0

        if nb_iter % args.eval_iter == 0 or nb_iter == args.total_iter:
            vl, vacc = validate_teacher_forced(trans_encoder, val_loader, text_enc, device)
            # Both decoders are reported. Selection uses categorial sampling, which
            # is what upstream's evaluation_transformer_test (the paper's reported
            # test protocol) uses: greedy almost never emits the end token here, so
            # it runs to the block limit and its metrics track length, not quality.
            gen = {}
            for name, cat in (('cat', True), ('greedy', False)):
                gen[name] = validate_generation(
                    trans_encoder, net, val_set.store, gen_items, text_enc, device,
                    args.max_tokens, args.nb_code, args.text_field, if_categorial=cat)
            mp, ratio, n_used, n_empty = gen['cat']
            gmp, gratio = gen['greedy'][0], gen['greedy'][1]

            writer.add_scalar('./Loss/val', vl, nb_iter)
            writer.add_scalar('./ACC/val', vacc, nb_iter)
            writer.add_scalar('./Gen/MPJPE_hands', mp['hands'], nb_iter)
            writer.add_scalar('./Gen/len_ratio', ratio, nb_iter)
            writer.add_scalar('./GenGreedy/MPJPE_hands', gmp['hands'], nb_iter)
            writer.add_scalar('./GenGreedy/len_ratio', gratio, nb_iter)
            logger.info(f"Eval. Iter {nb_iter} : val_loss {vl:.5f} val_acc {vacc:.4f} "
                        f"| CAT MPJPE all {mp['all']:.5f} body {mp['body']:.5f} "
                        f"hands {mp['hands']:.5f} len_ratio {ratio:.3f} "
                        f"| GREEDY hands {gmp['hands']:.5f} len_ratio {gratio:.3f} "
                        f"({n_used} scored, {n_empty} empty)")
            payload = {'trans': trans_encoder.state_dict(), 'args': vars(args),
                       'iter': nb_iter, 'val_loss': vl, 'val_acc': vacc,
                       'gen_mpjpe': mp, 'len_ratio': ratio,
                       'gen_mpjpe_greedy': gmp, 'len_ratio_greedy': gratio}
            torch.save(payload, os.path.join(args.out_dir, 'net_last.pth'))
            if mp['hands'] < best_hands:
                best_hands, best_iter = mp['hands'], nb_iter
                torch.save(payload, os.path.join(args.out_dir, 'net_best.pth'))
                logger.info(f"  --> new best (gen hands MPJPE {best_hands:.5f})")
            # Periodic immortal snapshots. Without these, "does training longer help?"
            # is unanswerable after the fact: net_last is overwritten and net_best is
            # picked by the flat/noisy 80-clip metric (it has landed anywhere from iter
            # 2000 to 28000). Snapshots let the question be scored at fixed iterations
            # on the real 300-clip protocol -- and on FGD / shuffle penalty, which
            # diverge from DTW.
            if args.save_every and nb_iter % args.save_every == 0:
                snap = os.path.join(args.out_dir, f'net_iter{nb_iter}.pth')
                torch.save(payload, snap)
                logger.info(f"  --> snapshot {os.path.basename(snap)}")

    logger.info(f"Done. best gen hands MPJPE {best_hands:.5f} @ iter {best_iter}")


if __name__ == '__main__':
    main()