"""
HDF5-based checkpointing for lattice field theory simulations.

This module provides functions to save and load simulation state to HDF5 files,
enabling interruption and resumption of long-running simulations.

Supported data structures:
- Lattice: Scalar field simulations
- CoupledLattice: Abelian Higgs model (scalar + U(1) gauge)
- GaugeLattice: Pure gauge field simulations
- FRWUniverse: Cosmological background state

Example usage:
    >>> from jaxlatt.io import save_checkpoint, load_checkpoint
    >>>
    >>> # Save simulation state
    >>> save_checkpoint("checkpoint.h5", lattice=lat, universe=frw, time=t, step=n)
    >>>
    >>> # Resume simulation
    >>> state = load_checkpoint("checkpoint.h5")
    >>> lat, frw, t, n = state["lattice"], state["universe"], state["time"], state["step"]
"""

from typing import Optional, Dict, Any, Union
from pathlib import Path
import numpy as np

try:
    import h5py
    HAS_H5PY = True
except ImportError:
    HAS_H5PY = False

from jaxlatt.core.lattice import Lattice, CoupledLattice, GaugeLattice
from jaxlatt.core.cosmology import FRWUniverse


def _check_h5py():
    """Raise ImportError if h5py is not available."""
    if not HAS_H5PY:
        raise ImportError(
            "h5py is required for HDF5 checkpointing. "
            "Install with: pip install h5py"
        )


def save_lattice(group: "h5py.Group", lattice: Lattice) -> None:
    """
    Save a Lattice to an HDF5 group.

    Args:
        group: HDF5 group to save into
        lattice: Lattice instance to save
    """
    group.attrs["type"] = "Lattice"
    group.attrs["ndim"] = lattice.ndim
    group.attrs["size"] = lattice.size
    group.attrs["length"] = lattice.length
    group.attrs["dx"] = lattice.dx

    group.create_dataset("field", data=np.asarray(lattice.field), compression="gzip")
    group.create_dataset("field_dot", data=np.asarray(lattice.field_dot), compression="gzip")


def load_lattice(group: "h5py.Group") -> Lattice:
    """
    Load a Lattice from an HDF5 group.

    Args:
        group: HDF5 group containing lattice data

    Returns:
        Reconstructed Lattice instance
    """
    import jax.numpy as jnp

    size = tuple(group.attrs["size"])
    length = tuple(group.attrs["length"])
    field = jnp.array(group["field"][:])
    field_dot = jnp.array(group["field_dot"][:])

    return Lattice(size=size, length=length, field=field, field_dot=field_dot)


def save_coupled_lattice(group: "h5py.Group", lattice: CoupledLattice) -> None:
    """
    Save a CoupledLattice to an HDF5 group.

    Args:
        group: HDF5 group to save into
        lattice: CoupledLattice instance to save
    """
    group.attrs["type"] = "CoupledLattice"
    group.attrs["size"] = lattice.size
    group.attrs["length"] = lattice.length
    group.attrs["dx"] = lattice.dx
    group.attrs["m"] = lattice.m
    group.attrs["lambda_"] = lattice.lambda_
    group.attrs["g"] = lattice.g

    # Complex arrays need special handling
    group.create_dataset("phi", data=np.asarray(lattice.phi), compression="gzip")
    group.create_dataset("pi", data=np.asarray(lattice.pi), compression="gzip")
    group.create_dataset("links", data=np.asarray(lattice.links), compression="gzip")
    group.create_dataset("E", data=np.asarray(lattice.E), compression="gzip")


def load_coupled_lattice(group: "h5py.Group") -> CoupledLattice:
    """
    Load a CoupledLattice from an HDF5 group.

    Args:
        group: HDF5 group containing lattice data

    Returns:
        Reconstructed CoupledLattice instance
    """
    import jax.numpy as jnp

    return CoupledLattice(
        phi=jnp.array(group["phi"][:]),
        pi=jnp.array(group["pi"][:]),
        links=jnp.array(group["links"][:]),
        E=jnp.array(group["E"][:]),
        m=float(group.attrs["m"]),
        lambda_=float(group.attrs["lambda_"]),
        g=float(group.attrs["g"]),
        dx=float(group.attrs["dx"]),
        size=tuple(group.attrs["size"]),
        length=tuple(group.attrs["length"]),
    )


def save_gauge_lattice(group: "h5py.Group", lattice: GaugeLattice) -> None:
    """
    Save a GaugeLattice to an HDF5 group.

    Args:
        group: HDF5 group to save into
        lattice: GaugeLattice instance to save
    """
    group.attrs["type"] = "GaugeLattice"
    group.attrs["size"] = lattice.size
    group.attrs["length"] = lattice.length
    group.attrs["dx"] = lattice.dx
    group.attrs["g"] = lattice.g

    group.create_dataset("links", data=np.asarray(lattice.links), compression="gzip")
    group.create_dataset("E", data=np.asarray(lattice.E), compression="gzip")


def load_gauge_lattice(group: "h5py.Group") -> GaugeLattice:
    """
    Load a GaugeLattice from an HDF5 group.

    Args:
        group: HDF5 group containing lattice data

    Returns:
        Reconstructed GaugeLattice instance
    """
    import jax.numpy as jnp

    return GaugeLattice(
        links=jnp.array(group["links"][:]),
        E=jnp.array(group["E"][:]),
        dx=float(group.attrs["dx"]),
        size=tuple(group.attrs["size"]),
        length=tuple(group.attrs["length"]),
        g=float(group.attrs["g"]),
    )


def save_frw_universe(group: "h5py.Group", universe: FRWUniverse) -> None:
    """
    Save an FRWUniverse to an HDF5 group.

    Args:
        group: HDF5 group to save into
        universe: FRWUniverse instance to save
    """
    group.attrs["type"] = "FRWUniverse"
    group.attrs["a"] = universe.a
    group.attrs["adot"] = universe.adot
    group.attrs["tau"] = universe.tau
    group.attrs["M_pl"] = universe.M_pl


def load_frw_universe(group: "h5py.Group") -> FRWUniverse:
    """
    Load an FRWUniverse from an HDF5 group.

    Args:
        group: HDF5 group containing universe data

    Returns:
        Reconstructed FRWUniverse instance
    """
    return FRWUniverse(
        a=float(group.attrs["a"]),
        adot=float(group.attrs["adot"]),
        tau=float(group.attrs["tau"]),
        M_pl=float(group.attrs["M_pl"]),
    )


def save_checkpoint(
    path: Union[str, Path],
    lattice: Optional[Union[Lattice, CoupledLattice, GaugeLattice]] = None,
    universe: Optional[FRWUniverse] = None,
    time: Optional[float] = None,
    step: Optional[int] = None,
    metadata: Optional[Dict[str, Any]] = None,
    overwrite: bool = True,
) -> None:
    """
    Save simulation checkpoint to HDF5 file.

    Args:
        path: Output file path
        lattice: Lattice state to save (Lattice, CoupledLattice, or GaugeLattice)
        universe: Cosmological background state (optional)
        time: Current simulation time
        step: Current timestep number
        metadata: Additional metadata dictionary (scalars and strings only)
        overwrite: Whether to overwrite existing file

    Example:
        >>> save_checkpoint(
        ...     "checkpoint.h5",
        ...     lattice=lat,
        ...     universe=frw,
        ...     time=10.5,
        ...     step=1000,
        ...     metadata={"dt": 0.01, "potential": "double_well"}
        ... )
    """
    _check_h5py()

    path = Path(path)
    mode = "w" if overwrite else "x"

    with h5py.File(path, mode) as f:
        # Save version info
        f.attrs["jaxlatt_version"] = "0.1.0"
        f.attrs["format_version"] = 1

        # Save simulation state
        if time is not None:
            f.attrs["time"] = time
        if step is not None:
            f.attrs["step"] = step

        # Save lattice
        if lattice is not None:
            grp = f.create_group("lattice")
            if isinstance(lattice, Lattice):
                save_lattice(grp, lattice)
            elif isinstance(lattice, CoupledLattice):
                save_coupled_lattice(grp, lattice)
            elif isinstance(lattice, GaugeLattice):
                save_gauge_lattice(grp, lattice)
            else:
                raise TypeError(f"Unknown lattice type: {type(lattice)}")

        # Save universe
        if universe is not None:
            grp = f.create_group("universe")
            save_frw_universe(grp, universe)

        # Save metadata
        if metadata is not None:
            grp = f.create_group("metadata")
            for key, value in metadata.items():
                if isinstance(value, (int, float, str, bool)):
                    grp.attrs[key] = value
                elif isinstance(value, np.ndarray):
                    grp.create_dataset(key, data=value)
                else:
                    # Try to convert to string
                    grp.attrs[key] = str(value)


def load_checkpoint(path: Union[str, Path]) -> Dict[str, Any]:
    """
    Load simulation checkpoint from HDF5 file.

    Args:
        path: Input file path

    Returns:
        Dictionary containing:
        - "lattice": Lattice/CoupledLattice/GaugeLattice (if saved)
        - "universe": FRWUniverse (if saved)
        - "time": Simulation time (if saved)
        - "step": Timestep number (if saved)
        - "metadata": Additional metadata dict (if saved)

    Example:
        >>> state = load_checkpoint("checkpoint.h5")
        >>> lat = state["lattice"]
        >>> t = state["time"]
        >>> step = state["step"]
    """
    _check_h5py()

    path = Path(path)
    result = {}

    with h5py.File(path, "r") as f:
        # Load simulation state
        if "time" in f.attrs:
            result["time"] = float(f.attrs["time"])
        if "step" in f.attrs:
            result["step"] = int(f.attrs["step"])

        # Load lattice
        if "lattice" in f:
            grp = f["lattice"]
            lattice_type = grp.attrs["type"]

            if lattice_type == "Lattice":
                result["lattice"] = load_lattice(grp)
            elif lattice_type == "CoupledLattice":
                result["lattice"] = load_coupled_lattice(grp)
            elif lattice_type == "GaugeLattice":
                result["lattice"] = load_gauge_lattice(grp)
            else:
                raise ValueError(f"Unknown lattice type: {lattice_type}")

        # Load universe
        if "universe" in f:
            result["universe"] = load_frw_universe(f["universe"])

        # Load metadata
        if "metadata" in f:
            grp = f["metadata"]
            metadata = dict(grp.attrs)
            for key in grp.keys():
                metadata[key] = grp[key][:]
            result["metadata"] = metadata

    return result
