#!/home/joaquin.m.belart/software/miniconda3/bin/python3.13
"""
geo_sr.py — Photogrammetric Super-Resolution CLI  [v2 — 2026 upgrade]

Trains and applies super-resolution models on aerial/orthophoto imagery while
preserving GeoTIFF georeferencing metadata.

Supports both RGB (3-channel) and panchromatic/single-channel workflows.
The mode is auto-detected from the HR training images.

ARCHITECTURES (--arch):
  rrdb  (default) : RRDB generator (Real-ESRGAN backbone) + U-Net discriminator
                    Wang et al., ICCVW 2021 [arXiv:2107.10833]
  hat             : Hybrid Attention Transformer generator
                    Chen et al., CVPR 2023 / TPAMI 2023 [arXiv:2309.05239]
                    +0.3–1.2 dB PSNR over SwinIR; wider receptive field via
                    window attention + channel attention in every block.

PREPROCESSING (train-denoiser / preprocess subcommands):
  Restormer-style U-Net denoiser for film grain, scanner noise, and motion blur
  removal BEFORE super-resolution.
  Zamir et al., CVPR 2022 [arXiv:2111.09881]

HISTORICAL DEGRADATION (--historical):
  Adds film-grain noise (signal-dependent), vignetting, scanner banding, and
  film-base fog to the synthetic LR pipeline — better domain match for scanned
  aerial photographs from the 1940s–1990s.

INSTALL:
  pip install torch torchvision opencv-python-headless pillow tqdm numpy
  pip install rasterio   # optional, for GeoTIFF georeferencing support

USAGE (SR — identical interface to v1):
  python geo_sr.py train --hr-dir ./hr_5cm --scale 2 --model ./models/sr.pth
  python geo_sr.py train --hr-dir ./hr_pan  --scale 2 --arch hat \\
      --model ./models/sr_hat.pth --historical
  python geo_sr.py apply --input-dir ./25cm --output-dir ./sr_x2 \\
      --model ./models/sr.pth

USAGE (denoiser):
  python geo_sr.py train-denoiser --hr-dir ./hr_5cm \\
      --model ./models/denoiser.pth --historical
  python geo_sr.py preprocess --input-dir ./scans --output-dir ./clean \\
      --model ./models/denoiser.pth
  # Combined: denoise then super-resolve
  python geo_sr.py apply --input-dir ./25cm --output-dir ./sr_x2 \\
      --model ./models/sr.pth --preprocess-model ./models/denoiser.pth
"""

import argparse
import math
import os
import random
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import List, Optional, Tuple

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader, Dataset
from torchvision import models
from tqdm import tqdm

# ─────────────────────────────────────────────────────────────────────────────
# Optional GeoTIFF support
# ─────────────────────────────────────────────────────────────────────────────
try:
    import rasterio
    HAS_RASTERIO = True
except ImportError:
    HAS_RASTERIO = False


# ─────────────────────────────────────────────────────────────────────────────
# RRDB Architecture  (Real-ESRGAN backbone — unchanged from v1)
# Wang et al., ICCVW 2021  [arXiv:2107.10833]
# ─────────────────────────────────────────────────────────────────────────────

class ResidualDenseBlock(nn.Module):
    """5-layer Residual Dense Block with local residual learning."""

    def __init__(self, num_feat: int = 64, num_grow_ch: int = 32):
        super().__init__()
        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, a=0.2)
                m.weight.data *= 0.1

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat([x, x1], dim=1)))
        x3 = self.lrelu(self.conv3(torch.cat([x, x1, x2], dim=1)))
        x4 = self.lrelu(self.conv4(torch.cat([x, x1, x2, x3], dim=1)))
        x5 = self.conv5(torch.cat([x, x1, x2, x3, x4], dim=1))
        return x5 * 0.2 + x


class RRDB(nn.Module):
    def __init__(self, num_feat: int = 64, num_grow_ch: int = 32):
        super().__init__()
        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)

    def forward(self, x):
        out = self.rdb1(x)
        out = self.rdb2(out)
        out = self.rdb3(out)
        return out * 0.2 + x


def _make_upsampler(num_feat: int, scale: int) -> nn.Sequential:
    """Pixel-shuffle upsampler for scale = 2, 3, 4."""
    layers: List[nn.Module] = []
    if scale == 4:
        for _ in range(2):
            layers += [
                nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1),
                nn.PixelShuffle(2),
                nn.LeakyReLU(0.2, inplace=True),
            ]
    elif scale in (2, 3):
        layers += [
            nn.Conv2d(num_feat, num_feat * scale * scale, 3, 1, 1),
            nn.PixelShuffle(scale),
            nn.LeakyReLU(0.2, inplace=True),
        ]
    else:
        raise ValueError(f"Unsupported scale: {scale}. Choose 2, 3 or 4.")
    return nn.Sequential(*layers)


class RRDBNet(nn.Module):
    """RRDB generator (Real-ESRGAN style). Supports scale 2/3/4, channels 1 or 3."""

    def __init__(self, in_channels=3, out_channels=3, num_feat=64,
                 num_block=23, num_grow_ch=32, scale=2):
        super().__init__()
        self.scale = scale
        self.conv_first = nn.Conv2d(in_channels, num_feat, 3, 1, 1)
        self.body = nn.Sequential(*[RRDB(num_feat, num_grow_ch) for _ in range(num_block)])
        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.upsample = _make_upsampler(num_feat, scale)
        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_last = nn.Conv2d(num_feat, out_channels, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        feat = self.conv_first(x)
        feat = feat + self.conv_body(self.body(feat))
        feat = self.upsample(feat)
        return self.conv_last(self.lrelu(self.conv_hr(feat)))


# ─────────────────────────────────────────────────────────────────────────────
# HAT Architecture  (Hybrid Attention Transformer)
# Chen et al., CVPR 2023 / TPAMI 2023  [arXiv:2309.05239]
#
# Key innovations over SwinIR:
#   1. Channel Attention (SE) inside every transformer block — recalibrates
#      inter-channel dependencies that window attention ignores.
#   2. Same-task pre-training (HAT-L) — not implemented here, but the
#      architecture is identical and can load those weights.
#   3. Wider effective receptive field from combining W-MSA + SW-MSA + CA.
# ─────────────────────────────────────────────────────────────────────────────

def _win_partition(x: torch.Tensor, ws: int) -> torch.Tensor:
    """(B, H, W, C)  →  (B*nW, ws, ws, C)"""
    B, H, W, C = x.shape
    x = x.view(B, H // ws, ws, W // ws, ws, C)
    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, ws, ws, C)


def _win_reverse(windows: torch.Tensor, ws: int, H: int, W: int) -> torch.Tensor:
    """(B*nW, ws, ws, C)  →  (B, H, W, C)"""
    B = int(windows.shape[0] / (H // ws) / (W // ws))
    x = windows.view(B, H // ws, W // ws, ws, ws, -1)
    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)


def _sw_attn_mask(H: int, W: int, ws: int, shift: int,
                  device: torch.device) -> Optional[torch.Tensor]:
    """Attention mask for shifted-window MSA. Returns (nW, ws², ws²) or None."""
    if shift == 0:
        return None
    img = torch.zeros(1, H, W, 1, device=device)
    slices_h = [slice(0, -ws), slice(-ws, -shift), slice(-shift, None)]
    slices_w = [slice(0, -ws), slice(-ws, -shift), slice(-shift, None)]
    cnt = 0
    for sh in slices_h:
        for sw in slices_w:
            img[:, sh, sw, :] = cnt
            cnt += 1
    wins = _win_partition(img, ws).view(-1, ws * ws)   # (nW, ws²)
    mask = wins.unsqueeze(1) - wins.unsqueeze(2)        # (nW, ws², ws²)
    return mask.masked_fill(mask != 0, -100.0).masked_fill(mask == 0, 0.0)


class _WindowAttention(nn.Module):
    """
    Window MSA with relative position bias.
    Supports regular (shift=0) and shifted (shift=ws//2) windows.
    """

    def __init__(self, dim: int, window_size: int, num_heads: int,
                 qkv_bias: bool = True):
        super().__init__()
        self.ws = window_size
        self.num_heads = num_heads
        self.scale = (dim // num_heads) ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)
        # Relative position bias table: (2*ws-1)² entries, num_heads values
        self.rpb = nn.Parameter(torch.zeros((2 * window_size - 1) ** 2, num_heads))
        nn.init.trunc_normal_(self.rpb, std=0.02)
        # Precompute flat index into rpb for each query-key pair
        c = torch.stack(torch.meshgrid(
            torch.arange(window_size), torch.arange(window_size), indexing='ij'
        )).flatten(1)                           # (2, ws²)
        rel = c[:, :, None] - c[:, None, :]    # (2, ws², ws²)
        rel = rel.permute(1, 2, 0).contiguous()
        rel[..., 0] += window_size - 1
        rel[..., 1] += window_size - 1
        rel[..., 0] *= 2 * window_size - 1
        self.register_buffer('rpb_idx', rel.sum(-1).long())   # (ws², ws²)

    def forward(self, x: torch.Tensor,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        B_, N, C = x.shape
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
        attn = (q * self.scale) @ k.transpose(-2, -1)
        # Relative position bias
        rpb = self.rpb[self.rpb_idx.view(-1)].view(N, N, self.num_heads)
        attn = attn + rpb.permute(2, 0, 1).unsqueeze(0)
        if mask is not None:
            nW = mask.shape[0]
            attn = (attn.view(B_ // nW, nW, self.num_heads, N, N)
                    + mask.unsqueeze(1).unsqueeze(0)).view(B_, self.num_heads, N, N)
        attn = torch.softmax(attn, dim=-1)
        return self.proj((attn @ v).transpose(1, 2).reshape(B_, N, C))


class _ChannelAttention(nn.Module):
    """
    SE-style channel attention for sequence tensors (B, L, C).
    Applied after each window-attention block in HAT; recalibrates channels
    that local windows cannot correlate.
    """

    def __init__(self, dim: int, reduction: int = 16):
        super().__init__()
        mid = max(dim // reduction, 4)
        self.fc = nn.Sequential(
            nn.Linear(dim, mid), nn.ReLU(inplace=True),
            nn.Linear(mid, dim), nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, L, C).  Returns scale * x."""
        scale = self.fc(x.mean(dim=1))      # (B, C)
        return x * scale.unsqueeze(1)


class _MLP(nn.Module):
    def __init__(self, dim: int, ratio: float = 4.0):
        super().__init__()
        mid = int(dim * ratio)
        self.net = nn.Sequential(nn.Linear(dim, mid), nn.GELU(), nn.Linear(mid, dim))

    def forward(self, x):
        return self.net(x)


class HATBlock(nn.Module):
    """
    One HAT transformer block:
      LN → W-MSA (or SW-MSA) → residual → LN → FFN → residual → CA
    Alternates regular / shifted windows based on `shift` flag.
    """

    def __init__(self, dim: int, num_heads: int, window_size: int = 8,
                 shift: bool = False, mlp_ratio: float = 4.0):
        super().__init__()
        self.ws = window_size
        self.shift = window_size // 2 if shift else 0
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.attn = _WindowAttention(dim, window_size, num_heads)
        self.ca = _ChannelAttention(dim)
        self.mlp = _MLP(dim, mlp_ratio)

    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        """x: (B, H*W, C)"""
        B, L, C = x.shape
        ws = self.ws
        # Pad to nearest multiple of ws
        ph = (ws - H % ws) % ws
        pw = (ws - W % ws) % ws
        Hp, Wp = H + ph, W + pw
        x_2d = x.view(B, H, W, C)
        if ph > 0 or pw > 0:
            x_2d = F.pad(x_2d, (0, 0, 0, pw, 0, ph))
        # Cyclic shift for SW-MSA
        shifted = torch.roll(x_2d, (-self.shift, -self.shift), (1, 2)) \
            if self.shift > 0 else x_2d
        # Window partition → attention → reverse
        wins = _win_partition(shifted, ws).view(-1, ws * ws, C)
        mask = _sw_attn_mask(Hp, Wp, ws, self.shift, x.device)
        attn_out = self.attn(self.norm1(wins), mask=mask)
        wins = wins + attn_out
        x_2d = _win_reverse(wins.view(-1, ws, ws, C), ws, Hp, Wp)
        # Reverse shift and remove padding
        if self.shift > 0:
            x_2d = torch.roll(x_2d, (self.shift, self.shift), (1, 2))
        if ph > 0 or pw > 0:
            x_2d = x_2d[:, :H, :W, :].contiguous()
        x = x + x_2d.view(B, L, C)
        # FFN
        x = x + self.mlp(self.norm2(x))
        # Channel attention (multiplicative, no extra residual — HAT §3.2)
        x = self.ca(x)
        return x


class _RHAG(nn.Module):
    """
    Residual Hybrid Attention Group: N HATBlocks + one aggregation conv.
    The conv learns to merge window-attention outputs; the group residual
    re-injects low-frequency content that attention may suppress.
    """

    def __init__(self, dim: int, depth: int, num_heads: int,
                 window_size: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.blocks = nn.ModuleList([
            HATBlock(dim, num_heads, window_size,
                     shift=(i % 2 == 1), mlp_ratio=mlp_ratio)
            for i in range(depth)
        ])
        self.conv = nn.Conv2d(dim, dim, 3, 1, 1)

    def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor:
        residual = x
        for blk in self.blocks:
            x = blk(x, H, W)
        # Reshape to (B, C, H, W) for conv, then back
        x_2d = x.view(-1, H, W, x.shape[-1]).permute(0, 3, 1, 2).contiguous()
        x_2d = self.conv(x_2d)
        x = x_2d.flatten(2).transpose(1, 2)    # (B, H*W, C)
        return x + residual


class HATNet(nn.Module):
    """
    HAT super-resolution network.
    Drop-in replacement for RRDBNet; same interface (in_channels, out_channels,
    num_feat, scale) plus HAT-specific kwargs.

    Default settings (num_feat=64, num_groups=6, depth=6, window_size=8)
    give a model ~2× larger than RRDB-23 but with +0.3–1.2 dB PSNR advantage.
    Use num_groups=4, depth=4 for a lighter variant.
    """

    def __init__(self, in_channels: int = 3, out_channels: int = 3,
                 num_feat: int = 64, num_groups: int = 6, depth: int = 6,
                 num_heads: int = 4, window_size: int = 8,
                 mlp_ratio: float = 4.0, scale: int = 2,
                 # ignored kwargs for API compat with RRDBNet callers:
                 num_block: int = 0, num_grow_ch: int = 0):
        super().__init__()
        self.scale = scale
        self.conv_first = nn.Conv2d(in_channels, num_feat, 3, 1, 1)
        self.groups = nn.ModuleList([
            _RHAG(num_feat, depth, num_heads, window_size, mlp_ratio)
            for _ in range(num_groups)
        ])
        self.norm = nn.LayerNorm(num_feat)
        self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.upsample = _make_upsampler(num_feat, scale)
        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_last = nn.Conv2d(num_feat, out_channels, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(0.2, inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        feat = self.conv_first(x)                           # (B, D, H, W)
        seq = feat.flatten(2).transpose(1, 2)               # (B, H*W, D)
        for grp in self.groups:
            seq = grp(seq, H, W)
        seq = self.norm(seq)
        x_2d = seq.view(B, H, W, feat.shape[1]).permute(0, 3, 1, 2).contiguous()
        feat = feat + self.conv_after_body(x_2d)
        feat = self.upsample(feat)
        return self.conv_last(self.lrelu(self.conv_hr(feat)))


# ─────────────────────────────────────────────────────────────────────────────
# Restormer Denoiser  (for preprocessing historical scans)
# Zamir et al., CVPR 2022  [arXiv:2111.09881]
#
# Key idea: Multi-Dconv Head Transposed Attention (MDTA) computes attention
# across the channel dimension — O(C²·HW) instead of O(HW²·C) — enabling
# full-resolution processing without tiling. The Gated-Dconv FFN (GDFN)
# provides additional spatial context via depthwise convolutions.
# ─────────────────────────────────────────────────────────────────────────────

class _MDTA(nn.Module):
    """Multi-Dconv Head Transposed Attention."""

    def __init__(self, dim: int, num_heads: int, bias: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
        self.qkv = nn.Conv2d(dim, dim * 3, 1, bias=bias)
        self.qkv_dw = nn.Conv2d(dim * 3, dim * 3, 3, 1, 1, groups=dim * 3, bias=bias)
        self.proj_out = nn.Conv2d(dim, dim, 1, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        hd = C // self.num_heads
        qkv = self.qkv_dw(self.qkv(x))
        q, k, v = qkv.chunk(3, dim=1)
        q = q.reshape(B, self.num_heads, hd, H * W)
        k = k.reshape(B, self.num_heads, hd, H * W)
        v = v.reshape(B, self.num_heads, hd, H * W)
        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)
        attn = (q @ k.transpose(-2, -1)) * self.temperature   # (B, nh, hd, hd)
        attn = attn.softmax(dim=-1)
        out = (attn @ v).reshape(B, C, H, W)
        return self.proj_out(out)


class _GDFN(nn.Module):
    """Gated-Dconv Feed-Forward Network."""

    def __init__(self, dim: int, ratio: float = 2.66, bias: bool = False):
        super().__init__()
        mid = int(dim * ratio)
        self.proj_in = nn.Conv2d(dim, mid * 2, 1, bias=bias)
        self.dw = nn.Conv2d(mid * 2, mid * 2, 3, 1, 1, groups=mid * 2, bias=bias)
        self.proj_out = nn.Conv2d(mid, dim, 1, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = self.dw(self.proj_in(x)).chunk(2, dim=1)
        return self.proj_out(F.gelu(x1) * x2)


class _RestormerBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, ffn_ratio: float = 2.66):
        super().__init__()
        # LayerNorm on (B,C,H,W) via permute
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.attn = _MDTA(dim, num_heads)
        self.ffn = _GDFN(dim, ffn_ratio)

    def _ln(self, norm, x):
        B, C, H, W = x.shape
        return norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self._ln(self.norm1, x))
        x = x + self.ffn(self._ln(self.norm2, x))
        return x


class RestormerDenoiser(nn.Module):
    """
    4-level Restormer U-Net for blind denoising of historical aerial scans.

    Handles film grain (signal-dependent structured noise), scanner banding,
    defocus blur, and motion blur — all degradations present in scanned
    mid-20th century aerial photography.

    Input/output: same spatial resolution (denoising, not SR).
    The network predicts the residual: output = input + learned_residual.
    """

    def __init__(self, in_channels: int = 1, dim: int = 48,
                 num_blocks: Tuple[int, ...] = (4, 6, 6, 8),
                 num_heads: Tuple[int, ...] = (1, 2, 4, 8),
                 ffn_ratio: float = 2.66):
        super().__init__()
        self.proj_in = nn.Conv2d(in_channels, dim, 3, 1, 1, bias=False)

        # Encoder
        self.enc1 = nn.Sequential(*[_RestormerBlock(dim, num_heads[0], ffn_ratio)
                                     for _ in range(num_blocks[0])])
        self.down1 = nn.Conv2d(dim, dim * 2, 2, stride=2, bias=False)
        self.enc2 = nn.Sequential(*[_RestormerBlock(dim * 2, num_heads[1], ffn_ratio)
                                     for _ in range(num_blocks[1])])
        self.down2 = nn.Conv2d(dim * 2, dim * 4, 2, stride=2, bias=False)
        self.enc3 = nn.Sequential(*[_RestormerBlock(dim * 4, num_heads[2], ffn_ratio)
                                     for _ in range(num_blocks[2])])
        self.down3 = nn.Conv2d(dim * 4, dim * 8, 2, stride=2, bias=False)

        # Bottleneck
        self.bottleneck = nn.Sequential(*[_RestormerBlock(dim * 8, num_heads[3], ffn_ratio)
                                           for _ in range(num_blocks[3])])

        # Decoder  (PixelShuffle upsamplers + skip connections)
        self.up3 = nn.Sequential(nn.Conv2d(dim * 8, dim * 16, 1, bias=False), nn.PixelShuffle(2))
        self.reduce3 = nn.Conv2d(dim * 8, dim * 4, 1, bias=False)   # after cat with enc3 skip
        self.dec3 = nn.Sequential(*[_RestormerBlock(dim * 4, num_heads[2], ffn_ratio)
                                     for _ in range(num_blocks[2])])

        self.up2 = nn.Sequential(nn.Conv2d(dim * 4, dim * 8, 1, bias=False), nn.PixelShuffle(2))
        self.reduce2 = nn.Conv2d(dim * 4, dim * 2, 1, bias=False)
        self.dec2 = nn.Sequential(*[_RestormerBlock(dim * 2, num_heads[1], ffn_ratio)
                                     for _ in range(num_blocks[1])])

        self.up1 = nn.Sequential(nn.Conv2d(dim * 2, dim * 4, 1, bias=False), nn.PixelShuffle(2))
        self.reduce1 = nn.Conv2d(dim * 2, dim, 1, bias=False)
        self.dec1 = nn.Sequential(*[_RestormerBlock(dim, num_heads[0], ffn_ratio)
                                     for _ in range(num_blocks[0])])

        self.proj_out = nn.Conv2d(dim, in_channels, 3, 1, 1, bias=False)

    def _pad(self, x: torch.Tensor, factor: int = 8) -> Tuple[torch.Tensor, int, int]:
        """Pad to nearest multiple of `factor` for the 3-level downsampling."""
        _, _, H, W = x.shape
        ph = (factor - H % factor) % factor
        pw = (factor - W % factor) % factor
        if ph > 0 or pw > 0:
            x = F.pad(x, (0, pw, 0, ph), mode='reflect')
        return x, ph, pw

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        inp = x
        x, ph, pw = self._pad(x, factor=8)
        x0 = self.proj_in(x)

        e1 = self.enc1(x0)
        e2 = self.enc2(self.down1(e1))
        e3 = self.enc3(self.down2(e2))
        b = self.bottleneck(self.down3(e3))

        d3 = self.reduce3(torch.cat([self.up3(b), e3], dim=1))
        d3 = self.dec3(d3)
        d2 = self.reduce2(torch.cat([self.up2(d3), e2], dim=1))
        d2 = self.dec2(d2)
        d1 = self.reduce1(torch.cat([self.up1(d2), e1], dim=1))
        d1 = self.dec1(d1)

        out = x + self.proj_out(d1)   # residual denoising
        # Remove padding
        _, _, H, W = inp.shape
        return out[:, :, :H, :W]


# ─────────────────────────────────────────────────────────────────────────────
# Architecture factory
# ─────────────────────────────────────────────────────────────────────────────

def build_generator(arch: str, in_channels: int, scale: int,
                    num_feat: int, num_blocks: int,
                    num_groups: int, num_heads: int, window_size: int) -> nn.Module:
    if arch == "hat":
        return HATNet(
            in_channels=in_channels, out_channels=in_channels,
            num_feat=num_feat, num_groups=num_groups,
            depth=num_blocks, num_heads=num_heads,
            window_size=window_size, scale=scale,
        )
    else:  # rrdb
        return RRDBNet(
            in_channels=in_channels, out_channels=in_channels,
            num_feat=num_feat, num_block=num_blocks,
            num_grow_ch=32, scale=scale,
        )


# ─────────────────────────────────────────────────────────────────────────────
# Discriminator — U-Net with spectral normalisation (unchanged from v1)
# ─────────────────────────────────────────────────────────────────────────────

class UNetDiscriminator(nn.Module):
    def __init__(self, in_channels: int = 3, num_feat: int = 64):
        super().__init__()
        sn = nn.utils.spectral_norm

        def blk(ic, oc, stride=1):
            return nn.Sequential(sn(nn.Conv2d(ic, oc, 3, stride, 1)),
                                 nn.LeakyReLU(0.2, inplace=True))

        self.e0 = blk(in_channels, num_feat)
        self.e1 = blk(num_feat, num_feat * 2, stride=2)
        self.e2 = blk(num_feat * 2, num_feat * 4, stride=2)
        self.e3 = blk(num_feat * 4, num_feat * 8, stride=2)
        self.d3 = blk(num_feat * 8, num_feat * 4)
        self.d2 = blk(num_feat * 8, num_feat * 2)
        self.d1 = blk(num_feat * 4, num_feat)
        self.d0 = blk(num_feat * 2, num_feat)
        self.out = sn(nn.Conv2d(num_feat, 1, 3, 1, 1))
        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)

    def forward(self, x):
        e0 = self.e0(x);  e1 = self.e1(e0);  e2 = self.e2(e1);  e3 = self.e3(e2)
        d3 = self.d3(e3)
        d2 = self.d2(torch.cat([self.up(d3), e2], dim=1))
        d1 = self.d1(torch.cat([self.up(d2), e1], dim=1))
        d0 = self.d0(torch.cat([self.up(d1), e0], dim=1))
        return self.out(d0)


# ─────────────────────────────────────────────────────────────────────────────
# Loss functions
# ─────────────────────────────────────────────────────────────────────────────

class VGGPerceptualLoss(nn.Module):
    """VGG-19 perceptual loss using relu3_4 and relu4_4 features."""

    def __init__(self):
        super().__init__()
        vgg = models.vgg19(weights=models.VGG19_Weights.DEFAULT)
        self.slice1 = nn.Sequential(*list(vgg.features)[:18]).eval()
        self.slice2 = nn.Sequential(*list(vgg.features)[18:27]).eval()
        for p in self.parameters():
            p.requires_grad = False
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std",  torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        if pred.shape[1] == 1:
            pred   = pred.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        pred   = (pred   - self.mean) / self.std
        target = (target - self.mean) / self.std
        f1_p = self.slice1(pred);   f1_t = self.slice1(target)
        f2_p = self.slice2(f1_p);   f2_t = self.slice2(f1_t)
        return F.l1_loss(f1_p, f1_t) + F.l1_loss(f2_p, f2_t)


# ─────────────────────────────────────────────────────────────────────────────
# Image I/O
# ─────────────────────────────────────────────────────────────────────────────

IMAGE_EXTS = {".tif", ".tiff", ".jpg", ".jpeg", ".png"}


def is_image(p: Path) -> bool:
    return p.suffix.lower() in IMAGE_EXTS


def detect_channels(directory: Path) -> int:
    for path in sorted(directory.rglob("*")):
        if not is_image(path):
            continue
        if HAS_RASTERIO and path.suffix.lower() in (".tif", ".tiff"):
            try:
                with rasterio.open(str(path)) as src:
                    return 1 if src.count < 3 else 3
            except Exception:
                pass
        img = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if img is not None:
            return 1 if (img.ndim == 2 or img.shape[2] == 1) else 3
    raise ValueError(f"No readable images found in {directory}")


def read_image(path: Path, num_channels: int) -> Optional[np.ndarray]:
    """
    Read as float32 (H, W, C) in [0, 1].  C = num_channels.
    GeoTIFFs get per-band 2nd/98th percentile stretch.
    """
    if HAS_RASTERIO and path.suffix.lower() in (".tif", ".tiff"):
        try:
            with rasterio.open(str(path)) as src:
                if num_channels == 1:
                    img = src.read(1).astype(np.float32)[:, :, np.newaxis]
                else:
                    if src.count < 3:
                        print(f"  Warning: {path.name} has {src.count} band(s), need 3 — skip")
                        return None
                    img = src.read([1, 2, 3]).astype(np.float32).transpose(1, 2, 0)
                for c in range(num_channels):
                    lo = np.percentile(img[:, :, c], 2)
                    hi = np.percentile(img[:, :, c], 98)
                    if hi > lo:
                        img[:, :, c] = np.clip((img[:, :, c] - lo) / (hi - lo), 0.0, 1.0)
                    else:
                        img[:, :, c] = 0.0
                return img
        except Exception as e:
            print(f"  [rasterio] {path.name}: {e} — falling back to OpenCV")

    if num_channels == 1:
        gray = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
        if gray is None:
            return None
        return (gray.astype(np.float32) / 255.0)[:, :, np.newaxis]
    else:
        bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
        if bgr is None:
            return None
        return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0


def histogram_match(src: np.ndarray, ref: np.ndarray) -> np.ndarray:
    """Per-channel histogram matching with 4096-bin quantisation."""
    BINS = 4096
    scale_q = BINS - 1
    out = src.copy()
    for c in range(src.shape[2]):
        s_q = (src[:, :, c] * scale_q).clip(0, scale_q).astype(np.uint16)
        r_q = (ref[:, :, c] * scale_q).clip(0, scale_q).astype(np.uint16)
        s_vals, s_inv, s_cnt = np.unique(s_q, return_inverse=True, return_counts=True)
        r_vals, r_cnt = np.unique(r_q, return_counts=True)
        s_cdf = np.cumsum(s_cnt).astype(np.float64); s_cdf /= s_cdf[-1]
        r_cdf = np.cumsum(r_cnt).astype(np.float64); r_cdf /= r_cdf[-1]
        lut = np.interp(s_cdf, r_cdf, r_vals.astype(np.float64))
        out[:, :, c] = (lut[s_inv].reshape(src.shape[:2]) / scale_q).astype(np.float32)
    return out


def save_sr_image(sr_img: np.ndarray, src_path: Path,
                  dst_path: Path, scale: int) -> None:
    """Save SR image, preserving GeoTIFF geotransform scaled by `scale`."""
    dst_path.parent.mkdir(parents=True, exist_ok=True)
    num_ch = sr_img.shape[2]
    if HAS_RASTERIO and src_path.suffix.lower() in (".tif", ".tiff"):
        dst_tif = dst_path.with_suffix(".tif")
        try:
            with rasterio.open(str(src_path)) as src:
                profile = src.profile.copy()
                tf = src.transform
                new_tf = rasterio.transform.Affine(
                    tf.a / scale, tf.b / scale, tf.c,
                    tf.d / scale, tf.e / scale, tf.f,
                )
                profile.update(driver="GTiff", width=sr_img.shape[1],
                                height=sr_img.shape[0], count=num_ch,
                                dtype="uint8", transform=new_tf, compress="lzw")
                img_u8 = (sr_img * 255).clip(0, 255).astype(np.uint8)
                with rasterio.open(str(dst_tif), "w", **profile) as dst:
                    for i in range(num_ch):
                        dst.write(img_u8[:, :, i], i + 1)
            return
        except Exception as e:
            print(f"  [rasterio] {dst_path.name}: {e} — falling back to PNG")
            dst_path = dst_path.with_suffix(".png")
    img_u8 = (sr_img * 255).clip(0, 255).astype(np.uint8)
    if num_ch == 1:
        cv2.imwrite(str(dst_path), img_u8[:, :, 0])
    else:
        cv2.imwrite(str(dst_path), cv2.cvtColor(img_u8, cv2.COLOR_RGB2BGR))


def save_image(img: np.ndarray, src_path: Path, dst_path: Path) -> None:
    """Save same-resolution image (denoiser output), copying GeoTIFF metadata."""
    dst_path.parent.mkdir(parents=True, exist_ok=True)
    num_ch = img.shape[2]
    if HAS_RASTERIO and src_path.suffix.lower() in (".tif", ".tiff"):
        dst_tif = dst_path.with_suffix(".tif")
        try:
            with rasterio.open(str(src_path)) as src:
                profile = src.profile.copy()
                profile.update(count=num_ch, dtype="uint8", compress="lzw")
                img_u8 = (img * 255).clip(0, 255).astype(np.uint8)
                with rasterio.open(str(dst_tif), "w", **profile) as dst:
                    for i in range(num_ch):
                        dst.write(img_u8[:, :, i], i + 1)
            return
        except Exception as e:
            print(f"  [rasterio] {dst_path.name}: {e} — falling back to PNG")
            dst_path = dst_path.with_suffix(".png")
    img_u8 = (img * 255).clip(0, 255).astype(np.uint8)
    cv2.imwrite(str(dst_path),
                img_u8[:, :, 0] if num_ch == 1
                else cv2.cvtColor(img_u8, cv2.COLOR_RGB2BGR))


# ─────────────────────────────────────────────────────────────────────────────
# Degradation pipeline
# ─────────────────────────────────────────────────────────────────────────────

def degrade(img: np.ndarray, scale: int,
            sigma_blur: float = 1.2, sigma_noise: float = 0.005,
            historical: bool = False) -> np.ndarray:
    """
    Simulate realistic LR degradation:
      Standard:   Gaussian blur → bicubic downsample → Gaussian noise
      Historical: (above) + film grain + vignetting + scanner banding +
                  film-base fog

    'Historical' mode better matches mid-20th century aerial scans from
    panchromatic film digitised at 15–25 µm on flatbed / drum scanners.
    """
    h, w = img.shape[:2]
    s = max(0.1, sigma_blur + random.uniform(-0.3, 0.3))
    ksize = int(2 * round(3 * s) + 1)
    blurred = cv2.GaussianBlur(img, (ksize, ksize), s)

    # Bicubic downsample
    lr = cv2.resize(blurred, (w // scale, h // scale), interpolation=cv2.INTER_CUBIC)
    if lr.ndim == 2:
        lr = lr[:, :, np.newaxis]

    # Standard Gaussian noise (sensor)
    if sigma_noise > 0:
        lr = np.clip(lr + np.random.normal(0, sigma_noise, lr.shape).astype(np.float32), 0, 1)

    if historical:
        lh, lw = lr.shape[:2]

        # Film grain: signal-dependent (brighter = finer grain, darker = coarser)
        # Approximation: grain std ∝ sqrt(luminance + eps) as in photographic emulsion
        grain_strength = random.uniform(0.008, 0.025)
        lum = lr.mean(axis=2, keepdims=True) if lr.shape[2] > 1 else lr
        local_sigma = grain_strength * np.sqrt(lum + 0.05)
        grain = np.random.normal(0, 1, lr.shape).astype(np.float32) * local_sigma
        lr = np.clip(lr + grain, 0, 1)

        # Vignetting: radial darkening toward image corners
        if random.random() > 0.4:
            vig_strength = random.uniform(0.05, 0.25)
            cy, cx = lh / 2, lw / 2
            yy, xx = np.mgrid[:lh, :lw].astype(np.float32)
            r = np.sqrt(((yy - cy) / cy) ** 2 + ((xx - cx) / cx) ** 2)
            vignette = np.clip(1.0 - vig_strength * r ** 2, 0.1, 1.0)
            lr = lr * vignette[:, :, np.newaxis]

        # Scanner banding: faint horizontal striping at fixed spatial frequency
        if random.random() > 0.6:
            band_amp = random.uniform(0.002, 0.010)
            freq = random.uniform(8, 40)  # cycles per image height
            y_coords = np.arange(lh, dtype=np.float32)
            banding = (band_amp * np.sin(2 * math.pi * freq * y_coords / lh)
                       ).astype(np.float32)
            lr = np.clip(lr + banding[:, np.newaxis, np.newaxis], 0, 1)

        # Film-base fog: slight additive grey offset (print-through / base+fog)
        fog = random.uniform(0.0, 0.03)
        lr = np.clip(lr + fog, 0, 1)

    return lr


# ─────────────────────────────────────────────────────────────────────────────
# Denoising degradation (for training RestormerDenoiser)
# ─────────────────────────────────────────────────────────────────────────────

def degrade_for_denoising(img: np.ndarray,
                          sigma_noise: float = 0.02,
                          historical: bool = False) -> np.ndarray:
    """
    Add noise/artifacts to a clean HR image to produce a noisy training pair.
    Same resolution as input — used for denoiser training (not SR).
    """
    noisy = img.copy()

    if historical:
        # Film grain
        grain_strength = random.uniform(0.010, 0.030)
        lum = noisy.mean(axis=2, keepdims=True) if noisy.shape[2] > 1 else noisy
        local_sigma = grain_strength * np.sqrt(lum + 0.05)
        noisy = np.clip(noisy + np.random.normal(0, 1, noisy.shape).astype(np.float32)
                        * local_sigma, 0, 1)
        # Scanner banding
        if random.random() > 0.5:
            h = noisy.shape[0]
            band_amp = random.uniform(0.003, 0.012)
            banding = (band_amp * np.sin(2 * math.pi * random.uniform(5, 30)
                                         * np.arange(h, dtype=np.float32) / h)
                       ).astype(np.float32)
            noisy = np.clip(noisy + banding[:, np.newaxis, np.newaxis], 0, 1)
        # Mild blur (defocus / motion)
        if random.random() > 0.5:
            s = random.uniform(0.3, 1.2)
            ksize = int(2 * round(3 * s) + 1)
            noisy = cv2.GaussianBlur(noisy, (ksize, ksize), s)
            if noisy.ndim == 2:
                noisy = noisy[:, :, np.newaxis]
        # Vignetting
        if random.random() > 0.5:
            h, w = noisy.shape[:2]
            cy, cx = h / 2, w / 2
            yy, xx = np.mgrid[:h, :w].astype(np.float32)
            r = np.sqrt(((yy - cy) / cy) ** 2 + ((xx - cx) / cx) ** 2)
            noisy = noisy * np.clip(1 - random.uniform(0.05, 0.2) * r ** 2, 0.1, 1)[:, :, None]
    else:
        # AWGN with random sigma
        sigma = random.uniform(sigma_noise * 0.5, sigma_noise * 2.0)
        noisy = np.clip(noisy + np.random.normal(0, sigma, noisy.shape).astype(np.float32), 0, 1)

    return noisy.astype(np.float32)


# ─────────────────────────────────────────────────────────────────────────────
# Datasets
# ─────────────────────────────────────────────────────────────────────────────

class SRDataset(Dataset):
    """
    SR training dataset. Generates synthetic LR/HR pairs from HR images.
    Optional LR directory for tone calibration via histogram matching.
    """

    def __init__(self, hr_dir: Path, scale: int, num_channels: int,
                 patch_size: int = 128, patches_per_image: int = 30,
                 lr_dir: Optional[Path] = None,
                 sigma_blur: float = 1.2, sigma_noise: float = 0.005,
                 historical: bool = False):
        self.scale = scale
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.patches_per_image = patches_per_image
        self.sigma_blur = sigma_blur
        self.sigma_noise = sigma_noise
        self.historical = historical

        hr_files = sorted([f for f in hr_dir.rglob("*") if is_image(f)])
        if not hr_files:
            raise ValueError(f"No images in {hr_dir}")
        print(f"  HR images found: {len(hr_files)}")

        self.hr_images: List[np.ndarray] = []
        for f in tqdm(hr_files, desc="  Loading HR images"):
            img = read_image(f, num_channels)
            if img is None:
                continue
            if img.shape[0] >= patch_size * scale and img.shape[1] >= patch_size * scale:
                self.hr_images.append(img)
            else:
                print(f"    Skip {f.name} (too small)")

        if not self.hr_images:
            raise ValueError(f"No usable HR images (need ≥{patch_size * scale}px).")
        print(f"  Usable HR images: {len(self.hr_images)}")

        self.lr_colour_ref: Optional[np.ndarray] = None
        if lr_dir is not None:
            lr_files = sorted([f for f in lr_dir.rglob("*") if is_image(f)])
            refs = [read_image(f, num_channels) for f in lr_files[:10]]
            refs = [r for r in refs if r is not None]
            if refs:
                patches = []
                for r in refs:
                    h, w = r.shape[:2]
                    for _ in range(5):
                        y = random.randint(0, max(0, h - 64))
                        x = random.randint(0, max(0, w - 64))
                        patches.append(r[y:y + 64, x:x + 64])
                self.lr_colour_ref = np.concatenate(patches, axis=0)
                print(f"  Tone reference built from {len(refs)} LR image(s)")

    def __len__(self):
        return len(self.hr_images) * self.patches_per_image

    def __getitem__(self, idx):
        hr = self.hr_images[idx // self.patches_per_image]
        h, w = hr.shape[:2]
        ph = self.patch_size * self.scale
        y = random.randint(0, h - ph)
        x = random.randint(0, w - ph)
        hr_patch = hr[y:y + ph, x:x + ph].copy()
        if random.random() > 0.5:
            hr_patch = np.fliplr(hr_patch).copy()
        if random.random() > 0.5:
            hr_patch = np.flipud(hr_patch).copy()
        hr_patch = np.rot90(hr_patch, random.randint(0, 3)).copy()
        lr_patch = degrade(hr_patch, self.scale, self.sigma_blur,
                           self.sigma_noise, self.historical)
        if self.lr_colour_ref is not None:
            lr_patch = histogram_match(lr_patch, self.lr_colour_ref)
        return (torch.from_numpy(lr_patch.transpose(2, 0, 1)).float(),
                torch.from_numpy(hr_patch.transpose(2, 0, 1)).float())


class DenoiserDataset(Dataset):
    """
    Denoiser training dataset: generates noisy/clean pairs at full HR resolution.
    """

    def __init__(self, hr_dir: Path, num_channels: int,
                 patch_size: int = 256, patches_per_image: int = 30,
                 sigma_noise: float = 0.02, historical: bool = False):
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.patches_per_image = patches_per_image
        self.sigma_noise = sigma_noise
        self.historical = historical

        hr_files = sorted([f for f in hr_dir.rglob("*") if is_image(f)])
        if not hr_files:
            raise ValueError(f"No images in {hr_dir}")
        print(f"  HR images found: {len(hr_files)}")

        self.hr_images: List[np.ndarray] = []
        for f in tqdm(hr_files, desc="  Loading HR images"):
            img = read_image(f, num_channels)
            if img is not None and img.shape[0] >= patch_size and img.shape[1] >= patch_size:
                self.hr_images.append(img)
        if not self.hr_images:
            raise ValueError("No usable HR images for denoiser training.")
        print(f"  Usable images: {len(self.hr_images)}")

    def __len__(self):
        return len(self.hr_images) * self.patches_per_image

    def __getitem__(self, idx):
        hr = self.hr_images[idx // self.patches_per_image]
        h, w = hr.shape[:2]
        ps = self.patch_size
        y = random.randint(0, h - ps)
        x = random.randint(0, w - ps)
        clean = hr[y:y + ps, x:x + ps].copy()
        if random.random() > 0.5:
            clean = np.fliplr(clean).copy()
        if random.random() > 0.5:
            clean = np.flipud(clean).copy()
        noisy = degrade_for_denoising(clean, self.sigma_noise, self.historical)
        return (torch.from_numpy(noisy.transpose(2, 0, 1)).float(),
                torch.from_numpy(clean.transpose(2, 0, 1)).float())


# ─────────────────────────────────────────────────────────────────────────────
# Checkpoint helpers
# ─────────────────────────────────────────────────────────────────────────────

def _load_checkpoint(model_path: Path, device: torch.device) -> dict:
    try:
        return torch.load(model_path, map_location=device, weights_only=True)
    except Exception:
        return torch.load(model_path, map_location=device, weights_only=False)


def _load_generator(model_path: Path, device: torch.device):
    """Load SR generator from checkpoint. Returns (model, scale, in_channels)."""
    ckpt = _load_checkpoint(model_path, device)
    scale       = ckpt.get("scale", 2)
    num_feat    = ckpt.get("num_feat", 64)
    num_blocks  = ckpt.get("num_blocks", 23)
    in_channels = ckpt.get("in_channels", 3)
    arch        = ckpt.get("arch", "rrdb")
    num_groups  = ckpt.get("num_groups", 6)
    num_heads   = ckpt.get("num_heads", 4)
    window_size = ckpt.get("window_size", 8)
    net_g = build_generator(arch, in_channels, scale, num_feat, num_blocks,
                             num_groups, num_heads, window_size).to(device)
    net_g.load_state_dict(ckpt["g_state"])
    net_g.eval()
    return net_g, scale, in_channels


def _load_denoiser(model_path: Path, device: torch.device) -> RestormerDenoiser:
    ckpt = _load_checkpoint(model_path, device)
    in_channels = ckpt.get("in_channels", 1)
    dim         = ckpt.get("dim", 48)
    model = RestormerDenoiser(in_channels=in_channels, dim=dim).to(device)
    model.load_state_dict(ckpt["state"])
    model.eval()
    return model


# ─────────────────────────────────────────────────────────────────────────────
# Training — SR
# ─────────────────────────────────────────────────────────────────────────────

def train(args: argparse.Namespace) -> None:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    hr_dir     = Path(args.hr_dir)
    lr_dir     = Path(args.lr_dir) if args.lr_dir else None
    model_path = Path(args.model)
    model_path.parent.mkdir(parents=True, exist_ok=True)

    print("\nDetecting channel count from HR directory...")
    num_channels = detect_channels(hr_dir)
    mode_label = "panchromatic (1-ch)" if num_channels == 1 else "RGB (3-ch)"
    print(f"  Mode: {mode_label}  |  Arch: {args.arch}"
          + ("  |  Historical degradation ON" if args.historical else ""))

    print(f"\nBuilding dataset (scale={args.scale}×, patch={args.patch_size}px LR)...")
    dataset = SRDataset(
        hr_dir=hr_dir, scale=args.scale, num_channels=num_channels,
        patch_size=args.patch_size, patches_per_image=args.patches_per_image,
        lr_dir=lr_dir, sigma_blur=args.sigma_blur, sigma_noise=args.sigma_noise,
        historical=args.historical,
    )
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                        num_workers=args.num_workers,
                        pin_memory=(device.type == "cuda"), drop_last=True,
                        persistent_workers=(args.num_workers > 0),
                        prefetch_factor=(2 if args.num_workers > 0 else None))

    n_gpus = torch.cuda.device_count() if device.type == "cuda" else 0
    print(f"\nBuilding {args.arch.upper()} generator...")
    net_g = build_generator(args.arch, num_channels, args.scale,
                             args.num_feat, args.num_blocks,
                             args.num_groups, args.num_heads,
                             args.window_size).to(device)
    net_d = UNetDiscriminator(in_channels=num_channels, num_feat=64).to(device)
    perc_loss_fn = VGGPerceptualLoss().to(device)
    if n_gpus > 1:
        print(f"  Multi-GPU: {n_gpus} GPUs via DataParallel")
        net_g = nn.DataParallel(net_g)
        net_d = nn.DataParallel(net_d)

    start_epoch = 0
    if model_path.exists() and not args.restart:
        print(f"Resuming from {model_path}")
        ckpt = _load_checkpoint(model_path, device)
        ckpt_arch = ckpt.get("arch", "rrdb")
        ckpt_ch   = ckpt.get("in_channels", 3)
        if ckpt_arch != args.arch or ckpt_ch != num_channels:
            print(f"  Warning: checkpoint arch/channels mismatch — starting fresh")
        else:
            net_g.load_state_dict(ckpt["g_state"])
            if "d_state" in ckpt and not args.no_gan:
                try:
                    net_d.load_state_dict(ckpt["d_state"])
                except Exception:
                    print("  Note: discriminator weights not loaded — restarting D")
            start_epoch = ckpt.get("epoch", 0)
            print(f"  Resumed from epoch {start_epoch}")

    opt_g = optim.Adam(net_g.parameters(), lr=args.lr_g, betas=(0.9, 0.99))
    opt_d = optim.Adam(net_d.parameters(), lr=args.lr_d, betas=(0.9, 0.99))
    milestones = [args.epochs // 2, args.epochs * 3 // 4]
    sched_g = optim.lr_scheduler.MultiStepLR(opt_g, milestones, gamma=0.5,
                                              last_epoch=start_epoch - 1)
    sched_d = optim.lr_scheduler.MultiStepLR(opt_d, milestones, gamma=0.5,
                                              last_epoch=start_epoch - 1)
    use_amp = (device.type == "cuda") and not args.no_amp
    scaler_g = GradScaler(enabled=use_amp)
    scaler_d = GradScaler(enabled=use_amp)
    gan_warmup = max(1, args.epochs // 10)

    print(f"\nTraining {args.epochs} epochs (from {start_epoch + 1})")
    print(f"  GAN warmup: first {gan_warmup} epochs L1+Perceptual only")
    print(f"  Loss weights: L1={args.w_l1}  Perc={args.w_perc}  Adv={args.w_adv}")
    print(f"  AMP: {'enabled' if use_amp else 'disabled'}")
    if args.no_gan:
        print("  GAN disabled (--no-gan)")

    for epoch in range(start_epoch, args.epochs):
        net_g.train(); net_d.train()
        use_gan = (not args.no_gan) and (epoch >= gan_warmup)
        sum_l1 = sum_perc = sum_adv = sum_d = 0.0
        n_batches = 0
        pbar = tqdm(loader, desc=f"Epoch {epoch + 1:>4}/{args.epochs}", ncols=90)

        for lr_t, hr_t in pbar:
            lr_t = lr_t.to(device, non_blocking=True)
            hr_t = hr_t.to(device, non_blocking=True)
            with autocast(enabled=use_amp):
                sr_t = net_g(lr_t)

            if use_gan:
                opt_d.zero_grad(set_to_none=True)
                with autocast(enabled=use_amp):
                    loss_d = 0.5 * (torch.mean((net_d(hr_t) - 1.0) ** 2)
                                    + torch.mean(net_d(sr_t.detach()) ** 2))
                scaler_d.scale(loss_d).backward()
                scaler_d.unscale_(opt_d)
                nn.utils.clip_grad_norm_(net_d.parameters(), 1.0)
                scaler_d.step(opt_d); scaler_d.update()
                sum_d += loss_d.item()

            opt_g.zero_grad(set_to_none=True)
            with autocast(enabled=use_amp):
                loss_l1   = F.l1_loss(sr_t, hr_t)
                loss_perc = perc_loss_fn(sr_t.clamp(0, 1), hr_t.clamp(0, 1))
                loss_g = args.w_l1 * loss_l1 + args.w_perc * loss_perc
                if use_gan:
                    loss_adv = torch.mean((net_d(sr_t) - 1.0) ** 2)
                    loss_g = loss_g + args.w_adv * loss_adv
                    sum_adv += loss_adv.item()

            scaler_g.scale(loss_g).backward()
            scaler_g.unscale_(opt_g)
            nn.utils.clip_grad_norm_(net_g.parameters(), 1.0)
            scaler_g.step(opt_g); scaler_g.update()
            sum_l1 += loss_l1.item(); sum_perc += loss_perc.item()
            n_batches += 1
            pbar.set_postfix(
                L1=f"{sum_l1 / n_batches:.4f}",
                Perc=f"{sum_perc / n_batches:.4f}",
                GAN=f"{sum_adv / n_batches:.4f}" if use_gan else "off",
            )

        sched_g.step(); sched_d.step()

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            g_state = (net_g.module if isinstance(net_g, nn.DataParallel) else net_g).state_dict()
            d_state = (net_d.module if isinstance(net_d, nn.DataParallel) else net_d).state_dict()
            torch.save({
                "epoch": epoch + 1, "g_state": g_state, "d_state": d_state,
                "scale": args.scale, "arch": args.arch,
                "num_feat": args.num_feat, "num_blocks": args.num_blocks,
                "num_groups": args.num_groups, "num_heads": args.num_heads,
                "window_size": args.window_size, "in_channels": num_channels,
            }, model_path)
            print(f"  ✓ Checkpoint saved → {model_path}")

    print("\nTraining complete.")


# ─────────────────────────────────────────────────────────────────────────────
# Training — Denoiser
# ─────────────────────────────────────────────────────────────────────────────

def train_denoiser(args: argparse.Namespace) -> None:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    hr_dir     = Path(args.hr_dir)
    model_path = Path(args.model)
    model_path.parent.mkdir(parents=True, exist_ok=True)

    print("\nDetecting channel count...")
    num_channels = detect_channels(hr_dir)
    mode_label = "panchromatic" if num_channels == 1 else "RGB"
    print(f"  Mode: {mode_label}"
          + ("  |  Historical degradation ON" if args.historical else ""))

    print("\nBuilding denoiser dataset...")
    dataset = DenoiserDataset(
        hr_dir=hr_dir, num_channels=num_channels,
        patch_size=args.patch_size, patches_per_image=args.patches_per_image,
        sigma_noise=args.sigma_noise, historical=args.historical,
    )
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                        num_workers=args.num_workers,
                        pin_memory=(device.type == "cuda"), drop_last=True,
                        persistent_workers=(args.num_workers > 0),
                        prefetch_factor=(2 if args.num_workers > 0 else None))

    print(f"\nBuilding RestormerDenoiser (dim={args.dim})...")
    model = RestormerDenoiser(in_channels=num_channels, dim=args.dim).to(device)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)

    start_epoch = 0
    if model_path.exists() and not args.restart:
        print(f"Resuming from {model_path}")
        ckpt = _load_checkpoint(model_path, device)
        try:
            model.load_state_dict(ckpt["state"])
            start_epoch = ckpt.get("epoch", 0)
            print(f"  Resumed from epoch {start_epoch}")
        except Exception as e:
            print(f"  Warning: {e} — starting fresh")

    perc_loss_fn = VGGPerceptualLoss().to(device)
    optimizer = optim.AdamW(model.parameters(), lr=args.lr_g, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs - start_epoch, eta_min=1e-6)
    use_amp = (device.type == "cuda") and not args.no_amp
    scaler = GradScaler(enabled=use_amp)

    print(f"\nTraining denoiser for {args.epochs} epochs")

    for epoch in range(start_epoch, args.epochs):
        model.train()
        sum_l1 = sum_perc = 0.0
        n_batches = 0
        pbar = tqdm(loader, desc=f"Epoch {epoch + 1:>4}/{args.epochs}", ncols=90)
        for noisy_t, clean_t in pbar:
            noisy_t = noisy_t.to(device, non_blocking=True)
            clean_t = clean_t.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with autocast(enabled=use_amp):
                pred = model(noisy_t)
                loss_l1   = F.l1_loss(pred, clean_t)
                loss_perc = perc_loss_fn(pred.clamp(0, 1), clean_t.clamp(0, 1))
                loss = loss_l1 + 0.1 * loss_perc
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer); scaler.update()
            sum_l1 += loss_l1.item(); sum_perc += loss_perc.item()
            n_batches += 1
            pbar.set_postfix(L1=f"{sum_l1/n_batches:.4f}",
                             Perc=f"{sum_perc/n_batches:.4f}")
        scheduler.step()

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            m_state = (model.module if isinstance(model, nn.DataParallel) else model).state_dict()
            torch.save({
                "epoch": epoch + 1, "state": m_state,
                "in_channels": num_channels, "dim": args.dim,
            }, model_path)
            print(f"  ✓ Checkpoint saved → {model_path}")

    print("\nDenoiser training complete.")


# ─────────────────────────────────────────────────────────────────────────────
# Inference helpers
# ─────────────────────────────────────────────────────────────────────────────

@torch.no_grad()
def sr_tile_inference(model: nn.Module, img: np.ndarray, scale: int,
                      tile_size: int = 512, overlap: int = 64,
                      device: torch.device = torch.device("cpu"),
                      tile_batch_size: int = 4) -> np.ndarray:
    """
    Apply SR model to an arbitrarily large image using overlapping tiles
    blended with a 2D Hanning window to eliminate seam artefacts.
    """
    h, w, num_ch = img.shape
    out = np.zeros((h * scale, w * scale, num_ch), dtype=np.float32)
    weight = np.zeros((h * scale, w * scale, 1), dtype=np.float32)
    step = tile_size - overlap
    win_1d = np.hanning(tile_size * scale).astype(np.float32)
    win_2d = np.outer(win_1d, win_1d)[:, :, np.newaxis]

    def tile_starts(length):
        starts = list(range(0, length - tile_size, step))
        starts.append(max(0, length - tile_size))
        return sorted(set(starts))

    ys = tile_starts(h) if h > tile_size else [0]
    xs = tile_starts(w) if w > tile_size else [0]
    positions = [(y, x) for y in ys for x in xs
                 if min(tile_size, h - y) >= 4 and min(tile_size, w - x) >= 4]
    use_amp = device.type == "cuda"

    for batch_start in range(0, len(positions), tile_batch_size):
        batch_pos = positions[batch_start: batch_start + tile_batch_size]
        tensors, real_sizes = [], []
        for y, x in batch_pos:
            tile = img[y: y + tile_size, x: x + tile_size]
            th, tw = tile.shape[:2]
            real_sizes.append((th, tw))
            if th < tile_size or tw < tile_size:
                pad = np.zeros((tile_size, tile_size, num_ch), dtype=np.float32)
                pad[:th, :tw] = tile
                tile = pad
            tensors.append(torch.from_numpy(tile.transpose(2, 0, 1)).float())
        batch_t = torch.stack(tensors).to(device, non_blocking=True)
        with autocast(enabled=use_amp):
            sr_batch = model(batch_t)
        sr_batch = sr_batch.float().cpu().numpy()
        for i, (y, x) in enumerate(batch_pos):
            th, tw = real_sizes[i]
            sr = np.clip(sr_batch[i].transpose(1, 2, 0), 0, 1)
            if sr.ndim == 2:
                sr = sr[:, :, np.newaxis]
            oy, ox = y * scale, x * scale
            sy, sx = th * scale, tw * scale
            sr = sr[:sy, :sx]
            w_crop = win_2d[:sy, :sx]
            out[oy: oy + sy, ox: ox + sx] += sr * w_crop
            weight[oy: oy + sy, ox: ox + sx] += w_crop

    return np.clip(out / np.maximum(weight, 1e-8), 0, 1)


@torch.no_grad()
def denoise_tile_inference(model: RestormerDenoiser, img: np.ndarray,
                            tile_size: int = 1024, overlap: int = 128,
                            device: torch.device = torch.device("cpu")) -> np.ndarray:
    """
    Apply Restormer denoiser to an arbitrarily large image using tiled inference.
    Same resolution as input.
    """
    h, w, num_ch = img.shape
    out = np.zeros_like(img)
    weight = np.zeros((h, w, 1), dtype=np.float32)
    step = tile_size - overlap
    win_1d = np.hanning(tile_size).astype(np.float32)
    win_2d = np.outer(win_1d, win_1d)[:, :, np.newaxis]

    def tile_starts(length):
        starts = list(range(0, length - tile_size, step))
        starts.append(max(0, length - tile_size))
        return sorted(set(starts))

    ys = tile_starts(h) if h > tile_size else [0]
    xs = tile_starts(w) if w > tile_size else [0]
    use_amp = device.type == "cuda"

    for y in ys:
        for x in xs:
            tile = img[y: y + tile_size, x: x + tile_size]
            th, tw = tile.shape[:2]
            if th < 8 or tw < 8:
                continue
            if th < tile_size or tw < tile_size:
                pad = np.zeros((tile_size, tile_size, num_ch), dtype=np.float32)
                pad[:th, :tw] = tile
                tile = pad
            t = torch.from_numpy(tile.transpose(2, 0, 1)).float().unsqueeze(0).to(device)
            with autocast(enabled=use_amp):
                pred = model(t)
            pred_np = np.clip(pred.squeeze(0).float().cpu().numpy().transpose(1, 2, 0), 0, 1)
            pred_np = pred_np[:th, :tw]
            w_crop = win_2d[:th, :tw]
            out[y: y + th, x: x + tw] += pred_np * w_crop
            weight[y: y + th, x: x + tw] += w_crop

    return np.clip(out / np.maximum(weight, 1e-8), 0, 1)


# ─────────────────────────────────────────────────────────────────────────────
# Apply SR
# ─────────────────────────────────────────────────────────────────────────────

def apply_sr(args: argparse.Namespace) -> None:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    model_path = Path(args.model)
    if not model_path.exists():
        print(f"Error: model not found at {model_path}", file=sys.stderr)
        sys.exit(1)

    print(f"Loading SR model from {model_path}...")
    net_g, scale, in_channels = _load_generator(model_path, device)
    mode_label = "panchromatic" if in_channels == 1 else "RGB"
    print(f"  Scale: {scale}×  |  Mode: {mode_label}  |  tile: {args.tile_size}px")

    # Optional denoiser preprocessing
    denoiser = None
    if args.preprocess_model:
        pp_path = Path(args.preprocess_model)
        if not pp_path.exists():
            print(f"Warning: preprocess model not found at {pp_path} — skipping denoising")
        else:
            print(f"Loading Restormer denoiser from {pp_path}...")
            denoiser = _load_denoiser(pp_path, device)
            print("  Denoiser preprocessing: ENABLED")

    input_dir  = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    files = sorted([f for f in input_dir.rglob("*") if is_image(f)])
    if not files:
        print(f"No images found in {input_dir}", file=sys.stderr)
        sys.exit(1)

    print(f"\nProcessing {len(files)} image(s) → {output_dir}\n")
    ok = skipped = 0
    io_executor = ThreadPoolExecutor(max_workers=args.io_workers)

    def _read(src): return src, read_image(src, in_channels)

    futures = {}
    file_iter = iter(files)
    prefetch_count = args.io_workers + 1
    for src_path in list(file_iter)[:prefetch_count]:
        futures[io_executor.submit(_read, src_path)] = src_path
    remaining = list(file_iter)

    with tqdm(total=len(files), desc="Super-resolving", ncols=90) as pbar:
        while futures:
            done = next(as_completed(futures))
            del futures[done]
            if remaining:
                nxt = remaining.pop(0)
                futures[io_executor.submit(_read, nxt)] = nxt

            src_path, img = done.result()
            dst_path = output_dir / src_path.relative_to(input_dir)

            if img is None:
                tqdm.write(f"  Warning: could not read {src_path.name} — skipped")
                skipped += 1
                pbar.update(1)
                continue

            # Optional denoising step
            if denoiser is not None:
                img = denoise_tile_inference(
                    denoiser, img,
                    tile_size=args.tile_size, overlap=args.tile_overlap,
                    device=device,
                )

            sr_img = sr_tile_inference(
                net_g, img, scale,
                tile_size=args.tile_size, overlap=args.tile_overlap,
                device=device, tile_batch_size=args.tile_batch_size,
            )
            io_executor.submit(save_sr_image, sr_img, src_path, dst_path, scale)
            ok += 1
            pbar.update(1)

    io_executor.shutdown(wait=True)
    print(f"\nDone.  {ok} image(s) saved → {output_dir}"
          + (f"  ({skipped} skipped)" if skipped else ""))


# ─────────────────────────────────────────────────────────────────────────────
# Preprocess (standalone denoising)
# ─────────────────────────────────────────────────────────────────────────────

def preprocess(args: argparse.Namespace) -> None:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    model_path = Path(args.model)
    if not model_path.exists():
        print(f"Error: model not found at {model_path}", file=sys.stderr)
        sys.exit(1)

    print(f"Loading Restormer denoiser from {model_path}...")
    model = _load_denoiser(model_path, device)
    in_channels = _load_checkpoint(model_path, device).get("in_channels", 1)
    mode_label = "panchromatic" if in_channels == 1 else "RGB"
    print(f"  Mode: {mode_label}  |  tile: {args.tile_size}px")

    input_dir  = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    files = sorted([f for f in input_dir.rglob("*") if is_image(f)])
    if not files:
        print(f"No images found in {input_dir}", file=sys.stderr)
        sys.exit(1)

    print(f"\nDenoising {len(files)} image(s) → {output_dir}\n")
    ok = skipped = 0
    io_executor = ThreadPoolExecutor(max_workers=2)

    def _read(src): return src, read_image(src, in_channels)

    futures = {}
    file_iter = iter(files)
    for src_path in list(file_iter)[:3]:
        futures[io_executor.submit(_read, src_path)] = src_path
    remaining = list(file_iter)

    with tqdm(total=len(files), desc="Denoising", ncols=90) as pbar:
        while futures:
            done = next(as_completed(futures))
            del futures[done]
            if remaining:
                nxt = remaining.pop(0)
                futures[io_executor.submit(_read, nxt)] = nxt

            src_path, img = done.result()
            dst_path = output_dir / src_path.relative_to(input_dir)

            if img is None:
                tqdm.write(f"  Warning: could not read {src_path.name} — skipped")
                skipped += 1
                pbar.update(1)
                continue

            clean = denoise_tile_inference(
                model, img,
                tile_size=args.tile_size, overlap=args.tile_overlap,
                device=device,
            )
            io_executor.submit(save_image, clean, src_path, dst_path)
            ok += 1
            pbar.update(1)

    io_executor.shutdown(wait=True)
    print(f"\nDone.  {ok} image(s) saved → {output_dir}"
          + (f"  ({skipped} skipped)" if skipped else ""))


# ─────────────────────────────────────────────────────────────────────────────
# CLI
# ─────────────────────────────────────────────────────────────────────────────

def _add_arch_args(g):
    """Shared architecture arguments for train subcommand."""
    g.add_argument("--arch", default="rrdb", choices=["rrdb", "hat"],
                   help="Generator architecture. 'hat' (Hybrid Attention Transformer) "
                        "yields +0.3–1.2 dB PSNR over SwinIR at higher compute cost.")
    g.add_argument("--num-feat", type=int, default=64, metavar="N",
                   help="Feature channel width (both RRDB and HAT).")
    g.add_argument("--num-blocks", type=int, default=23, metavar="N",
                   help="RRDB: number of RRDB blocks (23=full, 6=light). "
                        "HAT: HATBlocks per RHAG group (default 6).")
    g.add_argument("--num-groups", type=int, default=6, metavar="N",
                   help="HAT only: number of Residual Hybrid Attention Groups.")
    g.add_argument("--num-heads", type=int, default=4, metavar="N",
                   help="HAT only: attention heads per window (num-feat must be divisible).")
    g.add_argument("--window-size", type=int, default=8, metavar="N",
                   help="HAT only: window size for local self-attention. "
                        "Tile/patch size must be divisible by this value.")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="geo_sr",
        description="Photogrammetric Super-Resolution — RRDB or HAT backbone, "
                    "with optional Restormer denoiser preprocessing for historical aerials.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  # RRDB (v1 compatible, 2× SR):
  python geo_sr.py train --hr-dir ./hr_5cm --scale 2 --model ./models/sr_rrdb.pth

  # HAT (recommended for new training, historical aerials):
  python geo_sr.py train --hr-dir ./hr_pan --scale 2 --arch hat --historical \\
      --model ./models/sr_hat.pth

  # HAT, lighter (fewer groups/depth):
  python geo_sr.py train --hr-dir ./hr_pan --scale 2 --arch hat \\
      --num-groups 4 --num-blocks 4 --model ./models/sr_hat_light.pth

  # Train Restormer denoiser for historical scans:
  python geo_sr.py train-denoiser --hr-dir ./hr_pan --historical \\
      --model ./models/denoiser.pth

  # Apply SR only:
  python geo_sr.py apply --input-dir ./scans --output-dir ./sr_x2 \\
      --model ./models/sr_hat.pth

  # Denoise then SR (recommended for noisy historical aerials):
  python geo_sr.py apply --input-dir ./scans --output-dir ./sr_x2 \\
      --model ./models/sr_hat.pth --preprocess-model ./models/denoiser.pth

  # Standalone denoising only:
  python geo_sr.py preprocess --input-dir ./scans --output-dir ./clean \\
      --model ./models/denoiser.pth
""",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # ── train ──────────────────────────────────────────────────────────────
    tr = sub.add_parser("train", help="Train super-resolution model",
                        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    io_g = tr.add_argument_group("I/O")
    io_g.add_argument("--hr-dir", required=True, metavar="DIR")
    io_g.add_argument("--lr-dir", default=None, metavar="DIR",
                      help="Optional LR dir for tone calibration (no pixel pairing).")
    io_g.add_argument("--model", default="./models/geo_sr.pth", metavar="PATH")
    arch_g = tr.add_argument_group("Architecture")
    _add_arch_args(arch_g)
    arch_g.add_argument("--scale", type=int, default=2, choices=[2, 3, 4])
    train_g = tr.add_argument_group("Training")
    train_g.add_argument("--epochs", type=int, default=200)
    train_g.add_argument("--batch-size", type=int, default=8)
    train_g.add_argument("--patch-size", type=int, default=128, metavar="PX")
    train_g.add_argument("--patches-per-image", type=int, default=30)
    train_g.add_argument("--lr-g", type=float, default=1e-4)
    train_g.add_argument("--lr-d", type=float, default=1e-4)
    train_g.add_argument("--num-workers", type=int, default=0)
    train_g.add_argument("--save-every", type=int, default=10)
    train_g.add_argument("--restart", action="store_true")
    train_g.add_argument("--no-amp", action="store_true")
    loss_g = tr.add_argument_group("Loss weights")
    loss_g.add_argument("--w-l1",   type=float, default=0.01)
    loss_g.add_argument("--w-perc", type=float, default=1.0)
    loss_g.add_argument("--w-adv",  type=float, default=0.1)
    loss_g.add_argument("--no-gan", action="store_true")
    deg_g = tr.add_argument_group("Degradation")
    deg_g.add_argument("--sigma-blur",  type=float, default=1.2)
    deg_g.add_argument("--sigma-noise", type=float, default=0.005)
    deg_g.add_argument("--historical",  action="store_true",
                       help="Add film grain, vignetting, scanner banding, and film-base fog "
                            "to the synthetic degradation — recommended for aerial scans.")

    # ── train-denoiser ─────────────────────────────────────────────────────
    td = sub.add_parser("train-denoiser",
                        help="Train Restormer denoiser for historical aerial preprocessing",
                        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    td.add_argument("--hr-dir", required=True, metavar="DIR")
    td.add_argument("--model", default="./models/denoiser.pth", metavar="PATH")
    td.add_argument("--dim", type=int, default=48,
                    help="Restormer channel width. 48=small, 64=medium.")
    td.add_argument("--epochs", type=int, default=150)
    td.add_argument("--batch-size", type=int, default=4)
    td.add_argument("--patch-size", type=int, default=256, metavar="PX",
                    help="Training patch size (same resolution — no scale factor).")
    td.add_argument("--patches-per-image", type=int, default=30)
    td.add_argument("--lr-g",       type=float, default=2e-4)
    td.add_argument("--sigma-noise", type=float, default=0.02)
    td.add_argument("--num-workers", type=int, default=0)
    td.add_argument("--save-every",  type=int, default=10)
    td.add_argument("--restart", action="store_true")
    td.add_argument("--no-amp",  action="store_true")
    td.add_argument("--historical", action="store_true",
                    help="Use historical film/scan degradation model.")

    # ── apply ──────────────────────────────────────────────────────────────
    ap = sub.add_parser("apply", help="Apply SR model to a dataset",
                        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ap.add_argument("--input-dir",  required=True, metavar="DIR")
    ap.add_argument("--output-dir", required=True, metavar="DIR")
    ap.add_argument("--model",      required=True, metavar="PATH")
    ap.add_argument("--preprocess-model", default=None, metavar="PATH",
                    help="Optional Restormer denoiser checkpoint. "
                         "If provided, each image is denoised before SR. "
                         "Recommended for noisy historical aerial scans.")
    ap.add_argument("--tile-size",       type=int, default=512)
    ap.add_argument("--tile-overlap",    type=int, default=64)
    ap.add_argument("--tile-batch-size", type=int, default=4)
    ap.add_argument("--io-workers",      type=int, default=2)

    # ── preprocess ─────────────────────────────────────────────────────────
    pp = sub.add_parser("preprocess", help="Apply Restormer denoiser standalone",
                        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    pp.add_argument("--input-dir",  required=True, metavar="DIR")
    pp.add_argument("--output-dir", required=True, metavar="DIR")
    pp.add_argument("--model",      required=True, metavar="PATH")
    pp.add_argument("--tile-size",    type=int, default=1024)
    pp.add_argument("--tile-overlap", type=int, default=128)

    return parser


def main() -> None:
    if not HAS_RASTERIO:
        print("Note: rasterio not installed — GeoTIFF georeferencing will NOT be preserved.\n"
              "      Install with: pip install rasterio\n")
    parser = build_parser()
    args = parser.parse_args()
    if args.command == "train":
        train(args)
    elif args.command == "train-denoiser":
        train_denoiser(args)
    elif args.command == "apply":
        apply_sr(args)
    elif args.command == "preprocess":
        preprocess(args)


if __name__ == "__main__":
    main()