from pathlib import Path
import hashlib
import time
import h5py
import logging
import numpy as np
from tqdm.auto import tqdm

from .base import BaseLoader
from ..structs import HaloCatalog, Parameters
from ..structs.parameters import to_dict
from ..io import Handler

logger = logging.getLogger(__name__)


class Py21cmFastLoader(BaseLoader):
    """Loader for data generated by py21cmfast.

    Can either load pre-existing HDF5 files from a known directory, or
    generate them on the fly by calling :meth:`generate`.

    Args:
        parameters (Parameters): Simulation and cosmology parameters.
        file_root (Path, optional): Directory containing pre-generated
            ``haloes_z*.h5`` and ``densities_z*.h5`` files. If ``None``,
            call :meth:`generate` before loading data.
    """

    def __init__(self, parameters: Parameters, file_root: "Path | str" = None):
        super().__init__(parameters)
        self.file_root = Path(file_root) if file_root is not None else None
        if self.file_root is not None and self.file_root.is_dir():
            self._ensure_parameters_yaml(self.file_root)

    @property
    def redshifts(self) -> np.ndarray:
        """Snapshot redshifts — the sparse cosmo-sim grid.

        Uses ``parameters.cosmo_sim.snapshot_redshifts`` when set, otherwise falls
        back to the full ``parameters.solver.redshifts`` profile grid (backward-compatible).
        """
        snap = self.parameters.cosmo_sim.snapshot_redshifts
        return snap if snap is not None else self.parameters.solver.redshifts

    @property
    def input_tag(self) -> str:
        """Human-readable directory name for this py21cmfast dataset.

        Pass this to :class:`~beorn.io.Handler` as ``input_tag`` so that
        BEoRN output files embed the provenance of the input data:

        .. code-block:: python

            handler = Handler(output_dir, input_tag=loader.input_tag)
        """
        sim = self.parameters.simulation
        cosmo_sim = self.parameters.cosmo_sim
        dim = sim.Ncell * cosmo_sim.py21cmfast_high_res_factor
        cosmo_hash = hashlib.md5(str(to_dict(self.parameters.cosmology)).encode()).hexdigest()[:8]
        return f"py21cmfast_N{sim.Ncell}_D{dim}_L{sim.Lbox:.0f}_seed{cosmo_sim.random_seed}_{cosmo_hash}"

    def _simulation_info(self, file_root: Path) -> str:
        """Return a formatted multi-line summary of the py21cmfast setup (no redshift list)."""
        sim = self.parameters.simulation
        cosmo_sim = self.parameters.cosmo_sim
        cosmo = self.parameters.cosmology
        dim = sim.Ncell * cosmo_sim.py21cmfast_high_res_factor
        return (
            f'py21cmfast setup:\n'
            f'  Output directory : {file_root}\n'
            f'  Grid             : HII_DIM={sim.Ncell}, DIM={dim} (factor {cosmo_sim.py21cmfast_high_res_factor}x)\n'
            f'  Box size         : {sim.Lbox:.1f} Mpc/h ({sim.Lbox / cosmo.h:.1f} Mpc)\n'
            f'  Threads          : {sim.cores}\n'
            f'  Random seed      : {cosmo_sim.random_seed}\n'
            f'  Cosmology        : Om={cosmo.Om}, Ob={cosmo.Ob}, h={cosmo.h}, '
            f'sigma_8={cosmo.sigma_8}, ns={cosmo.ns}'
        )

    def _py21cmfast_yaml_dict(self) -> dict:
        """Return only the parameters needed to reproduce this py21cmfast run.

        Includes the cosmological inputs to :class:`py21cmfast.CosmoParams` and
        the grid/seed inputs to :class:`py21cmfast.UserParams`.  Astrophysical
        BEoRN parameters, redshift lists, and solver settings are intentionally
        omitted — redshifts can be inferred from the filenames on disk.
        """
        sim = self.parameters.simulation
        cosmo_sim = self.parameters.cosmo_sim
        cosmo = self.parameters.cosmology
        return {
            "cosmology": {
                "Om": cosmo.Om,
                "Ob": cosmo.Ob,
                "Ol": cosmo.Ol,
                "h": cosmo.h,
                "sigma_8": cosmo.sigma_8,
                "ns": cosmo.ns,
            },
            "simulation": {
                "Ncell": sim.Ncell,
                "Lbox": sim.Lbox,
                "cores": sim.cores,
            },
            "cosmo_sim": {
                "py21cmfast_high_res_factor": cosmo_sim.py21cmfast_high_res_factor,
                "random_seed": cosmo_sim.random_seed,
            },
        }

    def _ensure_parameters_yaml(self, directory: Path) -> None:
        """Write ``cosmo_params.yaml`` to *directory* if it does not already exist.

        Only the cosmological and grid/seed parameters used by py21cmfast are
        written — not the full BEoRN parameter set.
        """
        import yaml
        yaml_path = directory / "cosmo_params.yaml"
        if not yaml_path.exists():
            with yaml_path.open("w") as f:
                yaml.dump(self._py21cmfast_yaml_dict(), f, default_flow_style=False, sort_keys=False)
            logger.info(
                f"cosmo_params.yaml not found in {directory} — written from current parameters."
            )

    def generate(self, handler: Handler) -> None:
        """Run py21cmfast to generate halo catalogs and density fields.

        Produces ``haloes_z{z:.3f}.h5`` and ``densities_z{z:.3f}.h5`` files for
        every redshift in ``parameters.cosmo_sim.snapshot_redshifts`` under a subdirectory
        of ``handler.file_root``. The directory is keyed on cosmological and
        simulation parameters (not astrophysics, not the redshift list), so:

        - Changing astrophysical parameters reuses the same directory.
        - Adding or removing redshifts only generates the missing snapshots;
          existing ones are left untouched.
        - Changing cosmology or grid/seed creates a new directory.

        Sets ``self.file_root`` so that :meth:`load_halo_catalog` and
        :meth:`load_density_field` can be called immediately after.

        Args:
            handler (Handler): IO handler that determines the storage location.

        Raises:
            ImportError: If ``py21cmfast`` is not installed.
        """
        sim = self.parameters.simulation
        cosmo_sim = self.parameters.cosmo_sim
        cosmo = self.parameters.cosmology
        dim = sim.Ncell * cosmo_sim.py21cmfast_high_res_factor
        file_root = handler.file_root / self.input_tag

        all_redshifts = list(self.redshifts)
        present = [z for z in all_redshifts
                   if (file_root / f'haloes_z{z:.3f}.h5').exists() and
                      (file_root / f'densities_z{z:.3f}.h5').exists()]
        missing = [z for z in all_redshifts if z not in present]

        logger.info(self._simulation_info(file_root))

        if not missing:
            logger.info(f'All {len(present)} snapshots already cached. Skipping generation.')
            self.file_root = file_root
            self._ensure_parameters_yaml(file_root)
            return

        if present:
            logger.info(
                f'Found {len(present)}/{len(all_redshifts)} cached snapshots: '
                f'z = {", ".join(f"{z:.3f}" for z in present)}'
            )
        logger.info(
            f'Generating {len(missing)}/{len(all_redshifts)} missing snapshots: '
            f'z = {", ".join(f"{z:.3f}" for z in missing)}'
        )

        try:
            import py21cmfast as p21c
        except ImportError:
            raise ImportError(
                'Generating data with py21cmfast requires additional dependencies. '
                'Install beorn with the extra package set: `pip install beorn[extra]` '
                'or see https://github.com/21cmfast/21cmFAST/.'
            )

        start_time = time.process_time()
        file_root.mkdir(parents=True, exist_ok=True)
        self._ensure_parameters_yaml(file_root)

        user_params = p21c.UserParams(
            HII_DIM=sim.Ncell,
            DIM=dim,
            BOX_LEN=sim.Lbox / cosmo.h,
            USE_INTERPOLATION_TABLES=True,
            N_THREADS=sim.cores,
        )


        cosmo_params = p21c.CosmoParams(
            SIGMA_8=cosmo.sigma_8,
            hlittle=cosmo.h,
            OMm=cosmo.Om,
            OMb=cosmo.Ob,
            POWER_INDEX=cosmo.ns,
        )

        global_params = {
            "INITIAL_REDSHIFT": 300,
            "CLUMPING_FACTOR": 2.0,
        }

        logger.info('Computing initial conditions...')
        ic_start = time.process_time()
        IC = p21c.initial_conditions(
            user_params=user_params,
            cosmo_params=cosmo_params,
            random_seed=cosmo_sim.random_seed,
            direc=str(file_root),
        )
        logger.info(f'Initial conditions done in {time.process_time() - ic_start:.1f}s.')

        with p21c.global_params.use(**global_params):
            with tqdm(missing, desc='py21cmfast snapshots', unit='snapshot') as pbar:
                for redshift in pbar:
                    pbar.set_postfix(z=f"{redshift:.3f}", refresh=False)
                    halo_fname = f'haloes_z{redshift:.3f}.h5'
                    field_fname = f'densities_z{redshift:.3f}.h5'

                    z_start = time.process_time()
                    perturbed_field = p21c.perturb_field(
                        redshift=redshift, init_boxes=IC, direc=str(file_root), write=False
                    )
                    halo_list = p21c.perturb_halo_list(
                        redshift=redshift, init_boxes=IC, direc=str(file_root), write=False
                    )

                    halo_list.write(direc=file_root, fname=halo_fname)
                    perturbed_field.write(direc=file_root, fname=field_fname)
                    logger.debug(
                        f'z={redshift:.3f}: {len(halo_list.halo_masses)} halos, '
                        f'wrote {halo_fname} and {field_fname} '
                        f'({time.process_time() - z_start:.1f}s)'
                    )

        self.file_root = file_root
        logger.info(
            f'Finished generating {len(missing)} snapshots in {time.process_time() - start_time:.1f}s. '
            f'Directory: {file_root}'
        )

    def load_halo_catalog(self, redshift_index: int) -> HaloCatalog:
        self._require_file_root()
        redshift = self.redshifts[redshift_index]
        path = self.file_root / f'haloes_z{redshift:.3f}.h5'

        with h5py.File(path, 'r') as f:
            haloes = f['PerturbHaloField']
            m = np.asarray(haloes['halo_masses'])
            positions = np.asarray(haloes['halo_coords'])

        scaling = float(self.parameters.simulation.Lbox / self.parameters.simulation.Ncell)
        return HaloCatalog(
            masses=m * self.parameters.cosmology.h,
            positions=positions * scaling,
            parameters=self.parameters,
            redshift_index=redshift_index,
            redshift=float(redshift),
        )

    def load_density_field(self, redshift_index: int) -> np.ndarray:
        self._require_file_root()
        redshift = self.redshifts[redshift_index]
        path = self.file_root / f'densities_z{redshift:.3f}.h5'

        with h5py.File(path, 'r') as f:
            return f['PerturbedField']['density'][:]

    def load_rsd_fields(self, redshift_index: int):
        raise NotImplementedError("RSD fields are not yet implemented for Py21cmFastLoader.")

    def _require_file_root(self):
        if self.file_root is None:
            raise RuntimeError(
                "No data directory set. Either call generate(handler) to produce data "
                "with py21cmfast, or pass file_root to the constructor to load existing files."
            )
