# train_ddp_fno_encoderfno_ddpsafe_split_save_all_fields.py
# =========================================================
# ✅ DDP-safe split + training
# ✅ After training:
#    - Load best model
#    - Infer ALL train/test samples (DDP)
#    - Save per-sample errors + ALL E_pred/E_true + coords to MAT
#    - Also save an example sample MAT + PNG plot (optional)
#
# ✅ Adapted for NEW dataset: metalens_dataset.mat (NO standardization / NO denorm)
#    - Rimg:     (N,1,512,512)
#    - coords:   (M,2) in [-1,1]
#    - Ere_flat: (N,M) normalized
#    - Eim_flat: (N,M) normalized
# =========================================================

import os
import math
import datetime
import numpy as np
import torch
from tqdm.auto import tqdm
from scipy.io import savemat
import torch.nn as nn
from torch.utils.data import DataLoader, Subset, ConcatDataset
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
from pathlib import Path

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# Existing API (must be updated to read metalens_dataset.mat)
from getdata import build_datasets


# =========================================================
# DDP helpers
# =========================================================
def ddp_setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    return device, local_rank


def ddp_cleanup():
    if dist.is_available() and dist.is_initialized():
        dist.destroy_process_group()


def is_main_process():
    return (not dist.is_available()) or (not dist.is_initialized()) or dist.get_rank() == 0


def world_size():
    return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1


def broadcast_tensor_(t: torch.Tensor, src: int = 0):
    if world_size() > 1:
        dist.broadcast(t, src=src)
    return t


def all_gather_object(obj):
    """Gather python objects to every rank."""
    if world_size() == 1:
        return [obj]
    out = [None for _ in range(world_size())]
    dist.all_gather_object(out, obj)
    return out


# =========================================================
# Norm
# =========================================================
def make_gn(num_channels: int, max_groups: int = 8) -> nn.GroupNorm:
    for g in [max_groups, 4, 2, 1]:
        if num_channels % g == 0:
            return nn.GroupNorm(num_groups=g, num_channels=num_channels)
    return nn.GroupNorm(num_groups=1, num_channels=num_channels)


# =========================================================
# Loss: Relative L2 (on flattened M points)
# =========================================================
class RelL2Loss(nn.Module):
    def __init__(self, eps=1e-12):
        super().__init__()
        self.eps = eps

    def forward(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
        # pred/true: (B,M,2)
        num = torch.sum((pred - true) ** 2, dim=(1, 2))
        den = torch.sum(true ** 2, dim=(1, 2)) + self.eps
        return torch.mean(num / den)


def rel_l2_per_sample(pred: torch.Tensor, true: torch.Tensor, eps=1e-12) -> torch.Tensor:
    num = torch.sum((pred - true) ** 2, dim=(1, 2))
    den = torch.sum(true ** 2, dim=(1, 2)) + eps
    return num / den  # (B,)


# =========================================================
# FNO core (self-contained) -- FIXED FFTConv2d (modes clamp)
# =========================================================
class FFTConv2d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, modes: int):
        super().__init__()
        self.out_channels = out_channels
        self.modes = int(modes)
        self.w = nn.Parameter(
            torch.rand(in_channels, out_channels, 2 * self.modes, self.modes, 2)
            / (in_channels * out_channels)
        )

    @staticmethod
    def cmul2d(a, b):
        return torch.einsum("bixy,ioxy->boxy", a, b)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        device = x.device

        x_ft = torch.fft.rfft2(x)  # (B,C,H,Wf)
        Wf = x_ft.size(-1)

        my = min(self.modes, H // 2)
        mx = min(self.modes, Wf)

        out_ft = torch.zeros(B, self.out_channels, H, Wf, dtype=torch.cfloat, device=device)

        x_ft = torch.fft.fftshift(x_ft, dim=-2)
        cy = H // 2

        a = x_ft[..., cy - my: cy + my, :mx]  # (B,inC,2*my,mx)
        w_c = torch.view_as_complex(self.w)   # (inC,outC,2*modes,modes)
        b = w_c[..., self.modes - my: self.modes + my, :mx]
        out_ft[..., cy - my: cy + my, :mx] = self.cmul2d(a, b)

        out_ft = torch.fft.ifftshift(out_ft, dim=-2)
        x = torch.fft.irfft2(out_ft, s=(H, W))
        return x


class FNOBlock2d(nn.Module):
    def __init__(self, modes: int, width: int, norm: str = "gn", dropout: float = 0.10, layerscale=0.10):
        super().__init__()
        self.fftconv = FFTConv2d(width, width, modes)
        self.conv = nn.Conv2d(width, width, 1, bias=False)

        self.norm = nn.BatchNorm2d(width) if norm == "bn" else make_gn(width)
        self.act = nn.GELU()
        self.drop = nn.Dropout2d(p=float(dropout)) if dropout and dropout > 0 else nn.Identity()
        self.gamma = nn.Parameter(layerscale * torch.ones(1, width, 1, 1)) if layerscale and layerscale > 0 else None

    def forward(self, x):
        y = self.fftconv(x) + self.conv(x)
        x = x + (self.gamma * y if self.gamma is not None else y)
        x = self.norm(x)
        x = self.act(x)
        x = self.drop(x)
        return x


class FNOModel2d(nn.Module):
    def __init__(
        self,
        modes=6,
        width=24,
        blocks=4,
        padding=0,
        in_channels=3,
        out_channels=2,
        norm="gn",
        dropout=0.10,
        layerscale=0.10,
    ):
        super().__init__()
        self.conv_in = nn.Conv2d(in_channels, width, 1, bias=True)

        self.padding = int(padding)
        self.pad_in = nn.ConstantPad2d(self.padding, 0.0) if self.padding > 0 else nn.Identity()
        self.pad_out = nn.ConstantPad2d(-self.padding, 0.0) if self.padding > 0 else nn.Identity()

        self.fno_blocks = nn.Sequential(
            *[FNOBlock2d(modes, width, norm=norm, dropout=dropout, layerscale=layerscale) for _ in range(blocks)]
        )

        self.conv_out = nn.Sequential(
            nn.Conv2d(width, width, 1),
            nn.GELU(),
            nn.Dropout2d(dropout),
            nn.Conv2d(width, out_channels, 1),
        )

    def forward(self, x):
        x = self.conv_in(x)
        x = self.pad_in(x)
        x = self.fno_blocks(x)
        x = self.pad_out(x)
        x = self.conv_out(x)
        return x


# =========================================================
# Anti-alias Geometry Encoder
# =========================================================
class GeomEncoder(nn.Module):
    def __init__(self, out_ch=8, drop=0.10):
        super().__init__()
        self.prepool = nn.AvgPool2d(kernel_size=6, stride=6)  # 512 -> ~85

        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1, bias=False),
            make_gn(16), nn.GELU(), nn.Dropout2d(drop),

            nn.Conv2d(16, 24, 3, stride=2, padding=1, bias=False),  # ~85 -> ~43
            make_gn(24), nn.GELU(), nn.Dropout2d(drop),

            nn.Conv2d(24, 32, 3, stride=2, padding=1, bias=False),  # ~43 -> ~22
            make_gn(32), nn.GELU(), nn.Dropout2d(drop),

            nn.Conv2d(32, out_ch, 1, bias=True),
            make_gn(out_ch), nn.GELU(),
        )
        self.pool_to = None

    def forward(self, x_img, Ny, Nx):
        h = self.prepool(x_img)
        h = self.net(h)
        if (self.pool_to is None) or (self.pool_to.output_size != (Ny, Nx)):
            self.pool_to = nn.AdaptiveAvgPool2d((Ny, Nx))
        return self.pool_to(h)


# =========================================================
# Optional Fourier features (DEFAULT OFF)
# =========================================================
class FourierFeatures2D(nn.Module):
    def __init__(self, K=2, scale=1.0):
        super().__init__()
        self.K = int(K)
        self.scale = float(scale)
        freqs = torch.linspace(1.0, self.K, steps=self.K) * self.scale
        self.register_buffer("freqs", freqs, persistent=False)

    def forward(self, xy):
        x = xy[:, 0:1]
        y = xy[:, 1:2]
        f = self.freqs.view(1, self.K, 1, 1)
        x_proj = 2 * math.pi * x * f
        y_proj = 2 * math.pi * y * f
        return torch.cat(
            [torch.sin(x_proj), torch.cos(x_proj), torch.sin(y_proj), torch.cos(y_proj)],
            dim=1,
        )


# =========================================================
# Regressor: Encoder + (xy [+ optional Fourier]) + FNO
# =========================================================
class FNOFieldRegressor_EncoderFNO(nn.Module):
    def __init__(
        self,
        modes=6,
        width=24,
        blocks=4,
        padding=0,
        norm="gn",
        dropout=0.10,
        layerscale=0.10,
        geom_ch=8,
        use_fourier=False,
        fourier_K=2,
        fourier_scale=1.0,
        dropout_xy=0.05,
    ):
        super().__init__()
        self.encoder = GeomEncoder(out_ch=geom_ch, drop=0.10)

        self.use_fourier = bool(use_fourier)
        self.fourier = FourierFeatures2D(K=fourier_K, scale=fourier_scale) if self.use_fourier else None
        self.drop_xy = nn.Dropout2d(dropout_xy) if dropout_xy and dropout_xy > 0 else nn.Identity()

        in_ch = geom_ch + 2 + (4 * fourier_K if self.use_fourier else 0)

        self.fno = FNOModel2d(
            modes=modes,
            width=width,
            blocks=blocks,
            padding=padding,
            in_channels=in_ch,
            out_channels=2,
            norm=norm,
            dropout=dropout,
            layerscale=layerscale,
        )

        self._cached_key = None
        self._cached_xy = None  # (1,2,Ny,Nx)

    @staticmethod
    def _sanitize_coords(coords: torch.Tensor) -> torch.Tensor:
        # accept (B,M,2) or (B,1,M,2)
        if coords.dim() == 4:
            coords = coords[:, 0, :, :]
        return coords

    @staticmethod
    def _coords_to_xy_grid(coords: torch.Tensor):
        """
        coords: (B,M,2) or (M,2)
        Build xy grid (1,2,Ny,Nx) using unique sorted x and y.
        """
        coords0 = coords[0] if coords.dim() == 3 else coords  # (M,2)
        x = coords0[:, 0]
        y = coords0[:, 1]

        xs = torch.unique(x)
        ys = torch.unique(y)
        xs, _ = torch.sort(xs)
        ys, _ = torch.sort(ys)

        Ny, Nx = ys.numel(), xs.numel()
        yy, xx = torch.meshgrid(ys, xs, indexing="ij")
        xy = torch.stack([xx, yy], dim=0).unsqueeze(0)  # (1,2,Ny,Nx)
        return xy, int(Ny), int(Nx)

    def forward(self, x_img: torch.Tensor, coords: torch.Tensor):
        B = x_img.shape[0]
        device = x_img.device
        dtype = x_img.dtype

        coords = self._sanitize_coords(coords).to(device=device)
        xy, Ny, Nx = self._coords_to_xy_grid(coords)

        key = (Ny, Nx, device.type, device.index, str(dtype))
        if self._cached_key != key:
            self._cached_key = key
            self._cached_xy = xy.to(device=device, dtype=dtype)

        xy_b = self._cached_xy.expand(B, -1, -1, -1)
        xy_b = self.drop_xy(xy_b)

        geom_feat = self.encoder(x_img, Ny, Nx)

        if self.use_fourier:
            ff = self.fourier(xy_b)
            inp = torch.cat([geom_feat, xy_b, ff], dim=1)
        else:
            inp = torch.cat([geom_feat, xy_b], dim=1)

        out = self.fno(inp)  # (B,2,Ny,Nx)
        out = out.permute(0, 2, 3, 1).contiguous().view(B, Ny * Nx, 2)
        return out


# =========================================================
# Dataset loader: prefer build_dataset_full; else fallback concat
# =========================================================
def load_full_dataset_and_stats(matpath: str):
    """
    For the new dataset we do NOT need norm_stats.
    If build_dataset_full exists, it should return (full_ds, {}).
    Otherwise we concat train/test from build_datasets.
    """
    try:
        from getdata import build_dataset_full  # optional
        full_ds, norm_stats = build_dataset_full(matpath)
        if is_main_process():
            print("[Dataset] Using build_dataset_full().")
        return full_ds, norm_stats if norm_stats is not None else {}
    except Exception as e:
        if is_main_process():
            print("[Dataset INFO] build_dataset_full() not found or failed. Using build_datasets() and concat.")
            print(f"              Reason: {repr(e)}")
        train_ds, test_ds, _ = build_datasets(matpath)
        full_ds = ConcatDataset([train_ds, test_ds])
        return full_ds, {}


# =========================================================
# Indexed dataset wrapper (for per-sample saving)
# =========================================================
class IndexedDataset(torch.utils.data.Dataset):
    """
    returns: (i, x_img, coords, E_true)
    i is index within this ds (0..len(ds)-1)
    """
    def __init__(self, ds):
        self.ds = ds

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

    def __getitem__(self, i):
        x_img, coords, E_true = self.ds[i]
        return i, x_img, coords, E_true


# =========================================================
# Trainer
# =========================================================
class FNOOnly_DDP_Trainer:
    def __init__(
        self,
        model,
        device,
        batch_size=32,
        matpath="metalens_dataset.mat",
        lr=2e-4,
        weight_decay=3e-4,
        num_workers=4,
        pin_memory=True,
        warmup_epochs=5,
        min_lr_ratio=1e-3,
        early_stop_patience=30,
        grad_clip=1.0,
        split_seed=20260303,
        test_ratio=0.2,
    ):
        self.device = device
        self.matpath = matpath
        self.warmup_epochs = int(warmup_epochs)
        self.min_lr_ratio = float(min_lr_ratio)
        self.early_stop_patience = int(early_stop_patience)
        self.grad_clip = float(grad_clip)
        self.split_seed = int(split_seed)
        self.test_ratio = float(test_ratio)

        self.model = model.to(self.device)
        self.ddp_model = DDP(
            self.model,
            device_ids=[self.device.index],
            output_device=self.device.index,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )

        self.loss_fn = RelL2Loss()
        self.optimizer = torch.optim.AdamW(self.ddp_model.parameters(), lr=lr, weight_decay=weight_decay)

        self.scheduler = None
        self.total_epochs = None

        self.file_name = "fno_ddpsafe_final.pth"
        self.best_name = "fno_ddpsafe_best.pth"

        self.losses = []

        self.train_set, self.test_set, self.norm_stats = self.load_dataset_ddpsafe()

        self.train_sampler = DistributedSampler(self.train_set, shuffle=True, drop_last=False)
        self.test_sampler = DistributedSampler(self.test_set, shuffle=False, drop_last=False)

        self.batch_size = batch_size
        self.num_workers = num_workers
        self.pin_memory = pin_memory

        self.train_loader = DataLoader(
            self.train_set,
            batch_size=batch_size,
            sampler=self.train_sampler,
            shuffle=False,
            drop_last=False,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=(num_workers > 0),
        )

        self.test_loader = DataLoader(
            self.test_set,
            batch_size=batch_size,
            sampler=self.test_sampler,
            shuffle=False,
            drop_last=False,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=(num_workers > 0),
        )

        if is_main_process():
            Path("results").mkdir(parents=True, exist_ok=True)

    def load_dataset_ddpsafe(self):
        full_ds, norm_stats = load_full_dataset_and_stats(self.matpath)
        N = len(full_ds)

        n_test = int(round(N * self.test_ratio))
        n_test = max(1, min(N - 1, n_test))

        if is_main_process():
            g = torch.Generator()
            g.manual_seed(self.split_seed)
            perm = torch.randperm(N, generator=g, dtype=torch.int64)
        else:
            perm = torch.empty(N, dtype=torch.int64)

        perm = perm.to(self.device)
        broadcast_tensor_(perm, src=0)
        perm = perm.cpu()

        test_idx = perm[:n_test].tolist()
        train_idx = perm[n_test:].tolist()

        train_ds = Subset(full_ds, train_idx)
        test_ds = Subset(full_ds, test_idx)

        if is_main_process():
            print(
                f"[DDP Split] Full={N}, Train={len(train_ds)}, Test={len(test_ds)}, "
                f"seed={self.split_seed}, test_ratio={self.test_ratio}"
            )

        return train_ds, test_ds, norm_stats

    def load_model(self, path):
        try:
            state = torch.load(path, map_location="cpu", weights_only=True)
        except TypeError:
            state = torch.load(path, map_location="cpu")
        self.model.load_state_dict(state, strict=True)

    @staticmethod
    def _to_img_tensor(x_img: torch.Tensor) -> torch.Tensor:
        # dataset gives (1,512,512); dataloader gives (B,1,512,512) already,
        # but keep robust.
        if x_img.dim() == 2:
            x_img = x_img.unsqueeze(0).unsqueeze(0)
        elif x_img.dim() == 3:
            x_img = x_img.unsqueeze(0)
            if x_img.size(1) != 1:
                x_img = x_img[:, :1]
        elif x_img.dim() == 4:
            if x_img.size(1) != 1:
                x_img = x_img[:, :1]
        else:
            raise ValueError(f"Unexpected x_img dim: {x_img.shape}")
        return x_img

    @staticmethod
    def _sanitize_coords(coords: torch.Tensor) -> torch.Tensor:
        # coords may become (B,M,2) from dataloader stacking; ok.
        if coords.dim() == 4:
            coords = coords[:, 0, :, :]
        return coords

    @staticmethod
    def _sanitize_field(E_true: torch.Tensor) -> torch.Tensor:
        # Want (B,M,2)
        if E_true.dim() == 2:
            E_true = E_true.unsqueeze(0)
        if E_true.dim() == 4:
            E_true = E_true[:, 0, :, :]
        return E_true

    def _build_scheduler(self):
        assert self.total_epochs is not None

        def lr_lambda(epoch: int):
            if epoch < self.warmup_epochs:
                return float(epoch + 1) / float(self.warmup_epochs)
            t = (epoch - self.warmup_epochs) / max(1, (self.total_epochs - self.warmup_epochs))
            return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * 0.5 * (1.0 + np.cos(np.pi * t))

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lr_lambda)

    def get_data_loss(self, x_img, coords, E_true):
        E_pred = self.ddp_model(x_img, coords)
        return self.loss_fn(E_pred, E_true)

    @torch.no_grad()
    def test_loss(self):
        self.ddp_model.eval()
        loss_sum = torch.zeros((), device=self.device)
        count = torch.zeros((), device=self.device)

        for x_img, coords, E_true in self.test_loader:
            x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
            coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
            E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)

            loss = self.get_data_loss(x_img, coords, E_true)
            bsz = x_img.shape[0]
            loss_sum += loss.detach() * bsz
            count += bsz

        if world_size() > 1:
            dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
            dist.all_reduce(count, op=dist.ReduceOp.SUM)

        avg = (loss_sum / (count + 1e-12)).detach()
        return float(avg.item())

    def train(self, epochs=400):
        self.total_epochs = int(epochs)
        self._build_scheduler()

        start_time = datetime.datetime.now()
        best_test_loss = float("inf")
        bad = 0

        epoch_iter = range(self.total_epochs)
        if is_main_process():
            epoch_iter = tqdm(epoch_iter, desc="Training (DDP)")

        for epoch in epoch_iter:
            self.train_sampler.set_epoch(epoch)
            self.ddp_model.train()

            loss_sum = torch.zeros((), device=self.device)
            count = torch.zeros((), device=self.device)

            for x_img, coords, E_true in self.train_loader:
                x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
                coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
                E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)

                self.optimizer.zero_grad(set_to_none=True)
                loss = self.get_data_loss(x_img, coords, E_true)
                loss.backward()

                if self.grad_clip and self.grad_clip > 0:
                    torch.nn.utils.clip_grad_norm_(self.ddp_model.parameters(), self.grad_clip)

                self.optimizer.step()

                bsz = x_img.shape[0]
                loss_sum += loss.detach() * bsz
                count += bsz

            if world_size() > 1:
                dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
                dist.all_reduce(count, op=dist.ReduceOp.SUM)

            avg_train_loss = float((loss_sum / (count + 1e-12)).item())
            avg_test_loss = self.test_loss()
            self.scheduler.step()

            stop_flag = torch.zeros((), device=self.device, dtype=torch.int32)

            if is_main_process():
                self.losses.append([epoch, avg_train_loss, avg_test_loss])
                lr = self.optimizer.param_groups[0]["lr"]

                improved = avg_test_loss < best_test_loss
                if improved:
                    best_test_loss = avg_test_loss
                    bad = 0
                    torch.save(self.model.state_dict(), self.best_name)
                else:
                    bad += 1

                print(
                    f"Epoch {epoch} | train_loss={avg_train_loss:.6e} | "
                    f"test_loss={avg_test_loss:.6e} | best={best_test_loss:.6e} | "
                    f"bad={bad}/{self.early_stop_patience} | lr={lr:.3e}"
                )

                if bad >= self.early_stop_patience:
                    print(f"[Early stop] epoch={epoch}, best_test_loss={best_test_loss:.6e}")
                    stop_flag.fill_(1)

            if world_size() > 1:
                dist.broadcast(stop_flag, src=0)

            if int(stop_flag.item()) == 1:
                break

        if world_size() > 1:
            dist.barrier()

        if is_main_process():
            torch.save(self.model.state_dict(), self.file_name)
            print("Best model saved to:", self.best_name)
            print("Final model saved to:", self.file_name)
            print("Training Time:", (datetime.datetime.now() - start_time).total_seconds(), "s")

    # =========================================================
    # Save one example (NO denorm for new dataset)
    # =========================================================
    def save_example_pred_true_mat_and_plot(
        self,
        idx: int = 0,
        use_best: bool = True,
        mat_save: str = "results/fno_ddpsafe_pred_true_example.mat",
        fig_save: str = "results/fno_ddpsafe_field_plot_example.png",
        interp_method: str = "cubic",
    ):
        if not is_main_process():
            return

        if use_best and Path(self.best_name).exists():
            self.load_model(self.best_name)

        self.ddp_model.eval()

        x_img, coords, E_true = self.test_set[idx]

        if torch.is_tensor(coords) and coords.dim() == 3:
            coords = coords[0, :]
        coords_np = coords.detach().cpu().numpy()

        if torch.is_tensor(E_true) and E_true.dim() == 3:
            E_true = E_true[0, :, :]
        E_true_np = E_true.detach().cpu().numpy()[None, ...]  # (1,M,2)

        with torch.no_grad():
            x_img_t = self._to_img_tensor(x_img).to(self.device)
            coords_t = self._sanitize_coords(coords.unsqueeze(0) if coords.dim() == 2 else coords).to(self.device)
            E_pred_np = self.ddp_model(x_img_t, coords_t).detach().cpu().numpy()  # (1,M,2)

        Ere_pred, Eim_pred = E_pred_np[:, :, 0], E_pred_np[:, :, 1]
        Ere_true, Eim_true = E_true_np[:, :, 0], E_true_np[:, :, 1]

        savemat(
            mat_save,
            {
                "Ere_pred": Ere_pred.astype(np.float32),
                "Eim_pred": Eim_pred.astype(np.float32),
                "Ere_true": Ere_true.astype(np.float32),
                "Eim_true": Eim_true.astype(np.float32),
                "coords_xy": coords_np.astype(np.float32),
            },
        )
        print("Saved example mat:", mat_save)

        # plot
        X = coords_np[:, 0]
        Y = coords_np[:, 1]
        Xi, Yi = np.meshgrid(np.unique(X), np.unique(Y))

        def interp(z):
            Zi = griddata(points=(X, Y), values=z, xi=(Xi, Yi), method=interp_method)
            from scipy.ndimage import gaussian_filter
            return gaussian_filter(Zi, sigma=0.8)

        re_err = np.abs(Ere_pred[0] - Ere_true[0])
        im_err = np.abs(Eim_pred[0] - Eim_true[0])

        fields = [
            (interp(Ere_pred[0]), "E_real_pred"),
            (interp(Ere_true[0]), "E_real_true"),
            (interp(re_err),      "E_real_error"),
            (interp(Eim_pred[0]), "E_imag_pred"),
            (interp(Eim_true[0]), "E_imag_true"),
            (interp(im_err),      "E_imag_error"),
        ]

        fig, axes = plt.subplots(2, 3, figsize=(18, 9), sharex=True, sharey=True)
        for ax, (Zi, title) in zip(axes.ravel(), fields):
            pcm = ax.pcolormesh(Xi, Yi, Zi, shading="gouraud")
            ax.set_title(title)
            ax.set_xlabel("X")
            ax.set_ylabel("Y")
            ax.set_aspect("equal")
            fig.colorbar(pcm, ax=ax, shrink=0.8)

        plt.tight_layout()
        plt.savefig(fig_save, dpi=200)
        plt.close(fig)
        print("Saved example plot:", fig_save)

    # =========================================================
    # Infer ALL train/test, save errors + ALL fields + coords
    # =========================================================
    @torch.no_grad()
    def infer_all_and_save_errors_and_fields(
        self,
        use_best: bool = True,
        save_path: str = "results/fno_ddpsafe_all_errors_and_fields.mat",
        compute_mse_mae: bool = True,
        save_coords_per_sample: bool = False,
        dtype_save=np.float32,
    ):
        if use_best and Path(self.best_name).exists():
            if is_main_process():
                print(f"[Infer] Loading best model: {self.best_name}")
            self.load_model(self.best_name)

        self.ddp_model.eval()

        common_coords = None
        if is_main_process():
            _, coords0, _ = self.train_set[0]
            if torch.is_tensor(coords0) and coords0.dim() == 3:
                coords0 = coords0[0, :]
            common_coords = coords0.detach().cpu().numpy().astype(dtype_save)

        def build_infer_loader(ds):
            ds_idx = IndexedDataset(ds)
            sampler = DistributedSampler(ds_idx, shuffle=False, drop_last=False)  # will pad -> we dedup by idx later
            loader = DataLoader(
                ds_idx,
                batch_size=self.batch_size,
                sampler=sampler,
                shuffle=False,
                drop_last=False,
                num_workers=self.num_workers,
                pin_memory=self.pin_memory,
                persistent_workers=(self.num_workers > 0),
            )
            return loader, len(ds_idx)

        def run_one_split(tag: str, ds):
            loader, N = build_infer_loader(ds)

            local_idx, local_rel = [], []
            local_mse, local_mae = [], []
            local_pred, local_true = [], []
            local_coords = []

            for idx, x_img, coords, E_true in loader:
                x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
                coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
                E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)

                E_pred = self.ddp_model(x_img, coords)  # (B,M,2)

                rel = rel_l2_per_sample(E_pred, E_true).detach().cpu().numpy()

                if compute_mse_mae:
                    diff = (E_pred - E_true).detach()
                    mse = torch.mean(diff * diff, dim=(1, 2)).cpu().numpy()
                    mae = torch.mean(torch.abs(diff), dim=(1, 2)).cpu().numpy()

                local_idx.append(idx.cpu().numpy().astype(np.int64))
                local_rel.append(rel.astype(np.float64))
                if compute_mse_mae:
                    local_mse.append(mse.astype(np.float64))
                    local_mae.append(mae.astype(np.float64))

                local_pred.append(E_pred.detach().cpu().numpy().astype(dtype_save))
                local_true.append(E_true.detach().cpu().numpy().astype(dtype_save))

                if save_coords_per_sample:
                    local_coords.append(coords.detach().cpu().numpy().astype(dtype_save))

            if len(local_idx) == 0:
                pack = {
                    "idx": np.zeros((0,), np.int64),
                    "rel": np.zeros((0,), np.float64),
                    "mse": np.zeros((0,), np.float64) if compute_mse_mae else None,
                    "mae": np.zeros((0,), np.float64) if compute_mse_mae else None,
                    "pred": np.zeros((0, 0, 2), dtype_save),
                    "true": np.zeros((0, 0, 2), dtype_save),
                    "coords": np.zeros((0, 0, 2), dtype_save) if save_coords_per_sample else None,
                    "N": N,
                }
            else:
                pack = {
                    "idx": np.concatenate(local_idx, axis=0),
                    "rel": np.concatenate(local_rel, axis=0),
                    "mse": np.concatenate(local_mse, axis=0) if compute_mse_mae else None,
                    "mae": np.concatenate(local_mae, axis=0) if compute_mse_mae else None,
                    "pred": np.concatenate(local_pred, axis=0),
                    "true": np.concatenate(local_true, axis=0),
                    "coords": np.concatenate(local_coords, axis=0) if save_coords_per_sample else None,
                    "N": N,
                }

            packs = all_gather_object(pack)
            if not is_main_process():
                return None

            idx_all = np.concatenate([p["idx"] for p in packs], axis=0)
            rel_all = np.concatenate([p["rel"] for p in packs], axis=0)
            pred_all = np.concatenate([p["pred"] for p in packs], axis=0)
            true_all = np.concatenate([p["true"] for p in packs], axis=0)

            if compute_mse_mae:
                mse_all = np.concatenate([p["mse"] for p in packs], axis=0)
                mae_all = np.concatenate([p["mae"] for p in packs], axis=0)
            else:
                mse_all, mae_all = None, None

            if save_coords_per_sample:
                coords_all = np.concatenate([p["coords"] for p in packs if p["coords"] is not None], axis=0)
            else:
                coords_all = None

            order = np.argsort(idx_all)
            idx_all = idx_all[order]
            rel_all = rel_all[order]
            pred_all = pred_all[order]
            true_all = true_all[order]
            if compute_mse_mae:
                mse_all = mse_all[order]
                mae_all = mae_all[order]
            if save_coords_per_sample:
                coords_all = coords_all[order]

            # dedup padded indices from DistributedSampler
            _, first_pos = np.unique(idx_all, return_index=True)
            idx_all = idx_all[first_pos]
            rel_all = rel_all[first_pos]
            pred_all = pred_all[first_pos]
            true_all = true_all[first_pos]
            if compute_mse_mae:
                mse_all = mse_all[first_pos]
                mae_all = mae_all[first_pos]
            if save_coords_per_sample:
                coords_all = coords_all[first_pos]

            if idx_all.shape[0] != N:
                missing = set(range(N)) - set(idx_all.tolist())
                raise RuntimeError(
                    f"[{tag}] after dedup: got {idx_all.shape[0]} unique != N={N}, missing={len(missing)}"
                )

            out = {
                f"{tag}_idx": idx_all.astype(np.int64),
                f"{tag}_relL2": rel_all.astype(np.float64),
                f"{tag}_E_pred": pred_all.astype(dtype_save),
                f"{tag}_E_true": true_all.astype(dtype_save),
                f"{tag}_relL2_mean": float(np.mean(rel_all)),
                f"{tag}_relL2_std": float(np.std(rel_all)),
            }
            if compute_mse_mae:
                out[f"{tag}_mse"] = mse_all.astype(np.float64)
                out[f"{tag}_mae"] = mae_all.astype(np.float64)
                out[f"{tag}_mse_mean"] = float(np.mean(mse_all))
                out[f"{tag}_mae_mean"] = float(np.mean(mae_all))
            if save_coords_per_sample:
                out[f"{tag}_coords_xy_all"] = coords_all.astype(dtype_save)

            return out

        train_out = run_one_split("train", self.train_set)
        test_out = run_one_split("test", self.test_set)

        if is_main_process():
            mdict = {}
            mdict.update(train_out)
            mdict.update(test_out)

            if common_coords is not None:
                mdict["coords_xy"] = common_coords  # (M,2)

            mdict["meta_matpath"] = self.matpath
            mdict["meta_best_ckpt"] = self.best_name if use_best else ""
            mdict["meta_world_size"] = int(world_size())
            mdict["meta_compute_mse_mae"] = int(compute_mse_mae)
            mdict["meta_save_coords_per_sample"] = int(save_coords_per_sample)
            mdict["meta_dtype_save"] = str(dtype_save)

            savemat(save_path, mdict)
            print(f"[Infer] Saved ALL errors + ALL fields to: {save_path}")
            print(f"[Infer] train_relL2_mean={mdict['train_relL2_mean']:.6e} | test_relL2_mean={mdict['test_relL2_mean']:.6e}")


# =========================================================
# Main
# =========================================================
def main():
    device, local_rank = ddp_setup()
    torch.backends.cudnn.benchmark = True

    try:
        # -------- config --------
        epochs = 1000
        batch_size = 32

        # model
        modes = 6
        width = 24
        blocks = 4
        padding = 0
        dropout = 0.10
        layerscale = 0.10
        geom_ch = 8

        use_fourier = False
        fourier_K = 2
        fourier_scale = 1.0
        dropout_xy = 0.05

        # optim
        lr = 2e-4
        weight_decay = 3e-4

        trainer = FNOOnly_DDP_Trainer(
            model=FNOFieldRegressor_EncoderFNO(
                modes=modes,
                width=width,
                blocks=blocks,
                padding=padding,
                norm="gn",
                dropout=dropout,
                layerscale=layerscale,
                geom_ch=geom_ch,
                use_fourier=use_fourier,
                fourier_K=fourier_K,
                fourier_scale=fourier_scale,
                dropout_xy=dropout_xy,
            ),
            device=device,
            batch_size=batch_size,
            matpath="metalens_dataset.mat",   # ✅ NEW DATASET
            lr=lr,
            weight_decay=weight_decay,
            num_workers=4,
            pin_memory=True,
            warmup_epochs=5,
            min_lr_ratio=1e-3,
            early_stop_patience=30,
            grad_clip=1.0,
            split_seed=20260303,
            test_ratio=0.2,
        )

        trainer.train(epochs=epochs)

        # ✅ 1) Save one example (NO denorm)
        trainer.save_example_pred_true_mat_and_plot(
            idx=0,
            use_best=True,
            mat_save="results/fno_ddpsafe_pred_true_example.mat",
            fig_save="results/fno_ddpsafe_field_plot_example.png",
        )

        # ✅ 2) Save ALL train/test: errors + E_pred/E_true (+ coords)
        trainer.infer_all_and_save_errors_and_fields(
            use_best=True,
            save_path="results/fno_ddpsafe_all_errors_and_fields.mat",
            compute_mse_mae=True,
            save_coords_per_sample=False,  # coords are shared
            dtype_save=np.float32,
        )

    finally:
        ddp_cleanup()


if __name__ == "__main__":
    main()