# PI_code_test_FNO_inference.py
# =========================================================
# ✅ Inference ONLY for Encoder+FNO (no DDP, no training)
# ✅ Can read BOTH MATLAB v7 (scipy.io.loadmat) and v7.3 (HDF5 via h5py)
# ✅ Added: accurate inference timing (CUDA synchronized) + save timing into .mat
#
# Input .mat must contain (names can vary, see candidates below):
#   - Rimg  : (N,1,512,512) OR (512,512,N) OR (512,512) OR (1,512,512)
#   - coords: (M,2) in [-1,1]  (or (1,M,2))
#
# Output .mat:
#   - E_pred    : (N,M,2) float32  (Re, Im)
#   - coords_xy : (M,2) float32
#   - Rimg      : (N,1,512,512) float32
#   - meta_*    : timing + shapes + ckpt path
# =========================================================

import argparse
import math
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from scipy.io import loadmat, savemat


# -------------------------
# Utils
# -------------------------
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)


def _pick_key(d, keys):
    for k in keys:
        if k in d:
            return k
    return None


def _as_numpy_from_h5(obj):
    """
    Convert h5py dataset to numpy.
    MATLAB v7.3 stores arrays in column-major and often with transposed dims;
    We will reshape/transpose later in the Rimg handler.
    """
    import numpy as _np
    return _np.array(obj)


def load_mat_auto(mat_path: str):
    """
    Load .mat file:
      - MATLAB v7.2 and below: scipy.io.loadmat
      - MATLAB v7.3 (HDF5): h5py
    Returns: dict-like mapping name -> numpy array
    """
    mat_path = str(mat_path)
    try:
        data = loadmat(mat_path)
        # remove matlab meta keys if exist
        return {k: v for k, v in data.items() if not k.startswith("__")}
    except NotImplementedError:
        import h5py
        data = {}
        with h5py.File(mat_path, "r") as f:
            for k in f.keys():
                data[k] = _as_numpy_from_h5(f[k])
        return data


def load_meta_inference_mat(mat_in: str):
    """
    Returns:
      Rimg_np: (N,1,512,512) float32
      coords_np: (M,2) float32
    """
    data = load_mat_auto(mat_in)

    # candidates (you can add more if needed)
    k_r = _pick_key(data, ["Rimg", "rimg", "geom", "geometry", "M"])
    k_c = _pick_key(data, ["coords", "coord", "coords_xy", "xy"])

    if k_r is None or k_c is None:
        raise KeyError(
            f"[meta_inference] mat must contain geometry + coords.\n"
            f"Found keys: {list(data.keys())}\n"
            f"Geometry candidates: [Rimg,rimg,geom,geometry,M]\n"
            f"Coords candidates:   [coords,coord,coords_xy,xy]"
        )

    Rimg_np = np.array(data[k_r])
    coords_np = np.array(data[k_c])

    # ---- coords -> (M,2)
    coords_np = np.squeeze(coords_np)

    # MATLAB v7.3 (h5py) sometimes gives (2,M) for a (M,2) saved matrix
    if coords_np.ndim == 2 and coords_np.shape[0] == 2 and coords_np.shape[1] != 2:
        coords_np = coords_np.T

    if coords_np.ndim != 2 or coords_np.shape[1] != 2:
        raise ValueError(f"[meta_inference] coords must be (M,2). Got {coords_np.shape}")
    coords_np = coords_np.astype(np.float32)

    # ---- Rimg -> (N,1,512,512)
    Rimg_np = np.squeeze(Rimg_np)

    if Rimg_np.ndim == 2:
        # (512,512)
        Rimg_np = Rimg_np[None, None, :, :]

    elif Rimg_np.ndim == 3:
        # could be (512,512,N) OR (N,512,512) OR (1,512,512)
        if Rimg_np.shape[0] == 1 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512:
            Rimg_np = Rimg_np[:, None, :, :]  # (1,1,512,512)

        elif Rimg_np.shape[0] == 512 and Rimg_np.shape[1] == 512:
            # (512,512,N) -> (N,1,512,512)
            Rimg_np = np.transpose(Rimg_np, (2, 0, 1))[:, None, :, :]

        elif Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512:
            # (N,512,512) -> (N,1,512,512)
            Rimg_np = Rimg_np[:, None, :, :]

        else:
            raise ValueError(f"[meta_inference] Unrecognized Rimg 3D shape: {Rimg_np.shape}")

    elif Rimg_np.ndim == 4:
        # could be (N,1,512,512) OR (1,512,512,N) OR (512,512,1,N) etc.
        if Rimg_np.shape[1] == 1 and Rimg_np.shape[2] == 512 and Rimg_np.shape[3] == 512:
            # (N,1,512,512)
            pass

        elif Rimg_np.shape[0] == 1 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512:
            # (1,512,512,N) -> (N,1,512,512)
            Rimg_np = np.transpose(Rimg_np, (3, 0, 1, 2))

        elif Rimg_np.shape[0] == 512 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 1:
            # (512,512,1,N) -> (N,1,512,512)
            Rimg_np = np.transpose(Rimg_np, (3, 2, 0, 1))

        else:
            raise ValueError(f"[meta_inference] Unrecognized Rimg 4D shape: {Rimg_np.shape}")

    else:
        raise ValueError(f"[meta_inference] Unrecognized Rimg shape: {Rimg_np.shape}")

    if Rimg_np.shape[1] != 1 or Rimg_np.shape[-2:] != (512, 512):
        raise ValueError(f"[meta_inference] Rimg must be (N,1,512,512). Got {Rimg_np.shape}")

    return Rimg_np.astype(np.float32), coords_np.astype(np.float32)


# -------------------------
# FNO core
# -------------------------
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


# -------------------------
# 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 (OFF by default)
# -------------------------
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:
        if coords.dim() == 4:
            coords = coords[:, 0, :, :]
        return coords

    @staticmethod
    def _coords_to_xy_grid(coords: torch.Tensor):
        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, getattr(device, "index", None), 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


# -------------------------
# Inference
# -------------------------
@torch.no_grad()
def run_inference(
    mat_in: str,
    ckpt: str,
    mat_out: str,
    device: str,
    batch_size: int,
    modes: int,
    width: int,
    blocks: int,
    padding: int,
    dropout: float,
    layerscale: float,
    geom_ch: int,
    use_fourier: bool,
    fourier_K: int,
    fourier_scale: float,
    dropout_xy: float,
):
    Rimg_np, coords_np = load_meta_inference_mat(mat_in)
    N = int(Rimg_np.shape[0])
    M = int(coords_np.shape[0])

    dev = torch.device(device if torch.cuda.is_available() else "cpu")

    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,
    ).to(dev)

    # load weights
    try:
        state = torch.load(ckpt, map_location="cpu", weights_only=True)
    except TypeError:
        state = torch.load(ckpt, map_location="cpu")
    model.load_state_dict(state, strict=True)
    model.eval()

    coords_t = torch.from_numpy(coords_np).to(dev).unsqueeze(0)  # (1,M,2)

    # ---- timing helpers (CUDA safe)
    def _sync():
        if dev.type == "cuda":
            torch.cuda.synchronize(dev)

    # warmup (recommended)
    warmup_steps = min(2, (N + batch_size - 1) // batch_size)
    for wi in range(warmup_steps):
        s = wi * batch_size
        e = min(N, s + batch_size)
        x_img = torch.from_numpy(Rimg_np[s:e]).to(dev, dtype=torch.float32)
        coords_b = coords_t.expand(e - s, -1, -1).contiguous()
        _ = model(x_img, coords_b)
    _sync()

    preds = []
    t_total0 = time.perf_counter()
    forward_time_sum = 0.0
    num_samples_done = 0

    for s in range(0, N, batch_size):
        e = min(N, s + batch_size)
        x_img = torch.from_numpy(Rimg_np[s:e]).to(dev, dtype=torch.float32)
        coords_b = coords_t.expand(e - s, -1, -1).contiguous()

        _sync()
        t0 = time.perf_counter()
        pred = model(x_img, coords_b)  # (B,M,2)
        _sync()
        t1 = time.perf_counter()

        forward_time_sum += (t1 - t0)
        num_samples_done += (e - s)

        preds.append(pred.detach().cpu().numpy().astype(np.float32))

    t_total1 = time.perf_counter()
    total_time = t_total1 - t_total0

    avg_time_per_sample = total_time / max(num_samples_done, 1)
    avg_forward_per_sample = forward_time_sum / max(num_samples_done, 1)

    fps_total = num_samples_done / max(total_time, 1e-12)
    fps_forward = num_samples_done / max(forward_time_sum, 1e-12)

    E_pred = np.concatenate(preds, axis=0)  # (N,M,2)

    print("========== Inference Timing ==========")
    print(f"[Time] total elapsed           : {total_time:.6f} s")
    print(f"[Time] total forward (model)   : {forward_time_sum:.6f} s")
    print(f"[Time] avg elapsed / sample    : {avg_time_per_sample*1e3:.3f} ms")
    print(f"[Time] avg forward / sample    : {avg_forward_per_sample*1e3:.3f} ms")
    print(f"[Perf] throughput (total)      : {fps_total:.3f} samples/s")
    print(f"[Perf] throughput (forward)    : {fps_forward:.3f} samples/s")
    print("======================================")

    savemat(
        mat_out,
        {
            "E_pred": E_pred,
            "coords_xy": coords_np.astype(np.float32),
            "meta_ckpt": str(ckpt),
            "meta_N": np.int64(N),
            "meta_M": np.int64(M),

            # timing meta
            "meta_time_total_s": np.float64(total_time),
            "meta_time_forward_s": np.float64(forward_time_sum),
            "meta_ms_per_sample_total": np.float64(avg_time_per_sample * 1e3),
            "meta_ms_per_sample_forward": np.float64(avg_forward_per_sample * 1e3),
            "meta_throughput_total_sps": np.float64(fps_total),
            "meta_throughput_forward_sps": np.float64(fps_forward),
        },
    )

    print(f"[OK] mat_in  : {mat_in}")
    print(f"[OK] ckpt    : {ckpt}")
    print(f"[OK] mat_out : {mat_out}")
    print(f"[OK] shapes  : Rimg={Rimg_np.shape}, coords={coords_np.shape}, E_pred={E_pred.shape}")


def parse_args():
    p = argparse.ArgumentParser("Inference only: Encoder+FNO")
    p.add_argument("--mat_in", type=str, default="meta_inference.mat")
    p.add_argument("--ckpt", type=str, default="fno_ddpsafe_best.pth")
    p.add_argument("--mat_out", type=str, default="meta_inference_pred.mat")
    p.add_argument("--device", type=str, default="cuda:0")
    p.add_argument("--batch_size", type=int, default=16)

    # must match training hyperparameters
    p.add_argument("--modes", type=int, default=6)
    p.add_argument("--width", type=int, default=24)
    p.add_argument("--blocks", type=int, default=4)
    p.add_argument("--padding", type=int, default=0)
    p.add_argument("--dropout", type=float, default=0.10)
    p.add_argument("--layerscale", type=float, default=0.10)
    p.add_argument("--geom_ch", type=int, default=8)

    p.add_argument("--use_fourier", action="store_true")
    p.add_argument("--fourier_K", type=int, default=2)
    p.add_argument("--fourier_scale", type=float, default=1.0)
    p.add_argument("--dropout_xy", type=float, default=0.05)
    return p.parse_args()


def main():
    args = parse_args()

    if not Path(args.mat_in).exists():
        raise FileNotFoundError(f"Input mat not found: {args.mat_in}")
    if not Path(args.ckpt).exists():
        raise FileNotFoundError(f"Checkpoint not found: {args.ckpt}")

    run_inference(
        mat_in=args.mat_in,
        ckpt=args.ckpt,
        mat_out=args.mat_out,
        device=args.device,
        batch_size=args.batch_size,
        modes=args.modes,
        width=args.width,
        blocks=args.blocks,
        padding=args.padding,
        dropout=args.dropout,
        layerscale=args.layerscale,
        geom_ch=args.geom_ch,
        use_fourier=bool(args.use_fourier),
        fourier_K=args.fourier_K,
        fourier_scale=args.fourier_scale,
        dropout_xy=args.dropout_xy,
    )


if __name__ == "__main__":
    main()