"""
main_nba_mid.py — MID on the NBA dataset.

Only the data pipeline and encoder change from the original MID:
  - Data: nba_train/test.npy (11 players, T=10+20)
  - Encoder: GRU ego-encoder + social-transformer (no Trajectron)
  - Diffusion: MID's DiffusionTraj + TransformerConcatLinear (unchanged, 100 steps)

Diffusion target: relative future trajectory (fut_pos - last_obs_pos) / traj_scale.
At inference: pred_rel * traj_scale + last_obs_pos → absolute positions.

Usage:
    python main_nba_mid.py --data_dir ../data/nba/original
"""

import os
import sys
import time
import logging
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
try:
    from tensorboardX import SummaryWriter
except Exception:
    from torch.utils.tensorboard import SummaryWriter
from tqdm.auto import tqdm

# MID diffusion model (unchanged)
from models.diffusion import DiffusionTraj, VarianceSchedule, TransformerConcatLinear


# ---------------------------------------------------------------------------
# Constants  (match LED's NBA normalization)
# ---------------------------------------------------------------------------

OBS_LEN    = 10
PRED_LEN   = 20
NUM_AGENTS = 11
TRAJ_SCALE = 5.0
TRAJ_MEAN  = torch.FloatTensor([14.0, 7.5])   # court-space mean after /=(94/28)
K_EVAL     = 20                                # best-of-K samples


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------

class NBADatasetMID(Dataset):
    """NBA trajectory dataset for MID.

    Loads nba_{train,test}.npy  (shape: N, 30, 11, 2).
    Returns per-sequence tensors; normalization is done in preprocess_batch.
    """

    def __init__(self, data_dir: str, training: bool = True):
        super().__init__()
        fname = 'nba_train.npy' if training else 'nba_test.npy'
        path  = os.path.join(data_dir, fname)

        trajs = np.load(path).astype(np.float32)   # (N, 30, 11, 2)
        trajs /= (94.0 / 28.0)                     # normalise court units
        # train file has 32500 scenes, test file has 12500 — use all

        trajs = torch.from_numpy(trajs).permute(0, 2, 1, 3)  # (N, 11, 30, 2)

        self.pre = trajs[:, :, :OBS_LEN, :]        # (N, 11, 10, 2)
        self.fut = trajs[:, :, OBS_LEN:,  :]       # (N, 11, 20, 2)

    def __len__(self):
        return len(self.pre)

    def __getitem__(self, idx):
        return self.pre[idx], self.fut[idx]         # each (11, T, 2)


def nba_collate(batch):
    pre = torch.stack([b[0] for b in batch])       # (B, 11, 10, 2)
    fut = torch.stack([b[1] for b in batch])       # (B, 11, 20, 2)
    return pre, fut


# ---------------------------------------------------------------------------
# Data pre-processing  (mirrors LED trainer's data_preprocess)
# ---------------------------------------------------------------------------

def preprocess_batch(pre_motion: torch.Tensor,
                     fut_motion: torch.Tensor,
                     device: torch.device):
    """
    Args:
        pre_motion: [B, A, T_obs, 2]  court-unit-normalised absolute positions
        fut_motion: [B, A, T_fut, 2]

    Returns:
        past_6ch:    [B*A, T_obs, 6]  abs + rel + vel, each / traj_scale
        fut_rel:     [B*A, T_fut, 2]  future relative to last obs, / traj_scale
        social_mask: [B*A, B*A]       additive attention mask (0 / -inf)
        last_obs:    [B*A, 1, 2]      last observed position (un-scaled)
    """
    B, A = pre_motion.shape[:2]
    traj_mean = TRAJ_MEAN.to(device)              # [2]

    pre      = pre_motion.reshape(B * A, OBS_LEN, 2)
    fut      = fut_motion.reshape(B * A, PRED_LEN, 2)
    last_obs = pre[:, -1:, :]                      # [B*A, 1, 2]

    abs_xy = (pre - traj_mean) / TRAJ_SCALE
    rel_xy = (pre - last_obs)  / TRAJ_SCALE
    vel_xy = torch.cat([rel_xy[:, 1:] - rel_xy[:, :-1],
                        torch.zeros_like(rel_xy[:, :1])], dim=1)

    past_6ch = torch.cat([abs_xy, rel_xy, vel_xy], dim=-1)   # [B*A, T, 6]
    fut_rel  = (fut - last_obs) / TRAJ_SCALE                  # [B*A, T_fut, 2]

    # block-diagonal mask: agents within the same scene can attend to each other
    mask = torch.full((B * A, B * A), float('-inf'), device=device)
    for i in range(B):
        s, e = i * A, (i + 1) * A
        mask[s:e, s:e] = 0.0

    return past_6ch, fut_rel, mask, last_obs


# ---------------------------------------------------------------------------
# NBA Encoder  (replaces Trajectron)
# ---------------------------------------------------------------------------

class _STEncoder(nn.Module):
    """Per-agent spatio-temporal encoder: Conv1D + GRU over 6-channel past.

    Input:  [N, T, 6]
    Output: [N, hidden]
    """

    def __init__(self, in_channels: int = 6, hidden: int = 256):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.gru  = nn.GRU(32, hidden, num_layers=1, batch_first=True)
        nn.init.kaiming_normal_(self.conv.weight)
        nn.init.kaiming_normal_(self.gru.weight_ih_l0)
        nn.init.kaiming_normal_(self.gru.weight_hh_l0)
        nn.init.zeros_(self.conv.bias)
        nn.init.zeros_(self.gru.bias_ih_l0)
        nn.init.zeros_(self.gru.bias_hh_l0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:   # [N, T, 6] → [N, 256]
        h = self.relu(self.conv(x.transpose(1, 2)))        # [N, 32, T]
        _, state = self.gru(h.transpose(1, 2))             # [1, N, hidden]
        return state.squeeze(0)                             # [N, hidden]


class _SocialTransformer(nn.Module):
    """Cross-agent social attention.

    Treats all N=B*11 agents as a sequence of length N with batch size 1.
    The block-diagonal mask restricts attention to within the same scene.

    Input:  x_flat [N, T*6], mask [N, N]
    Output: [N, hidden]
    """

    def __init__(self, past_len: int = OBS_LEN, hidden: int = 256):
        super().__init__()
        self.proj    = nn.Linear(past_len * 6, hidden, bias=False)
        layer        = nn.TransformerEncoderLayer(
                           d_model=hidden, nhead=2,
                           dim_feedforward=hidden, batch_first=False)
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)

    def forward(self, x_flat: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # x_flat: [N, T*6],  mask: [N, N]
        h = self.proj(x_flat).unsqueeze(1)       # [N, 1, hidden]
        # TransformerEncoder expects [seq_len, batch, dim] → treat N as seq_len
        h = h + self.encoder(h, mask)            # [N, 1, hidden]
        return h.squeeze(1)                      # [N, hidden]


class NBAEncoder(nn.Module):
    """NBA social encoder → context [B*A, encoder_dim] for DiffusionTraj.

    Architecture:
        ego_embed    = st_encoder(past_6ch)          [N, 256]
        social_embed = social_transformer(past_6ch)  [N, 256]
        context      = Linear(512 → encoder_dim)     [N, encoder_dim]
    """

    def __init__(self, encoder_dim: int = 256, past_len: int = OBS_LEN):
        super().__init__()
        self.ego_encoder    = _STEncoder(in_channels=6, hidden=256)
        self.social_encoder = _SocialTransformer(past_len=past_len, hidden=256)
        self.fusion         = nn.Linear(512, encoder_dim)

    def forward(self, past_6ch: torch.Tensor,
                social_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            past_6ch:    [N, T, 6]
            social_mask: [N, N]  additive mask (0 / -inf)
        Returns:
            [N, encoder_dim]
        """
        ego    = self.ego_encoder(past_6ch)                       # [N, 256]
        social = self.social_encoder(
                     past_6ch.reshape(past_6ch.size(0), -1),      # [N, T*6]
                     social_mask)                                  # [N, 256]
        return self.fusion(torch.cat([ego, social], dim=-1))      # [N, encoder_dim]


# ---------------------------------------------------------------------------
# Trainer
# ---------------------------------------------------------------------------

class Trainer:

    def __init__(self, args):
        self.args   = args
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self._build_dirs()
        self._build_data()
        self._build_model()
        self._build_optimizer()

    def _build_dirs(self):
        self.exp_dir = os.path.join('experiments', self.args.exp_name)
        os.makedirs(self.exp_dir, exist_ok=True)
        self.tb_log  = SummaryWriter(log_dir=self.exp_dir)

        log_path = os.path.join(
            self.exp_dir,
            'nba_{}.log'.format(time.strftime('%Y-%m-%d-%H-%M')))
        self.log = logging.getLogger(self.args.exp_name)
        self.log.setLevel(logging.INFO)
        self.log.addHandler(logging.FileHandler(log_path))
        self.log.addHandler(logging.StreamHandler(sys.stdout))
        self.log.info(f"Args: {self.args}")

    def _build_data(self):
        train_dset = NBADatasetMID(self.args.data_dir, training=True)
        test_dset  = NBADatasetMID(self.args.data_dir, training=False)

        self.train_loader = DataLoader(
            train_dset, batch_size=self.args.batch_size,
            shuffle=True,  num_workers=4,
            collate_fn=nba_collate, pin_memory=True)
        self.test_loader  = DataLoader(
            test_dset,  batch_size=self.args.eval_batch_size,
            shuffle=False, num_workers=4,
            collate_fn=nba_collate, pin_memory=True)

        self.log.info(
            f"Train: {len(train_dset)} scenes  Test: {len(test_dset)} scenes")

    def _build_model(self):
        self.encoder = NBAEncoder(
            encoder_dim = self.args.encoder_dim,
            past_len    = OBS_LEN,
        ).to(self.device)

        net = TransformerConcatLinear(
            point_dim   = 2,
            context_dim = self.args.encoder_dim,
            tf_layer    = self.args.tf_layer,
            residual    = False,
        )
        self.diffusion = DiffusionTraj(
            net       = net,
            var_sched = VarianceSchedule(
                num_steps = 100,
                beta_T    = 5e-2,
                mode      = 'linear',
            ),
        ).to(self.device)

        n_enc  = sum(p.numel() for p in self.encoder.parameters())
        n_diff = sum(p.numel() for p in self.diffusion.parameters())
        self.log.info(f"Encoder params:   {n_enc:,}")
        self.log.info(f"Diffusion params: {n_diff:,}")

    def _build_optimizer(self):
        params = (list(self.encoder.parameters()) +
                  list(self.diffusion.parameters()))
        self.optimizer = optim.Adam(params, lr=self.args.lr)
        self.scheduler = optim.lr_scheduler.ExponentialLR(
            self.optimizer, gamma=0.98)

    # ------------------------------------------------------------------

    def train(self):
        for epoch in range(1, self.args.epochs + 1):
            self.encoder.train()
            self.diffusion.train()

            total_loss, count = 0.0, 0
            pbar = tqdm(self.train_loader, ncols=90)
            for pre, fut in pbar:
                pre = pre.to(self.device)
                fut = fut.to(self.device)

                past_6ch, fut_rel, mask, _ = preprocess_batch(
                    pre, fut, self.device)

                context = self.encoder(past_6ch, mask)
                loss    = self.diffusion.get_loss(fut_rel, context)

                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(
                    list(self.encoder.parameters()) +
                    list(self.diffusion.parameters()), 1.0)
                self.optimizer.step()

                total_loss += loss.item()
                count      += 1
                pbar.set_description(
                    f"Epoch {epoch} loss={total_loss/count:.4f}")

            self.scheduler.step()
            avg_loss = total_loss / count
            self.tb_log.add_scalar('loss/train', avg_loss, epoch)
            self.log.info(f"Epoch {epoch}  train_loss={avg_loss:.4f}")

            if epoch % self.args.eval_every == 0:
                ade, fde = self.evaluate()
                self.tb_log.add_scalar('metric/ADE', ade, epoch)
                self.tb_log.add_scalar('metric/FDE', fde, epoch)
                self.log.info(
                    f"Epoch {epoch}  ADE={ade:.4f}  FDE={fde:.4f}")

                torch.save({
                    'encoder':   self.encoder.state_dict(),
                    'diffusion': self.diffusion.state_dict(),
                    'epoch':     epoch,
                }, os.path.join(self.exp_dir, f'nba_epoch{epoch}.pt'))

    @torch.no_grad()
    def evaluate(self):
        self.encoder.eval()
        self.diffusion.eval()

        ade_sum, fde_sum, n_agents = 0.0, 0.0, 0

        for pre, fut in tqdm(self.test_loader, ncols=90, desc='Eval'):
            pre = pre.to(self.device)
            fut = fut.to(self.device)
            B   = pre.size(0)

            past_6ch, _, mask, last_obs = preprocess_batch(
                pre, fut, self.device)
            context = self.encoder(past_6ch, mask)      # [B*11, enc_dim]

            # [K, B*11, T_fut, 2]
            pred_rel = self.diffusion.sample(
                num_points    = PRED_LEN,
                context       = context,
                sample        = K_EVAL,
                bestof        = True,
                sampling      = self.args.sampling,
                step          = self.args.sampling_step,
            )

            # absolute positions
            pred_abs = pred_rel * TRAJ_SCALE + last_obs.unsqueeze(0)   # [K, B*11, T, 2]
            fut_abs  = fut.reshape(B * NUM_AGENTS, PRED_LEN, 2)        # [B*11, T, 2]

            dist     = (pred_abs - fut_abs.unsqueeze(0)).norm(dim=-1)  # [K, B*11, T]

            # minADE: pick best single trajectory (mean over T, then min over K)
            ade_per_mode = dist.mean(dim=-1)                           # [K, B*11]
            ade_sum  += ade_per_mode.min(dim=0).values.sum().item()

            # minFDE: pick best mode at last frame
            fde_sum  += dist[:, :, -1].min(dim=0).values.sum().item()
            n_agents += B * NUM_AGENTS

        ade = ade_sum / n_agents
        fde = fde_sum / n_agents
        return ade, fde


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser()

    # Data
    p.add_argument('--data_dir', type=str, default='../data/nba/original',
                   help='Directory containing nba_train.npy / nba_test.npy')

    # Experiment
    p.add_argument('--exp_name', type=str, default='mid_nba')

    # Training
    p.add_argument('--epochs',          type=int,   default=100)
    p.add_argument('--batch_size',      type=int,   default=32)
    p.add_argument('--eval_batch_size', type=int,   default=64)
    p.add_argument('--lr',              type=float, default=1e-3)
    p.add_argument('--eval_every',      type=int,   default=5)

    # Model
    p.add_argument('--encoder_dim', type=int, default=256)
    p.add_argument('--tf_layer',    type=int, default=3)

    # Sampling
    p.add_argument('--sampling',      type=str, default='ddim',
                   choices=['ddpm', 'ddim'])
    p.add_argument('--sampling_step', type=int, default=10,
                   help='Inference steps (100=full DDPM, <100=DDIM stride)')

    return p.parse_args()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == '__main__':
    args    = parse_args()
    trainer = Trainer(args)
    trainer.train()