#!/usr/bin/env python3
"""
fine_tune.py — Fine-tune BiSeNetV2 on custom sidewalk data.

Walk-and-label workflow:
  1. Record a ROS2 bag while walking the route:
       ros2 bag record /camera/color/image_raw -o sidewalk_route
  2. Extract frames from the bag:
       python3 fine_tune.py --extract-frames sidewalk_route/ --output-dir data/raw/ --every-n 5
  3. Label 50+ frames in LabelMe (or CVAT):
       labelme data/raw/  (save as JSON)
  4. Convert labels to Cityscapes-format masks:
       python3 fine_tune.py --convert-labels data/raw/ --output-dir data/labels/
  5. Fine-tune:
       python3 fine_tune.py --train --images data/raw/ --labels data/labels/ \
           --weights /mnt/nvme/saltybot/models/bisenetv2_cityscapes.pth \
           --output /mnt/nvme/saltybot/models/bisenetv2_custom.pth
  6. Build new TRT engine:
       python3 build_engine.py --weights /mnt/nvme/saltybot/models/bisenetv2_custom.pth

LabelMe class names (must match exactly in labelme JSON):
  sidewalk  → Cityscapes ID 1 → TRAV_SIDEWALK
  road      → Cityscapes ID 0 → TRAV_ROAD
  grass     → Cityscapes ID 8 → TRAV_GRASS (use 'vegetation')
  obstacle  → Cityscapes ID 2 → TRAV_OBSTACLE (use 'building' for static)
  person    → Cityscapes ID 11
  car       → Cityscapes ID 13

Requirements:
  pip install torch torchvision albumentations labelme opencv-python
"""

import argparse
import json
import os
import sys

import numpy as np

# ── Label → Cityscapes ID mapping for custom annotation ──────────────────────

LABEL_TO_CITYSCAPES = {
    "road":        0,
    "sidewalk":    1,
    "building":    2,
    "wall":        3,
    "fence":       4,
    "pole":        5,
    "traffic light": 6,
    "traffic sign":  7,
    "vegetation":  8,
    "grass":       8,   # alias
    "terrain":     9,
    "sky":         10,
    "person":      11,
    "rider":       12,
    "car":         13,
    "truck":       14,
    "bus":         15,
    "train":       16,
    "motorcycle":  17,
    "bicycle":     18,
    # saltybot-specific aliases
    "kerb":        1,   # map kerb → sidewalk
    "pavement":    1,   # British English
    "path":        1,
    "tarmac":      0,   # map tarmac → road
    "shrub":       8,   # map shrub → vegetation
}

DEFAULT_CLASS = 255  # unlabelled pixels → ignore in loss


def parse_args():
    p = argparse.ArgumentParser(description="Fine-tune BiSeNetV2 on custom data")
    p.add_argument("--extract-frames", metavar="BAG_DIR",
                   help="Extract frames from a ROS2 bag")
    p.add_argument("--output-dir",     default="data/raw",
                   help="Output directory for extracted frames or converted labels")
    p.add_argument("--every-n",        type=int, default=5,
                   help="Extract every N-th frame from bag")
    p.add_argument("--convert-labels", metavar="LABELME_DIR",
                   help="Convert LabelMe JSON annotations to Cityscapes masks")
    p.add_argument("--train",          action="store_true",
                   help="Run fine-tuning loop")
    p.add_argument("--images",         default="data/raw",
                   help="Directory of training images")
    p.add_argument("--labels",         default="data/labels",
                   help="Directory of label masks (PNG, Cityscapes IDs)")
    p.add_argument("--weights",        required=False,
                   default="/mnt/nvme/saltybot/models/bisenetv2_cityscapes.pth",
                   help="Path to pretrained BiSeNetV2 weights")
    p.add_argument("--output",
                   default="/mnt/nvme/saltybot/models/bisenetv2_custom.pth",
                   help="Output path for fine-tuned weights")
    p.add_argument("--epochs",         type=int, default=20)
    p.add_argument("--lr",             type=float, default=1e-4)
    p.add_argument("--batch-size",     type=int, default=4)
    p.add_argument("--input-size",     nargs=2, type=int, default=[512, 256],
                   help="Model input size: W H")
    p.add_argument("--eval",           action="store_true",
                   help="Evaluate mIoU on labelled data instead of training")
    return p.parse_args()


# ── Frame extraction from ROS2 bag ────────────────────────────────────────────

def extract_frames(bag_dir: str, output_dir: str, every_n: int = 5):
    """Extract /camera/color/image_raw frames from a ROS2 bag."""
    try:
        import rclpy
        from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions
        from rclpy.serialization import deserialize_message
        from sensor_msgs.msg import Image
        from cv_bridge import CvBridge
        import cv2
    except ImportError:
        print("ERROR: rosbag2_py / rclpy not available. Must run in ROS2 environment.")
        sys.exit(1)

    os.makedirs(output_dir, exist_ok=True)
    bridge = CvBridge()

    reader = SequentialReader()
    reader.open(
        StorageOptions(uri=bag_dir, storage_id="sqlite3"),
        ConverterOptions("", ""),
    )

    topic = "/camera/color/image_raw"
    count = 0
    saved = 0

    while reader.has_next():
        topic_name, data, ts = reader.read_next()
        if topic_name != topic:
            continue
        count += 1
        if count % every_n != 0:
            continue

        msg = deserialize_message(data, Image)
        bgr = bridge.imgmsg_to_cv2(msg, "bgr8")
        fname = os.path.join(output_dir, f"frame_{ts:020d}.jpg")
        cv2.imwrite(fname, bgr, [cv2.IMWRITE_JPEG_QUALITY, 95])
        saved += 1
        if saved % 10 == 0:
            print(f"  Saved {saved} frames...")

    print(f"Extracted {saved} frames from {count} total to {output_dir}")


# ── LabelMe JSON → Cityscapes mask ───────────────────────────────────────────

def convert_labelme_to_masks(labelme_dir: str, output_dir: str):
    """Convert LabelMe polygon annotations to per-pixel Cityscapes-ID PNG masks."""
    try:
        import cv2
    except ImportError:
        print("ERROR: opencv-python not installed. pip install opencv-python")
        sys.exit(1)

    os.makedirs(output_dir, exist_ok=True)
    json_files = [f for f in os.listdir(labelme_dir) if f.endswith(".json")]

    if not json_files:
        print(f"No .json annotation files found in {labelme_dir}")
        sys.exit(1)

    print(f"Converting {len(json_files)} annotation files...")
    for jf in json_files:
        with open(os.path.join(labelme_dir, jf)) as f:
            anno = json.load(f)

        h = anno["imageHeight"]
        w = anno["imageWidth"]
        mask = np.full((h, w), DEFAULT_CLASS, dtype=np.uint8)

        for shape in anno.get("shapes", []):
            label = shape["label"].lower().strip()
            cid = LABEL_TO_CITYSCAPES.get(label)
            if cid is None:
                print(f"  WARNING: unknown label '{label}' in {jf} — skipping")
                continue

            pts = np.array(shape["points"], dtype=np.int32)
            cv2.fillPoly(mask, [pts], cid)

        out_name = os.path.splitext(jf)[0] + "_mask.png"
        cv2.imwrite(os.path.join(output_dir, out_name), mask)

    print(f"Masks saved to {output_dir}")


# ── Training loop ─────────────────────────────────────────────────────────────

def train(args):
    """Fine-tune BiSeNetV2 with cross-entropy loss + ignore_index=255."""
    try:
        import torch
        import torch.nn as nn
        import torch.optim as optim
        from torch.utils.data import Dataset, DataLoader
        import cv2
        import albumentations as A
        from albumentations.pytorch import ToTensorV2
    except ImportError as e:
        print(f"ERROR: missing dependency — {e}")
        print("pip install torch torchvision albumentations opencv-python")
        sys.exit(1)

    # -- Dataset ---------------------------------------------------------------
    class SegDataset(Dataset):
        def __init__(self, img_dir, lbl_dir, transform=None):
            self.imgs = sorted([
                os.path.join(img_dir, f)
                for f in os.listdir(img_dir)
                if f.endswith((".jpg", ".png"))
            ])
            self.lbls = sorted([
                os.path.join(lbl_dir, f)
                for f in os.listdir(lbl_dir)
                if f.endswith(".png")
            ])
            if len(self.imgs) != len(self.lbls):
                raise ValueError(
                    f"Image/label count mismatch: {len(self.imgs)} vs {len(self.lbls)}"
                )
            self.transform = transform

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

        def __getitem__(self, idx):
            img = cv2.cvtColor(cv2.imread(self.imgs[idx]), cv2.COLOR_BGR2RGB)
            lbl = cv2.imread(self.lbls[idx], cv2.IMREAD_GRAYSCALE)
            if self.transform:
                aug = self.transform(image=img, mask=lbl)
                img, lbl = aug["image"], aug["mask"]
            return img, lbl.long()

    iw, ih = args.input_size
    transform = A.Compose([
        A.Resize(ih, iw),
        A.HorizontalFlip(p=0.5),
        A.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1, p=0.5),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2(),
    ])

    dataset = SegDataset(args.images, args.labels, transform=transform)
    loader  = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                         num_workers=2, pin_memory=True)
    print(f"Dataset: {len(dataset)} samples  batch={args.batch_size}")

    # -- Model -----------------------------------------------------------------
    sys.path.insert(0, "/tmp/BiSeNet")
    try:
        from lib.models.bisenetv2 import BiSeNetV2
    except ImportError:
        os.system("git clone --depth 1 https://github.com/CoinCheung/BiSeNet.git /tmp/BiSeNet")
        from lib.models.bisenetv2 import BiSeNetV2

    net = BiSeNetV2(n_classes=19)
    state = torch.load(args.weights, map_location="cpu")
    state_dict = state.get("state_dict", state.get("model", state))
    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
    net.load_state_dict(state_dict, strict=False)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    net = net.to(device)
    print(f"Training on: {device}")

    # -- Optimiser (low LR to preserve Cityscapes knowledge) -------------------
    optimiser = optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=args.epochs)
    criterion = nn.CrossEntropyLoss(ignore_index=DEFAULT_CLASS)

    # -- Loop ------------------------------------------------------------------
    best_loss = float("inf")
    for epoch in range(1, args.epochs + 1):
        net.train()
        total_loss = 0.0

        for imgs, lbls in loader:
            imgs = imgs.to(device, dtype=torch.float32)
            lbls = lbls.to(device)

            out = net(imgs)
            logits = out[0] if isinstance(out, (list, tuple)) else out

            loss = criterion(logits, lbls)
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
            total_loss += loss.item()

        scheduler.step()
        avg = total_loss / len(loader)
        print(f"  Epoch {epoch:3d}/{args.epochs}  loss={avg:.4f}  lr={scheduler.get_last_lr()[0]:.2e}")

        if avg < best_loss:
            best_loss = avg
            torch.save(net.state_dict(), args.output)
            print(f"  Saved best model → {args.output}")

    print(f"\nFine-tuning done. Best loss: {best_loss:.4f}  Saved: {args.output}")
    print("\nNext step: rebuild TRT engine:")
    print(f"  python3 build_engine.py --weights {args.output}")


# ── mIoU evaluation ───────────────────────────────────────────────────────────

def evaluate_miou(args):
    """Compute mIoU on labelled validation data."""
    try:
        import torch
        import cv2
    except ImportError:
        print("ERROR: torch/cv2 not available")
        sys.exit(1)

    from saltybot_segmentation.seg_utils import cityscapes_to_traversability

    NUM_CLASSES = 19
    confusion = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int64)

    sys.path.insert(0, "/tmp/BiSeNet")
    from lib.models.bisenetv2 import BiSeNetV2
    import torch

    net = BiSeNetV2(n_classes=NUM_CLASSES)
    state_dict = {k.replace("module.", ""): v
                  for k, v in torch.load(args.weights, map_location="cpu").items()}
    net.load_state_dict(state_dict, strict=False)
    net.eval().cuda() if torch.cuda.is_available() else net.eval()
    device = next(net.parameters()).device

    iw, ih = args.input_size
    img_files = sorted([f for f in os.listdir(args.images) if f.endswith((".jpg", ".png"))])
    lbl_files = sorted([f for f in os.listdir(args.labels) if f.endswith(".png")])

    from saltybot_segmentation.seg_utils import preprocess_for_inference, unpad_mask

    with torch.no_grad():
        for imgf, lblf in zip(img_files, lbl_files):
            img = cv2.imread(os.path.join(args.images, imgf))
            lbl = cv2.imread(os.path.join(args.labels, lblf), cv2.IMREAD_GRAYSCALE)
            H, W = img.shape[:2]

            blob, scale, pad_l, pad_t = preprocess_for_inference(img, iw, ih)
            t_blob = torch.from_numpy(blob).to(device)
            out = net(t_blob)
            logits = out[0] if isinstance(out, (list, tuple)) else out
            pred = logits[0].argmax(0).cpu().numpy().astype(np.uint8)
            pred = unpad_mask(pred, H, W, scale, pad_l, pad_t)

            valid = lbl != DEFAULT_CLASS
            np.add.at(confusion, (lbl[valid].ravel(), pred[valid].ravel()), 1)

    # Per-class IoU
    iou_per_class = []
    for c in range(NUM_CLASSES):
        tp = confusion[c, c]
        fp = confusion[:, c].sum() - tp
        fn = confusion[c, :].sum() - tp
        denom = tp + fp + fn
        iou_per_class.append(tp / denom if denom > 0 else float("nan"))

    valid_iou = [v for v in iou_per_class if not np.isnan(v)]
    miou = np.mean(valid_iou)

    CITYSCAPES_NAMES = ["road", "sidewalk", "building", "wall", "fence", "pole",
                        "traffic light", "traffic sign", "vegetation", "terrain",
                        "sky", "person", "rider", "car", "truck", "bus", "train",
                        "motorcycle", "bicycle"]
    print("\nmIoU per class:")
    for i, (name, iou) in enumerate(zip(CITYSCAPES_NAMES, iou_per_class)):
        marker = " *" if not np.isnan(iou) else ""
        print(f"  {i:2d} {name:<16} {iou*100:5.1f}%{marker}")
    print(f"\nmIoU: {miou*100:.2f}%  ({len(valid_iou)}/{NUM_CLASSES} classes present)")


def main():
    args = parse_args()

    if args.extract_frames:
        extract_frames(args.extract_frames, args.output_dir, args.every_n)
    elif args.convert_labels:
        convert_labelme_to_masks(args.convert_labels, args.output_dir)
    elif args.eval:
        evaluate_miou(args)
    elif args.train:
        train(args)
    else:
        print("Specify one of: --extract-frames, --convert-labels, --train, --eval")
        sys.exit(1)


if __name__ == "__main__":
    main()
