# infer_ddp_save_all_dataset_with_inputs.py
# =========================================================
# ✅ DDP / single-GPU inference for FULL dataset
# ✅ Save per-sample: global_idx, Rimg(input), coords, E_pred, E_true, errors
# ✅ DDP-safe: each rank writes shard -> rank0 merges, sorts by idx, dedups
# Output: results/full_infer_with_inputs.h5
# =========================================================

import os
import shutil
import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, ConcatDataset
from torch.utils.data.distributed import DistributedSampler
from pathlib import Path
import h5py

# ---------------------------------------------------------
# Import from your training script (must be in same folder or PYTHONPATH)
# It must have "if __name__ == '__main__': main()" guard (you already do).
# ---------------------------------------------------------
from PI_code_test_FNO import (
    ddp_setup,
    ddp_cleanup,
    is_main_process,
    world_size,
    broadcast_tensor_,
    load_full_dataset_and_stats,
    IndexedDataset,
    rel_l2_per_sample,
    FNOFieldRegressor_EncoderFNO,
)


# ---------------------------------------------------------
# Small helpers (same behavior as trainer)
# ---------------------------------------------------------
def _to_img_tensor(x_img: torch.Tensor) -> torch.Tensor:
    # want (B,1,512,512)
    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: {tuple(x_img.shape)}")
    return x_img


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


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 load_model_strict(model: torch.nn.Module, ckpt_path: str):
    ckpt_path = str(ckpt_path)
    try:
        state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    except TypeError:
        state = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(state, strict=True)


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

    try:
        # =========================
        # Config (edit as needed)
        # =========================
        matpath = "metalens_dataset.mat"

        # ckpt
        ckpt_path = "fno_ddpsafe_best.pth"   # or "fno_ddpsafe_final.pth"

        # dataloader
        batch_size = 16     # inference: you can increase if memory allows
        num_workers = 4
        pin_memory = True

        # output
        out_dir = Path("results")
        out_dir.mkdir(parents=True, exist_ok=True)
        shards_dir = out_dir / "infer_shards"
        final_h5 = out_dir / "full_infer_with_inputs.h5"

        # =========================
        # Build FULL dataset
        # =========================
        full_ds, _ = load_full_dataset_and_stats(matpath)
        N = len(full_ds)

        if is_main_process():
            print(f"[Infer] FULL dataset size: N={N}")

        # =========================
        # Build model (must match training)
        # =========================
        model = FNOFieldRegressor_EncoderFNO(
            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,
        ).to(device)

        # DDP wrap
        ddp_model = torch.nn.parallel.DistributedDataParallel(
            model,
            device_ids=[device.index],
            output_device=device.index,
            find_unused_parameters=False,
            broadcast_buffers=False,
        )

        if is_main_process():
            print(f"[Infer] Loading ckpt: {ckpt_path}")
        load_model_strict(model, ckpt_path)
        ddp_model.eval()

        # =========================
        # coords (shared) save once (rank0)
        # =========================
        coords_shared = None
        if is_main_process():
            x0, coords0, E0 = full_ds[0]
            if torch.is_tensor(coords0) and coords0.dim() == 3:
                coords0 = coords0[0, :]
            coords_shared = coords0.detach().cpu().numpy().astype(np.float32)

        # =========================
        # DataLoader with indices
        # =========================
        ds_idx = IndexedDataset(full_ds)  # returns: (i, x_img, coords, E_true)
        sampler = DistributedSampler(ds_idx, shuffle=False, drop_last=False)
        loader = DataLoader(
            ds_idx,
            batch_size=batch_size,
            sampler=sampler,
            shuffle=False,
            drop_last=False,
            num_workers=num_workers,
            pin_memory=pin_memory,
            persistent_workers=(num_workers > 0),
        )

        # =========================
        # Each rank writes its shard to avoid giant all_gather
        # =========================
        if is_main_process():
            if shards_dir.exists():
                shutil.rmtree(shards_dir)
            shards_dir.mkdir(parents=True, exist_ok=True)
        if world_size() > 1:
            dist.barrier()

        shard_path = shards_dir / f"shard_rank{dist.get_rank() if dist.is_initialized() else 0}.h5"
        if shard_path.exists():
            shard_path.unlink()

        # We append to HDF5 with resizable datasets
        with h5py.File(shard_path, "w") as f:
            # Create empty resizable datasets
            f.create_dataset("idx", shape=(0,), maxshape=(None,), dtype="int64", chunks=True)
            f.create_dataset("Rimg", shape=(0, 1, 512, 512), maxshape=(None, 1, 512, 512),
                             dtype="float32", chunks=(1, 1, 512, 512))
            # E_pred/E_true are (B,M,2) => store as (n, M, 2)
            # We don't know M at creation time until first batch -> create lazily
            dset_Ep = None
            dset_Et = None
            f.create_dataset("relL2", shape=(0,), maxshape=(None,), dtype="float64", chunks=True)
            f.create_dataset("mse", shape=(0,), maxshape=(None,), dtype="float64", chunks=True)
            f.create_dataset("mae", shape=(0,), maxshape=(None,), dtype="float64", chunks=True)

            n_written = 0
            for idx, x_img, coords, E_true in loader:
                idx_np = idx.cpu().numpy().astype(np.int64)

                x_img = _to_img_tensor(x_img).to(device, non_blocking=True)
                coords = _sanitize_coords(coords).to(device, non_blocking=True)
                E_true = _sanitize_field(E_true).to(device, non_blocking=True)

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

                rel = rel_l2_per_sample(E_pred, E_true).detach().cpu().numpy().astype(np.float64)
                diff = (E_pred - E_true).detach()
                mse = torch.mean(diff * diff, dim=(1, 2)).cpu().numpy().astype(np.float64)
                mae = torch.mean(torch.abs(diff), dim=(1, 2)).cpu().numpy().astype(np.float64)

                x_np = x_img.detach().cpu().numpy().astype(np.float32)      # (B,1,512,512)
                Ep_np = E_pred.detach().cpu().numpy().astype(np.float32)    # (B,M,2)
                Et_np = E_true.detach().cpu().numpy().astype(np.float32)    # (B,M,2)

                B = idx_np.shape[0]
                M = Ep_np.shape[1]

                # Lazy create E datasets once we know M
                if dset_Ep is None:
                    dset_Ep = f.create_dataset(
                        "E_pred", shape=(0, M, 2), maxshape=(None, M, 2),
                        dtype="float32", chunks=(1, M, 2)
                    )
                    dset_Et = f.create_dataset(
                        "E_true", shape=(0, M, 2), maxshape=(None, M, 2),
                        dtype="float32", chunks=(1, M, 2)
                    )

                # Resize and append
                new_n = n_written + B
                f["idx"].resize((new_n,))
                f["Rimg"].resize((new_n, 1, 512, 512))
                f["relL2"].resize((new_n,))
                f["mse"].resize((new_n,))
                f["mae"].resize((new_n,))
                dset_Ep.resize((new_n, M, 2))
                dset_Et.resize((new_n, M, 2))

                f["idx"][n_written:new_n] = idx_np
                f["Rimg"][n_written:new_n, ...] = x_np
                f["relL2"][n_written:new_n] = rel
                f["mse"][n_written:new_n] = mse
                f["mae"][n_written:new_n] = mae
                dset_Ep[n_written:new_n, ...] = Ep_np
                dset_Et[n_written:new_n, ...] = Et_np

                n_written = new_n

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

        # =========================
        # Rank0 merges shards -> final_h5
        # =========================
        if is_main_process():
            print(f"[Infer] Shards saved in: {shards_dir}")

            shard_files = sorted(shards_dir.glob("shard_rank*.h5"))
            if len(shard_files) == 0:
                raise RuntimeError("No shard files found.")

            # Load all shard arrays (still much safer than all_gather on GPU)
            all_idx = []
            all_R = []
            all_rel = []
            all_mse = []
            all_mae = []
            all_Ep = []
            all_Et = []

            M_final = None
            for sf in shard_files:
                with h5py.File(sf, "r") as f:
                    idx = f["idx"][:]
                    Rimg = f["Rimg"][:]
                    rel = f["relL2"][:]
                    mse = f["mse"][:]
                    mae = f["mae"][:]
                    Ep = f["E_pred"][:]
                    Et = f["E_true"][:]

                if M_final is None:
                    M_final = Ep.shape[1]
                else:
                    if Ep.shape[1] != M_final:
                        raise RuntimeError(f"M mismatch among shards: {sf} has M={Ep.shape[1]} vs {M_final}")

                all_idx.append(idx)
                all_R.append(Rimg)
                all_rel.append(rel)
                all_mse.append(mse)
                all_mae.append(mae)
                all_Ep.append(Ep)
                all_Et.append(Et)

            idx_all = np.concatenate(all_idx, axis=0)
            R_all = np.concatenate(all_R, axis=0)
            rel_all = np.concatenate(all_rel, axis=0)
            mse_all = np.concatenate(all_mse, axis=0)
            mae_all = np.concatenate(all_mae, axis=0)
            Ep_all = np.concatenate(all_Ep, axis=0)
            Et_all = np.concatenate(all_Et, axis=0)

            # Sort by idx
            order = np.argsort(idx_all)
            idx_all = idx_all[order]
            R_all = R_all[order]
            rel_all = rel_all[order]
            mse_all = mse_all[order]
            mae_all = mae_all[order]
            Ep_all = Ep_all[order]
            Et_all = Et_all[order]

            # Dedup (DistributedSampler may pad)
            uniq_idx, first_pos = np.unique(idx_all, return_index=True)
            idx_all = idx_all[first_pos]
            R_all = R_all[first_pos]
            rel_all = rel_all[first_pos]
            mse_all = mse_all[first_pos]
            mae_all = mae_all[first_pos]
            Ep_all = Ep_all[first_pos]
            Et_all = Et_all[first_pos]

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

            # Save final H5
            if final_h5.exists():
                final_h5.unlink()

            with h5py.File(final_h5, "w") as f:
                f.create_dataset("idx", data=idx_all.astype(np.int64))
                f.create_dataset("Rimg", data=R_all.astype(np.float32), compression="gzip", compression_opts=4)
                f.create_dataset("coords_xy", data=coords_shared.astype(np.float32))
                f.create_dataset("E_pred", data=Ep_all.astype(np.float32), compression="gzip", compression_opts=4)
                f.create_dataset("E_true", data=Et_all.astype(np.float32), compression="gzip", compression_opts=4)
                f.create_dataset("relL2", data=rel_all.astype(np.float64))
                f.create_dataset("mse", data=mse_all.astype(np.float64))
                f.create_dataset("mae", data=mae_all.astype(np.float64))

                f.attrs["meta_matpath"] = str(matpath)
                f.attrs["meta_ckpt"] = str(ckpt_path)
                f.attrs["meta_world_size"] = int(world_size())
                f.attrs["N"] = int(N)
                f.attrs["M"] = int(Ep_all.shape[1])

            print(f"[Infer] Final saved: {final_h5}")
            print(f"[Infer] relL2 mean={rel_all.mean():.6e}, std={rel_all.std():.6e}")
            print(f"[Infer] mse mean={mse_all.mean():.6e} | mae mean={mae_all.mean():.6e}")

            # Optional: remove shards
            # shutil.rmtree(shards_dir)

    finally:
        ddp_cleanup()


if __name__ == "__main__":
    main()