import h5py
import numpy as np
import torch
from torch.utils.data import Dataset


# ----------------------------
# Helpers: robust mat(v7.3) read
# ----------------------------
def _read_array(f, key):
    if key not in f:
        raise KeyError(f"Key '{key}' not found. Available keys: {list(f.keys())}")
    return np.array(f[key])


def _ensure_coords_m2(coords):
    # coords expected: (m,2). Some mat files read as (2,m)
    if coords.ndim != 2:
        raise ValueError(f"coords must be 2D, got {coords.shape}")
    if coords.shape[1] == 2:
        return coords
    if coords.shape[0] == 2:
        return coords.T
    raise ValueError(f"coords should be (m,2) or (2,m), got {coords.shape}")


def _ensure_X_nchw(X, H=512, W=512):
    """
    Expect geometry:
      (N,1,H,W)  OR (H,W,1,N) OR (1,H,W,N)
    """
    if X.ndim != 4:
        raise ValueError(f"X must be 4D, got {X.shape}")

    # (N,1,H,W)
    if X.shape[0] > 1 and X.shape[1] == 1 and X.shape[2] == H and X.shape[3] == W:
        return X

    # (H,W,1,N) -> (N,1,H,W)
    if X.shape[0] == H and X.shape[1] == W and X.shape[2] == 1:
        return np.transpose(X, (3, 2, 0, 1))

    # (1,H,W,N) -> (N,1,H,W)
    if X.shape[0] == 1 and X.shape[1] == H and X.shape[2] == W:
        return np.transpose(X, (3, 0, 1, 2))

    raise ValueError(f"Unrecognized X layout: {X.shape}")


def _ensure_Y_Nm(Y, N_expected=None, m_expected=None):
    # expect (N,m) or (m,N)
    if Y.ndim != 2:
        raise ValueError(f"Y must be 2D, got {Y.shape}")

    if N_expected is not None and Y.shape[0] == N_expected:
        return Y
    if N_expected is not None and Y.shape[1] == N_expected:
        return Y.T

    if m_expected is not None and Y.shape[1] == m_expected:
        return Y
    if m_expected is not None and Y.shape[0] == m_expected:
        return Y.T

    return Y


# ----------------------------
# Dataset (NO normalization here!)
# ----------------------------
class MetalensDataset(Dataset):
    def __init__(self, X_geom, coords, Y_re, Y_im):
        """
        X_geom: (N,1,512,512)  already in {0,1} or float
        coords: (m,2)          already normalized to [-1,1]
        Y_re:   (N,m)          already normalized
        Y_im:   (N,m)          already normalized
        """
        # cast geometry to float32 for convs
        self.X = torch.tensor(X_geom, dtype=torch.float32)
        self.coords = torch.tensor(coords, dtype=torch.float32)     # shared for all samples
        self.Yre = torch.tensor(Y_re, dtype=torch.float32)
        self.Yim = torch.tensor(Y_im, dtype=torch.float32)

    def __len__(self):
        return self.X.shape[0]

    def __getitem__(self, idx):
        x = self.X[idx]  # (1,512,512)
        y = torch.stack([self.Yre[idx], self.Yim[idx]], dim=-1)  # (m,2)
        return x, self.coords, y


# ----------------------------
# Build datasets
# ----------------------------
def build_datasets(
    mat_path="metalens_dataset.mat",
    train_ratio=0.8,
    keys=None,
):
    """
    Read your new dataset:
      Rimg     : (N,1,512,512)  (or equivalent layout)
      coords   : (m,2)
      Ere_flat : (N,m)
      Eim_flat : (N,m)
    No norm_stats needed.
    """
    if keys is None:
        keys = dict(
            X="Rimg",
            coords="coords",
            Yre="Ere_flat",
            Yim="Eim_flat",
        )

    with h5py.File(mat_path, "r") as f:
        X = _read_array(f, keys["X"])
        coords = _read_array(f, keys["coords"])
        Yre = _read_array(f, keys["Yre"])
        Yim = _read_array(f, keys["Yim"])

    coords = _ensure_coords_m2(coords)
    X = _ensure_X_nchw(X, 512, 512)

    N = X.shape[0]
    m = coords.shape[0]

    Yre = _ensure_Y_Nm(Yre, N_expected=N, m_expected=m)
    Yim = _ensure_Y_Nm(Yim, N_expected=N, m_expected=m)

    if Yre.shape != (N, m) or Yim.shape != (N, m):
        raise ValueError(f"Y shapes mismatch: Yre={Yre.shape}, Yim={Yim.shape}, expected {(N, m)}")

    if N < 2:
        raise ValueError(f"N too small: {N}")

    n_train = int(train_ratio * N)
    n_train = min(max(1, n_train), N - 1)

    X_train, X_test = X[:n_train], X[n_train:]
    Yre_train, Yre_test = Yre[:n_train], Yre[n_train:]
    Yim_train, Yim_test = Yim[:n_train], Yim[n_train:]

    train_ds = MetalensDataset(X_train, coords, Yre_train, Yim_train)
    test_ds = MetalensDataset(X_test, coords, Yre_test, Yim_test)

    print(f"[build_datasets] X: {X.shape}, coords: {coords.shape}, Y: {Yre.shape}")
    print(f"[build_datasets] Split train={len(train_ds)}, test={len(test_ds)}")

    # norm_stats 不再需要：返回空 dict
    norm_stats = {}
    return train_ds, test_ds, norm_stats