#!/usr/bin/env python3
"""
End-to-end training script for FSD-Level5-CoT on SADC driving data.

This script:
  1. Downloads a subset of the SADC dataset (streaming → disk)
  2. Builds the FSD model from fsd_model/
  3. Trains end-to-end with gradient accumulation, warmup, eval, logging
  4. Pushes the trained model to Hugging Face Hub

Dataset: jHaselberger/SADC-Situation-Awareness-for-Driver-Centric-Driving-Style-Adaptation
Model:   Reality123b/FSD-Level5-CoT

Usage:
    # Default (5000 train, 1000 val, 5 epochs)
    python train_sadc_e2e.py

    # Custom
    python train_sadc_e2e.py --train_samples 10000 --val_samples 2000 --epochs 10 --batch_size 4

    # Quick test run
    python train_sadc_e2e.py --train_samples 100 --val_samples 50 --epochs 1
"""

import os
import sys
import time
import json
import math
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Config defaults
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

DATASET_NAME = "jHaselberger/SADC-Situation-Awareness-for-Driver-Centric-Driving-Style-Adaptation"
HUB_MODEL_ID = "Reality123b/FSD-Level5-CoT"

# Model architecture
BEV_SIZE = 100
BEV_FEATURE_DIM = 128
PLANNING_D_MODEL = 128
IMG_H, IMG_W = 120, 160
NUM_WAYPOINTS = 20
COT_ACTOR_QUERIES = 32
COT_ROAD_QUERIES = 16

# Speed constant
MAX_SPEED_MS = 20.0 * 0.44704  # 20 mph → m/s

ROAD_TYPE_MAP = {
    "misc": 0, "rural": 1, "federal": 2, "highway": 3,
    "city": 4, "parking": 5, "intersection": 6,
}


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Step 1: Download SADC Subset
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def download_sadc_subset(train_samples, val_samples, output_dir, train_split, val_split):
    """Download a manageable subset of SADC via streaming."""
    from datasets import load_dataset, Dataset as HFDataset

    os.makedirs(output_dir, exist_ok=True)
    train_path = os.path.join(output_dir, "train")
    val_path = os.path.join(output_dir, "val")

    # Check if already downloaded
    if os.path.exists(train_path) and os.path.exists(val_path):
        print(f"[Download] Found existing subset at {output_dir}, skipping download.")
        from datasets import load_from_disk
        return load_from_disk(train_path), load_from_disk(val_path)

    # Train
    print(f"[Download] Streaming {train_samples} train samples from '{train_split}'...")
    ds_stream = load_dataset(DATASET_NAME, split=train_split, streaming=True)
    train_rows = []
    for i, row in enumerate(ds_stream):
        if i >= train_samples:
            break
        train_rows.append(row)
        if (i + 1) % 1000 == 0:
            print(f"  ... {i + 1}/{train_samples}")
    train_ds = HFDataset.from_list(train_rows)
    train_ds.save_to_disk(train_path)
    print(f"  Saved {len(train_ds)} train samples.")

    # Val
    print(f"[Download] Streaming {val_samples} val samples from '{val_split}'...")
    ds_stream = load_dataset(DATASET_NAME, split=val_split, streaming=True)
    val_rows = []
    for i, row in enumerate(ds_stream):
        if i >= val_samples:
            break
        val_rows.append(row)
        if (i + 1) % 500 == 0:
            print(f"  ... {i + 1}/{val_samples}")
    val_ds = HFDataset.from_list(val_rows)
    val_ds.save_to_disk(val_path)
    print(f"  Saved {len(val_ds)} val samples.")

    return train_ds, val_ds


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Step 2: Dataset wrapper
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class SADCDrivingDataset(Dataset):
    """Wraps SADC HF dataset → FSD model inputs + targets."""

    def __init__(self, hf_dataset, img_size=(IMG_H, IMG_W)):
        self.ds = hf_dataset
        self.img_h, self.img_w = img_size

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

    def __getitem__(self, idx):
        row = self.ds[idx]

        # ── Image ──
        img = row.get("frame", None)
        if img is None:
            img_tensor = torch.zeros(3, self.img_h, self.img_w)
        else:
            from torchvision import transforms
            transform = transforms.Compose([
                transforms.Resize((self.img_h, self.img_w)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ])
            try:
                if hasattr(img, "convert"):
                    img = img.convert("RGB")
                img_tensor = transform(img)
            except Exception:
                img_tensor = torch.zeros(3, self.img_h, self.img_w)

        # Replicate to 6 virtual cameras with slight noise
        camera_images = img_tensor.unsqueeze(0).expand(6, -1, -1, -1).clone()
        for i in range(1, 6):
            camera_images[i] += torch.randn_like(camera_images[i]) * 0.01

        # ── Ego state ──
        speed_ms = float(row.get("v_kmph", 0.0)) / 3.6
        ax = float(row.get("ax_mpss", 0.0))
        steering = float(row.get("steering_rack_pos_m", 0.0))
        yaw_rate = float(row.get("yaw_rate_radps", 0.0))
        lane_center = float(row.get("d_lanecenter_m", 0.0))
        curvature = float(row.get("lane_curvature_radpm", 0.0))

        ego_state = torch.tensor([
            speed_ms, ax, steering, yaw_rate, 0.0, lane_center,
        ], dtype=torch.float32)

        # ── Navigation command ──
        road_type = str(row.get("road_type", "misc"))
        nav_cmd = ROAD_TYPE_MAP.get(road_type, 0)

        # ── Camera intrinsics / extrinsics (synthetic) ──
        K = torch.zeros(6, 3, 3)
        K[:, 0, 0] = 200.0
        K[:, 1, 1] = 200.0
        K[:, 0, 2] = self.img_w / 2
        K[:, 1, 2] = self.img_h / 2
        K[:, 2, 2] = 1.0

        E = torch.eye(4).unsqueeze(0).expand(6, -1, -1).clone()
        yaw_offsets = [-45, 45, -135, 135, -90, 90]
        for i, yaw_deg in enumerate(yaw_offsets):
            yaw_r = math.radians(yaw_deg)
            E[i, 0, 0] = math.cos(yaw_r)
            E[i, 0, 1] = -math.sin(yaw_r)
            E[i, 1, 0] = math.sin(yaw_r)
            E[i, 1, 1] = math.cos(yaw_r)

        # ── Ultrasonic (simulated) ──
        base_dist = max(0.5, abs(lane_center))
        us_distances = torch.ones(20, 1) * base_dist
        us_distances[:7] = torch.clamp(torch.randn(7, 1) * 0.5 + 3.0, 0.3, 5.0)
        us_distances[7:14] = torch.clamp(torch.randn(7, 1) * 0.5 + 3.5, 0.3, 5.0)
        us_distances[14:17] = torch.clamp(torch.tensor([[base_dist]] * 3) + torch.randn(3, 1) * 0.2, 0.3, 5.0)
        us_distances[17:20] = torch.clamp(torch.tensor([[base_dist]] * 3) + torch.randn(3, 1) * 0.2, 0.3, 5.0)

        us_placements = torch.zeros(20, 6)
        for i in range(7):
            us_placements[i] = torch.tensor([2.25, (i - 3) * 0.3, 0.4, (i - 3) * 10, 0, 0])
        for i in range(7):
            us_placements[7 + i] = torch.tensor([-2.25, (i - 3) * 0.3, 0.4, 180 + (i - 3) * 10, 0, 0])
        for i in range(3):
            us_placements[14 + i] = torch.tensor([(1 - i) * 1.0, 0.9, 0.6, -90, 0, 0])
            us_placements[17 + i] = torch.tensor([(1 - i) * 1.0, -0.9, 0.6, 90, 0, 0])

        # ── Ground truth targets ──
        gt_steering = torch.tensor(steering * 20.0)
        gt_throttle = torch.tensor(max(0.0, ax / 3.0)).clamp(0, 1)
        gt_brake = torch.tensor(max(0.0, -ax / 8.0)).clamp(0, 1)

        gt_waypoints = torch.zeros(NUM_WAYPOINTS, 4)
        for t in range(NUM_WAYPOINTS):
            dt = (t + 1) * 0.5
            gt_waypoints[t, 0] = speed_ms * dt
            gt_waypoints[t, 1] = -lane_center * min(1.0, dt / 3.0)
            gt_waypoints[t, 2] = curvature * speed_ms * dt
            gt_waypoints[t, 3] = min(speed_ms, MAX_SPEED_MS)

        if abs(steering) > 0.3:
            gt_behavior = 1 if steering > 0 else 2
        elif abs(ax) < 0.1 and speed_ms < 0.5:
            gt_behavior = 5
        else:
            gt_behavior = 0

        bev = BEV_SIZE
        gt_seg = torch.zeros(bev, bev, dtype=torch.long)
        gt_seg[bev // 4 : 3 * bev // 4, :] = 1

        gt_heatmap = torch.zeros(10, bev, bev)

        gt_occ = torch.zeros(1, bev, bev)
        gt_occ[:, : bev // 4, :] = 1.0
        gt_occ[:, 3 * bev // 4 :, :] = 1.0

        inputs = {
            "camera_images": camera_images,
            "camera_intrinsics": K,
            "camera_extrinsics": E,
            "ultrasonic_distances": us_distances,
            "ultrasonic_placements": us_placements,
            "ego_state": ego_state,
            "nav_command": torch.tensor(nav_cmd, dtype=torch.long),
        }

        targets = {
            "gt_steering": gt_steering,
            "gt_throttle": gt_throttle,
            "gt_brake": gt_brake,
            "gt_waypoints": gt_waypoints,
            "gt_behavior": torch.tensor(gt_behavior, dtype=torch.long),
            "gt_segmentation": gt_seg,
            "gt_heatmap": gt_heatmap,
            "gt_occupancy": gt_occ,
        }

        return inputs, targets


def collate_fn(batch):
    inputs_list, targets_list = zip(*batch)
    inputs = {k: torch.stack([d[k] for d in inputs_list]) for k in inputs_list[0]}
    targets = {k: torch.stack([d[k] for d in targets_list]) for k in targets_list[0]}
    return inputs, targets


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Step 3: Training
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

@torch.no_grad()
def evaluate(model, loss_fn, val_loader, device, max_batches=50):
    model.eval()
    losses = []
    for i, (inputs, targets) in enumerate(val_loader):
        if i >= max_batches:
            break
        inputs = {k: v.to(device, non_blocking=True) for k, v in inputs.items()}
        targets = {k: v.to(device, non_blocking=True) for k, v in targets.items()}
        try:
            output = model(**inputs)
            l = loss_fn(output, targets)
            losses.append(l["total"].item())
        except RuntimeError:
            continue
    return np.mean(losses) if losses else float("inf")


def train(args, train_ds, val_ds):
    """Build model and run training loop."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n[Train] Device: {device}")
    if device.type == "cuda":
        print(f"  GPU: {torch.cuda.get_device_name()}")
        print(f"  VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")

    # ── Tracking ──
    HAS_TRACKIO = False
    try:
        import trackio
        trackio.init(project="fsd-level5-cot", name="sadc-e2e-training")
        HAS_TRACKIO = True
        print("  Trackio initialized ✓")
    except Exception as e:
        print(f"  Trackio not available: {e}")

    # ── Datasets + Loaders ──
    train_dataset = SADCDrivingDataset(train_ds)
    val_dataset = SADCDrivingDataset(val_ds)

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=True,
    )
    print(f"  Train batches/epoch: {len(train_loader)}")
    print(f"  Val batches: {len(val_loader)}")

    # ── Build model ──
    print("\n[Train] Building FSD model...")
    script_dir = os.path.dirname(os.path.abspath(__file__))
    if script_dir not in sys.path:
        sys.path.insert(0, script_dir)

    from fsd_model.config import VehicleConfig
    from fsd_model.model import FullSelfDrivingModel, FSDLoss

    config = VehicleConfig()
    model = FullSelfDrivingModel(
        vehicle_config=config,
        bev_size=BEV_SIZE,
        bev_resolution=0.5,
        bev_feature_dim=BEV_FEATURE_DIM,
        num_object_classes=10,
        num_seg_classes=7,
        num_waypoints=NUM_WAYPOINTS,
        planning_d_model=PLANNING_D_MODEL,
        future_steps=6,
        num_forecast_modes=6,
        forecast_steps=12,
        num_behaviors=10,
        enable_cot=True,
        cot_num_actor_queries=COT_ACTOR_QUERIES,
        cot_num_road_queries=COT_ROAD_QUERIES,
    ).to(device)

    param_info = model.count_parameters()
    total_params = param_info["total"]
    print(f"  Total parameters: {total_params:,}")

    # ── Loss ──
    loss_fn = FSDLoss(
        learnable_weights=True,
        w_detection=0.5,
        w_segmentation=1.0,
        w_occupancy=1.0,
        w_motion=0.5,
        w_behavior=1.0,
        w_trajectory=3.0,
        w_control=3.0,
        w_safety=2.0,
    ).to(device)

    # ── Optimizer + Scheduler ──
    all_params = list(model.parameters()) + list(loss_fn.parameters())
    optimizer = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=args.weight_decay)

    total_steps = len(train_loader) * args.epochs // args.grad_accum
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer,
        max_lr=args.lr,
        total_steps=total_steps + 10,
        pct_start=0.1,
        anneal_strategy="cos",
    )

    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()

    # ── Training loop ──
    effective_batch = args.batch_size * args.grad_accum
    print(f"\n[Train] Starting: {args.epochs} epochs, effective batch={effective_batch}")
    print(f"  Total optimizer steps: ~{total_steps}")

    global_step = 0
    best_val_loss = float("inf")
    t0 = time.time()

    for epoch in range(args.epochs):
        model.train()
        epoch_losses = []
        optimizer.zero_grad()

        for batch_idx, (inputs, targets) in enumerate(train_loader):
            inputs = {k: v.to(device, non_blocking=True) for k, v in inputs.items()}
            targets = {k: v.to(device, non_blocking=True) for k, v in targets.items()}

            try:
                output = model(**inputs)
                losses = loss_fn(output, targets)
                loss = losses["total"] / args.grad_accum
            except RuntimeError as e:
                if "out of memory" in str(e):
                    torch.cuda.empty_cache()
                    print(f"  OOM at batch {batch_idx}, skipping")
                    continue
                raise

            loss.backward()

            if (batch_idx + 1) % args.grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(all_params, args.max_grad_norm)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                global_step += 1

            total_loss_val = losses["total"].item()
            epoch_losses.append(total_loss_val)

            # Logging
            if (batch_idx + 1) % args.log_every == 0:
                elapsed = time.time() - t0
                lr = scheduler.get_last_lr()[0]
                avg_loss = np.mean(epoch_losses[-args.log_every :])
                ctrl = losses.get("control", torch.tensor(0.0)).item()
                traj = losses.get("trajectory", torch.tensor(0.0)).item()
                seg = losses.get("segmentation", torch.tensor(0.0)).item()
                safety = losses.get("safety", torch.tensor(0.0)).item()

                print(
                    f"  [E{epoch+1}/{args.epochs}][{batch_idx+1}/{len(train_loader)}] "
                    f"loss={avg_loss:.4f} ctrl={ctrl:.4f} traj={traj:.4f} "
                    f"seg={seg:.4f} safety={safety:.4f} lr={lr:.2e} t={elapsed:.0f}s"
                )

                if HAS_TRACKIO:
                    trackio.log({
                        "train/loss": avg_loss,
                        "train/control_loss": ctrl,
                        "train/trajectory_loss": traj,
                        "train/segmentation_loss": seg,
                        "train/safety_loss": safety,
                        "train/lr": lr,
                        "train/epoch": epoch + batch_idx / len(train_loader),
                    })

            # Periodic eval
            if global_step > 0 and global_step % args.eval_every == 0:
                val_loss = evaluate(model, loss_fn, val_loader, device)
                print(f"  ── EVAL step {global_step}: val_loss={val_loss:.4f} (best={best_val_loss:.4f})")
                if HAS_TRACKIO:
                    trackio.log({"val/loss": val_loss, "val/step": global_step})
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    save_checkpoint(model, args.save_dir, "best")
                    print(f"  ── Saved best model (val_loss={val_loss:.4f})")
                model.train()

        # End-of-epoch eval
        val_loss = evaluate(model, loss_fn, val_loader, device)
        avg_epoch_loss = np.mean(epoch_losses) if epoch_losses else float("inf")
        print(
            f"\n  Epoch {epoch+1}/{args.epochs}: "
            f"train_loss={avg_epoch_loss:.4f}  val_loss={val_loss:.4f}"
        )
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            save_checkpoint(model, args.save_dir, "best")
            print(f"  ── New best model (val_loss={val_loss:.4f})")

    # ── Final save ──
    total_time = time.time() - t0
    print(f"\n{'='*60}")
    print(f"Training complete in {total_time/60:.1f} min")
    print(f"Best val loss: {best_val_loss:.4f}")
    save_checkpoint(model, args.save_dir, "final")

    # ── Push to Hub ──
    if args.push_to_hub:
        print(f"\n[Hub] Pushing model to {args.hub_model_id}...")
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            api.upload_folder(
                folder_path=os.path.join(args.save_dir, "best"),
                repo_id=args.hub_model_id,
                path_in_repo="trained_model",
                commit_message=f"Trained model (best val_loss={best_val_loss:.4f})",
            )
            print(f"  ✓ Pushed to {args.hub_model_id}/trained_model")
        except Exception as e:
            print(f"  Push failed: {e}")

    # ── Save metadata ──
    meta = {
        "dataset": DATASET_NAME,
        "train_samples": len(train_ds),
        "val_samples": len(val_ds),
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "grad_accum": args.grad_accum,
        "lr": args.lr,
        "best_val_loss": best_val_loss,
        "total_params": total_params,
        "training_time_min": total_time / 60,
        "device": str(device),
    }
    meta_path = os.path.join(args.save_dir, "training_meta.json")
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)
    print(f"  Metadata saved to {meta_path}")

    if args.push_to_hub:
        try:
            api.upload_file(
                path_or_fileobj=meta_path,
                path_in_repo="trained_model/training_meta.json",
                repo_id=args.hub_model_id,
            )
        except Exception:
            pass

    print("\nDone! ✓")
    return best_val_loss


def save_checkpoint(model, save_dir, tag):
    path = os.path.join(save_dir, tag)
    os.makedirs(path, exist_ok=True)
    if hasattr(model, "save_pretrained"):
        model.save_pretrained(path)
    else:
        torch.save(model.state_dict(), os.path.join(path, "model.pt"))


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  Main
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def parse_args():
    p = argparse.ArgumentParser(description="End-to-end FSD-Level5-CoT training on SADC")

    # Data
    p.add_argument("--train_samples", type=int, default=5000)
    p.add_argument("--val_samples", type=int, default=1000)
    p.add_argument("--train_split", type=str, default="pretrain_train")
    p.add_argument("--val_split", type=str, default="pretrain_val")
    p.add_argument("--data_dir", type=str, default="./sadc_subset")

    # Training
    p.add_argument("--epochs", type=int, default=5)
    p.add_argument("--batch_size", type=int, default=8)
    p.add_argument("--grad_accum", type=int, default=4)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--weight_decay", type=float, default=1e-4)
    p.add_argument("--max_grad_norm", type=float, default=5.0)
    p.add_argument("--num_workers", type=int, default=4)

    # Logging / eval
    p.add_argument("--log_every", type=int, default=10)
    p.add_argument("--eval_every", type=int, default=500)

    # Saving
    p.add_argument("--save_dir", type=str, default="./checkpoints")
    p.add_argument("--push_to_hub", action="store_true", default=True)
    p.add_argument("--no_push_to_hub", action="store_false", dest="push_to_hub")
    p.add_argument("--hub_model_id", type=str, default=HUB_MODEL_ID)

    return p.parse_args()


def main():
    args = parse_args()

    print("=" * 60)
    print("  FSD-Level5-CoT  ·  End-to-End Training on SADC")
    print("=" * 60)
    print(f"  Train samples: {args.train_samples}")
    print(f"  Val samples:   {args.val_samples}")
    print(f"  Epochs:        {args.epochs}")
    print(f"  Batch size:    {args.batch_size} × {args.grad_accum} accum = {args.batch_size * args.grad_accum}")
    print(f"  LR:            {args.lr}")
    print(f"  Push to Hub:   {args.push_to_hub} → {args.hub_model_id}")
    print("=" * 60)

    # Step 1: Download
    train_ds, val_ds = download_sadc_subset(
        train_samples=args.train_samples,
        val_samples=args.val_samples,
        output_dir=args.data_dir,
        train_split=args.train_split,
        val_split=args.val_split,
    )

    # Step 2+3: Train
    best_val = train(args, train_ds, val_ds)

    return best_val


if __name__ == "__main__":
    main()