"""
EL Defect Detection — Training Script for RTX 4060 (8GB VRAM)

Model: U-Net++ with EfficientNet-B4 encoder + scSE attention
Dataset: E-SCDD (snt-ubix/e-scdd) — 903 images, 512x512
Loss: 0.5 * Dice + 0.5 * Focal (handles severe class imbalance)
Classes: 0=background, 1=busbar, 2=crack, 3=dark/inactive, 4=other_defects

Usage:
    pip install torch torchvision segmentation-models-pytorch albumentations \
                huggingface-hub scikit-image scipy opencv-python-headless pillow
    python train.py
"""

import os
import sys
import json
import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from pathlib import Path
from PIL import Image

import segmentation_models_pytorch as smp
import albumentations as A
from albumentations.pytorch import ToTensorV2


# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════

class Config:
    # Data
    DATA_DIR = "./data"                  # Will download here
    OUTPUT_DIR = "./output"

    # Model — U-Net++ with EfficientNet-B4 is SOTA for thin-crack segmentation
    # Dense skip connections preserve fine details that plain U-Net misses
    ARCHITECTURE = "UnetPlusPlus"        # UnetPlusPlus > Unet for thin structures
    ENCODER = "efficientnet-b4"          # Best accuracy/size ratio, 20.9M params
    ENCODER_WEIGHTS = "imagenet"
    IN_CHANNELS = 1                      # EL images are grayscale
    NUM_CLASSES = 5                      # bg, busbar, crack, dark, other_defects

    # Training — tuned for RTX 4060 (8GB VRAM)
    IMG_SIZE = 512                       # E-SCDD native resolution
    BATCH_SIZE = 4                       # Safe for 8GB with AMP
    NUM_EPOCHS = 100
    ENCODER_LR = 1e-4                    # Lower LR for pretrained encoder
    DECODER_LR = 5e-4                    # Higher LR for random decoder
    WEIGHT_DECAY = 1e-4
    USE_AMP = True                       # Mixed precision — halves VRAM usage
    NUM_WORKERS = 4
    GRADIENT_CLIP = 1.0

    # Loss
    DICE_WEIGHT = 0.5
    FOCAL_WEIGHT = 0.5
    FOCAL_GAMMA = 2.0

    # Hub
    HUB_MODEL_ID = None                  # Set to "username/model-name" to push
    PUSH_TO_HUB = False

    # Class names
    CLASS_NAMES = ["background", "busbar", "crack", "dark", "other_defect"]


# ═══════════════════════════════════════════════════════════════
# CLASS MAPPING: E-SCDD 30 classes → 5 classes
# ═══════════════════════════════════════════════════════════════

# Mask pixel values in E-SCDD are integers 0-29 (Label column in CSV)
# We remap to 5 meaningful classes:
#   0 = background (all spacing, borders, padding, text, clamp, frame, jbox)
#   1 = busbar (label 9)
#   2 = crack (label 14=crack, label 10=crack_rbn_edge)
#   3 = dark/inactive (label 11=inactive, label 17=dead_cell, label 20=edge_dark)
#   4 = other_defect (rings, material, gridline, splice, corrosion, belt_mark, etc.)

LABEL_REMAP = np.zeros(30, dtype=np.uint8)  # default: everything → 0 (background)

# Background features (labels 0-8, 21-24, 29)
# Already 0 by default

# Busbar
LABEL_REMAP[9] = 1    # busbars → busbar

# Crack (HIGH IMPORTANCE)
LABEL_REMAP[10] = 2   # crack_rbn_edge → crack
LABEL_REMAP[14] = 2   # crack → crack

# Dark/Inactive (HIGH IMPORTANCE)
LABEL_REMAP[11] = 3   # inactive → dark
LABEL_REMAP[17] = 3   # dead_cell → dark
LABEL_REMAP[20] = 3   # edge_dark → dark

# Other defects (MEDIUM IMPORTANCE)
LABEL_REMAP[12] = 4   # rings
LABEL_REMAP[13] = 4   # material
LABEL_REMAP[15] = 4   # gridline defect
LABEL_REMAP[16] = 4   # splice
LABEL_REMAP[18] = 4   # corrosion_rbn
LABEL_REMAP[19] = 4   # belt_mark
LABEL_REMAP[25] = 4   # scuff
LABEL_REMAP[26] = 4   # corrosion_cell
LABEL_REMAP[27] = 4   # brightening
LABEL_REMAP[28] = 4   # star


# ═══════════════════════════════════════════════════════════════
# DATASET
# ═══════════════════════════════════════════════════════════════

class ESCDDDataset(Dataset):
    """
    E-SCDD dataset: 512x512 EL images (RGBA) + grayscale masks (L, values 0-29).
    """

    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_dir = Path(img_dir)
        self.mask_dir = Path(mask_dir)
        self.transform = transform

        # Match images to masks by filename
        img_files = {f.stem: f for f in sorted(self.img_dir.glob("*.png"))}
        mask_files = {f.stem: f for f in sorted(self.mask_dir.glob("*.png"))}

        self.pairs = []
        for stem in img_files:
            if stem in mask_files:
                self.pairs.append((img_files[stem], mask_files[stem]))

        print(f"  {img_dir}: {len(self.pairs)} image-mask pairs")

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

    def __getitem__(self, idx):
        img_path, mask_path = self.pairs[idx]

        # Load image — RGBA, convert to grayscale
        img = np.array(Image.open(img_path).convert("L"), dtype=np.float32)

        # Load mask — grayscale, pixel value = class label (0-29)
        mask = np.array(Image.open(mask_path), dtype=np.uint8)

        # Remap 30 → 5 classes using lookup table
        mask = LABEL_REMAP[np.clip(mask, 0, 29)]

        # Apply augmentations
        if self.transform:
            augmented = self.transform(image=img, mask=mask)
            img = augmented["image"]    # (1, H, W) float tensor
            mask = augmented["mask"]    # (H, W) long tensor
        else:
            img = torch.from_numpy(img).unsqueeze(0) / 255.0
            mask = torch.from_numpy(mask).long()

        return img, mask


def get_train_transforms(img_size=512):
    return A.Compose([
        A.RandomCrop(img_size, img_size, p=1.0),
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.5),
        A.RandomRotate90(p=0.5),
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
        A.GaussNoise(std_range=(0.02, 0.1), p=0.3),
        A.GridDistortion(num_steps=5, distort_limit=0.3, p=0.3),
        A.Normalize(mean=[0.0], std=[1.0], max_pixel_value=255.0),
        ToTensorV2(),
    ])


def get_val_transforms(img_size=512):
    return A.Compose([
        A.CenterCrop(img_size, img_size, p=1.0),
        A.Normalize(mean=[0.0], std=[1.0], max_pixel_value=255.0),
        ToTensorV2(),
    ])


# ═══════════════════════════════════════════════════════════════
# DOWNLOAD DATASET
# ═══════════════════════════════════════════════════════════════

def download_dataset(data_dir):
    """Download E-SCDD from HuggingFace Hub."""
    train_img = os.path.join(data_dir, "el_images_train")
    if os.path.exists(train_img) and len(os.listdir(train_img)) > 100:
        print("Dataset already downloaded.")
        return

    print("Downloading E-SCDD dataset from HuggingFace Hub...")
    from huggingface_hub import snapshot_download
    snapshot_download(
        repo_id="snt-ubix/e-scdd",
        repo_type="dataset",
        local_dir=data_dir,
    )
    print(f"Downloaded to {data_dir}")


# ═══════════════════════════════════════════════════════════════
# METRICS
# ═══════════════════════════════════════════════════════════════

def compute_metrics(pred_logits, target, num_classes=5):
    """Compute per-class IoU and Dice."""
    pred = torch.argmax(pred_logits, dim=1)  # (B, H, W)

    ious, dices = [], []
    for c in range(num_classes):
        pred_c = (pred == c)
        target_c = (target == c)

        intersection = (pred_c & target_c).float().sum()
        union = (pred_c | target_c).float().sum()

        iou = (intersection + 1e-6) / (union + 1e-6)
        dice = (2 * intersection + 1e-6) / (pred_c.float().sum() + target_c.float().sum() + 1e-6)

        ious.append(iou.item())
        dices.append(dice.item())

    return {
        "mean_iou": np.mean(ious),
        "mean_dice": np.mean(dices),
        "per_class_iou": dict(zip(Config.CLASS_NAMES, ious)),
        "per_class_dice": dict(zip(Config.CLASS_NAMES, dices)),
    }


# ═══════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════

def train():
    cfg = Config()
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")

    # ── Download data ────────────────────────────────────────
    download_dataset(cfg.DATA_DIR)

    # ── Create datasets ──────────────────────────────────────
    print("\nLoading datasets...")
    train_ds = ESCDDDataset(
        os.path.join(cfg.DATA_DIR, "el_images_train"),
        os.path.join(cfg.DATA_DIR, "el_masks_train"),
        transform=get_train_transforms(cfg.IMG_SIZE),
    )
    val_ds = ESCDDDataset(
        os.path.join(cfg.DATA_DIR, "el_images_val"),
        os.path.join(cfg.DATA_DIR, "el_masks_val"),
        transform=get_val_transforms(cfg.IMG_SIZE),
    )

    train_loader = DataLoader(
        train_ds, batch_size=cfg.BATCH_SIZE, shuffle=True,
        num_workers=cfg.NUM_WORKERS, pin_memory=True, drop_last=True,
    )
    val_loader = DataLoader(
        val_ds, batch_size=cfg.BATCH_SIZE, shuffle=False,
        num_workers=cfg.NUM_WORKERS, pin_memory=True,
    )

    # ── Compute class weights from training data ─────────────
    print("\nComputing class distribution...")
    class_pixels = np.zeros(cfg.NUM_CLASSES, dtype=np.float64)
    for i in range(min(len(train_ds), 200)):  # Sample 200 images
        _, mask = train_ds[i]
        if isinstance(mask, torch.Tensor):
            mask = mask.numpy()
        for c in range(cfg.NUM_CLASSES):
            class_pixels[c] += (mask == c).sum()

    total = class_pixels.sum()
    class_freq = class_pixels / total
    print("Class distribution:")
    for i, name in enumerate(cfg.CLASS_NAMES):
        print(f"  {name}: {class_freq[i]*100:.2f}%  ({int(class_pixels[i]):,} px)")

    # ── Create model ─────────────────────────────────────────
    print(f"\nCreating {cfg.ARCHITECTURE} + {cfg.ENCODER}...")
    ModelClass = getattr(smp, cfg.ARCHITECTURE)
    model = ModelClass(
        encoder_name=cfg.ENCODER,
        encoder_weights=cfg.ENCODER_WEIGHTS,
        in_channels=cfg.IN_CHANNELS,
        classes=cfg.NUM_CLASSES,
        decoder_attention_type="scse",
    )
    model = model.to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")

    # ── Loss: Dice + Focal (handles class imbalance) ─────────
    dice_loss = smp.losses.DiceLoss(mode="multiclass", from_logits=True, smooth=1.0)
    focal_loss = smp.losses.FocalLoss(mode="multiclass", gamma=cfg.FOCAL_GAMMA)

    def criterion(pred, target):
        return cfg.DICE_WEIGHT * dice_loss(pred, target) + cfg.FOCAL_WEIGHT * focal_loss(pred, target)

    # ── Optimizer with differential LR ───────────────────────
    optimizer = AdamW([
        {"params": model.encoder.parameters(), "lr": cfg.ENCODER_LR},
        {"params": model.decoder.parameters(), "lr": cfg.DECODER_LR},
        {"params": model.segmentation_head.parameters(), "lr": cfg.DECODER_LR},
    ], weight_decay=cfg.WEIGHT_DECAY)

    scheduler = CosineAnnealingLR(optimizer, T_max=cfg.NUM_EPOCHS, eta_min=1e-6)
    scaler = torch.amp.GradScaler(enabled=cfg.USE_AMP)

    # ── Training loop ────────────────────────────────────────
    best_val_dice = 0.0
    history = {"train_loss": [], "val_loss": [], "val_dice": [], "val_iou": []}

    print(f"\n{'='*60}")
    print(f"Starting training: {cfg.NUM_EPOCHS} epochs")
    print(f"{'='*60}\n")

    for epoch in range(cfg.NUM_EPOCHS):
        t_start = time.time()

        # ── Train ────────────────────────────────────────────
        model.train()
        train_loss = 0.0

        for batch_idx, (images, masks) in enumerate(train_loader):
            images = images.to(device)
            masks = masks.to(device)

            optimizer.zero_grad()

            with torch.amp.autocast(device_type="cuda", enabled=cfg.USE_AMP):
                logits = model(images)
                loss = criterion(logits, masks)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.GRADIENT_CLIP)
            scaler.step(optimizer)
            scaler.update()

            train_loss += loss.item()

        train_loss /= len(train_loader)
        scheduler.step()

        # ── Validate ─────────────────────────────────────────
        model.eval()
        val_loss = 0.0
        all_ious, all_dices = [], []

        with torch.no_grad():
            for images, masks in val_loader:
                images = images.to(device)
                masks = masks.to(device)

                with torch.amp.autocast(device_type="cuda", enabled=cfg.USE_AMP):
                    logits = model(images)
                    loss = criterion(logits, masks)

                val_loss += loss.item()
                metrics = compute_metrics(logits, masks, cfg.NUM_CLASSES)
                all_ious.append(metrics["mean_iou"])
                all_dices.append(metrics["mean_dice"])

        val_loss /= len(val_loader)
        val_dice = np.mean(all_dices)
        val_iou = np.mean(all_ious)

        t_elapsed = time.time() - t_start
        lr_enc = optimizer.param_groups[0]["lr"]
        lr_dec = optimizer.param_groups[1]["lr"]

        print(f"Epoch {epoch+1:3d}/{cfg.NUM_EPOCHS} | "
              f"train_loss={train_loss:.4f} | val_loss={val_loss:.4f} | "
              f"val_dice={val_dice:.4f} | val_iou={val_iou:.4f} | "
              f"lr_enc={lr_enc:.6f} | {t_elapsed:.1f}s")

        # Per-class dice every 10 epochs
        if (epoch + 1) % 10 == 0:
            # Run full validation for per-class metrics
            all_per_class = {name: [] for name in cfg.CLASS_NAMES}
            with torch.no_grad():
                for images, masks in val_loader:
                    images, masks = images.to(device), masks.to(device)
                    with torch.amp.autocast(device_type="cuda", enabled=cfg.USE_AMP):
                        logits = model(images)
                    m = compute_metrics(logits, masks, cfg.NUM_CLASSES)
                    for name in cfg.CLASS_NAMES:
                        all_per_class[name].append(m["per_class_dice"][name])
            print("  Per-class Dice:")
            for name in cfg.CLASS_NAMES:
                print(f"    {name:20s}: {np.mean(all_per_class[name]):.4f}")

        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["val_dice"].append(val_dice)
        history["val_iou"].append(val_iou)

        # ── Save best model ──────────────────────────────────
        if val_dice > best_val_dice:
            best_val_dice = val_dice
            save_path = os.path.join(cfg.OUTPUT_DIR, "best_model.pth")
            torch.save({
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_dice": val_dice,
                "val_iou": val_iou,
                "architecture": cfg.ARCHITECTURE,
                "encoder": cfg.ENCODER,
                "num_classes": cfg.NUM_CLASSES,
                "img_size": cfg.IMG_SIZE,
                "class_names": cfg.CLASS_NAMES,
                "label_remap": LABEL_REMAP.tolist(),
            }, save_path)
            print(f"  → Best model saved (dice={val_dice:.4f})")

        # Periodic checkpoint every 25 epochs
        if (epoch + 1) % 25 == 0:
            ckpt_path = os.path.join(cfg.OUTPUT_DIR, f"checkpoint_ep{epoch+1}.pth")
            torch.save({"epoch": epoch+1, "model_state_dict": model.state_dict()}, ckpt_path)

    # ── Save final model + history ───────────────────────────
    final_path = os.path.join(cfg.OUTPUT_DIR, "final_model.pth")
    torch.save({
        "epoch": cfg.NUM_EPOCHS,
        "model_state_dict": model.state_dict(),
        "val_dice": history["val_dice"][-1],
        "val_iou": history["val_iou"][-1],
        "architecture": cfg.ARCHITECTURE,
        "encoder": cfg.ENCODER,
        "num_classes": cfg.NUM_CLASSES,
        "img_size": cfg.IMG_SIZE,
        "class_names": cfg.CLASS_NAMES,
        "label_remap": LABEL_REMAP.tolist(),
        "history": history,
    }, final_path)

    with open(os.path.join(cfg.OUTPUT_DIR, "history.json"), "w") as f:
        json.dump(history, f, indent=2)

    print(f"\n{'='*60}")
    print(f"Training complete! Best val dice: {best_val_dice:.4f}")
    print(f"Models saved to {cfg.OUTPUT_DIR}/")
    print(f"{'='*60}")

    # ── Push to Hub ──────────────────────────────────────────
    if cfg.PUSH_TO_HUB and cfg.HUB_MODEL_ID:
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            api.create_repo(cfg.HUB_MODEL_ID, exist_ok=True)
            api.upload_folder(
                folder_path=cfg.OUTPUT_DIR,
                repo_id=cfg.HUB_MODEL_ID,
                commit_message=f"Trained model (dice={best_val_dice:.4f})",
            )
            print(f"Pushed to hub: {cfg.HUB_MODEL_ID}")
        except Exception as e:
            print(f"Hub push failed: {e}")


if __name__ == "__main__":
    train()