import os
import h5py
import numpy as np
from ewokscore import Task


class ExportHDF5(
    Task,
    input_names=[
        "sample_name",
        "tth",
        "det",
        "tth_bg",
        "det_bg",
        "q",
        "R",
        "qraw",
        "Rraw",
        "z",
        "delta_rho",
        "delta_rho_bl",
        "delta_rho_nc",
        "detm",
        "detm_bg",
        "tthc",
        "tthcr",
        "tthc_bg",
        "tthcr_bg",
        "detmn",
        "detmn_bg",
        "detmnp",
        "detmnpr",
        "detmnp_bg",
        "detmnpr_bg",
        "detmnpri_bg",
        "fpc",
        "baseline",
        "minima",
        "phase",
        "phase_bl",
        "Rq4_bl",
        "Rq4_sm",
        "filename",
    ],
    optional_input_names=[
        "start_time_str",
        "end_time_str",
    ],
    output_names=["filename"],
):
    """
    Export to an HDF5 file organized like NeXus:
    - /entry (NXentry) with sample, times
    - /entry/instrument/detector (NXdetector) for raw counts
    - /entry/data/* (NXdata) for plot-ready signals (R(tth), R(q), Rraw(qraw))
    - /entry/process (NXprocess) for intermediate/derived arrays
    """

    @staticmethod
    def _ensure_group(parent: h5py.Group, path: str) -> h5py.Group:
        g = parent
        for part in path.strip("/").split("/"):
            if part:
                g = g.require_group(part)
        return g

    @staticmethod
    def _write_1d(
        parent: h5py.Group,
        name: str,
        data,
        *,
        units: str | None = None,
        long_name: str | None = None,
    ) -> None:
        if data is None:
            return
        if isinstance(data, np.ndarray):
            arr = data
        else:
            arr = np.asarray(data)
        if arr.ndim != 1:
            return
        dset = parent.create_dataset(name, data=arr)
        if units is not None:
            dset.attrs["units"] = units
        if long_name is not None:
            dset.attrs["long_name"] = long_name

    @staticmethod
    def _nxclass(g: h5py.Group, cls: str) -> None:
        g.attrs["NX_class"] = cls

    def run(self):
        inp = self.inputs

        # Decide output filename
        out = inp.filename or f"{(inp.sample_name or 'output')}.h5"
        root_dir = os.path.dirname(out)
        if root_dir and not os.path.isdir(root_dir):
            os.makedirs(root_dir, exist_ok=True)

        with h5py.File(out, "w") as f:
            # Root default → entry
            f.attrs["default"] = "entry"

            # /entry (NXentry)
            entry = self._ensure_group(f, "entry")
            self._nxclass(entry, "NXentry")

            # Title & times
            title = inp.sample_name or "unknown"
            entry.create_dataset("title", data=np.bytes_(title))
            if getattr(inp, "start_time_str", None):
                entry.create_dataset("start_time", data=np.bytes_(inp.start_time_str))
            if getattr(inp, "end_time_str", None):
                entry.create_dataset("end_time", data=np.bytes_(inp.end_time_str))

            # /entry/sample (NXsample)
            sample = self._ensure_group(entry, "sample")
            self._nxclass(sample, "NXsample")
            sample.create_dataset("name", data=np.bytes_(title))

            # /entry/instrument/detector (NXdetector) for raw counts
            detgrp = self._ensure_group(entry, "instrument/detector")
            self._nxclass(detgrp, "NXdetector")
            self._write_1d(detgrp, "det", inp.det, units="counts")
            self._write_1d(detgrp, "det_bg", inp.det_bg, units="counts")

            # /entry/data: NXdata groups for canonical plots
            dataroot = self._ensure_group(entry, "data")
            default_nxdata = None

            r = None if inp.R is None else np.asarray(inp.R)
            q = None if inp.q is None else np.asarray(inp.q)
            rraw = None if inp.Rraw is None else np.asarray(inp.Rraw)
            qraw = None if inp.qraw is None else np.asarray(inp.qraw)
            tth = None if inp.tth is None else np.asarray(inp.tth)
            tthcr = None if inp.tthcr is None else np.asarray(inp.tthcr)
            z = None if inp.z is None else np.asarray(inp.z)
            delta_rho = None if inp.delta_rho is None else np.asarray(inp.delta_rho)
            delta_rho_bl = (
                None if inp.delta_rho_bl is None else np.asarray(inp.delta_rho_bl)
            )
            delta_rho_nc = (
                None if inp.delta_rho_nc is None else np.asarray(inp.delta_rho_nc)
            )

            # R(tth)
            tth_from_q = None
            if (
                r is not None
                and q is not None
                and r.ndim == 1
                and q.ndim == 1
                and r.shape[0] == q.shape[0]
                and qraw is not None
                and tthcr is not None
                and qraw.ndim == 1
                and tthcr.ndim == 1
                and qraw.shape[0] == tthcr.shape[0]
                and qraw.shape[0] > 0
            ):
                wref = 4.0 * np.pi * np.sin(np.deg2rad(tthcr / 2.0))
                valid = np.isfinite(wref) & np.isfinite(qraw) & (qraw != 0)
                if np.any(valid):
                    wavelength = np.median(wref[valid] / qraw[valid])
                    if np.isfinite(wavelength) and wavelength > 0:
                        x = np.clip(q * wavelength / (4.0 * np.pi), -1.0, 1.0)
                        cand = 2.0 * np.rad2deg(np.arcsin(x))
                        if np.all(np.isfinite(cand)):
                            tth_from_q = cand
            if (
                tth_from_q is None
                and r is not None
                and tth is not None
                and r.ndim == 1
                and tth.ndim == 1
                and r.shape[0] == tth.shape[0]
            ):
                tth_from_q = tth

            if (
                r is not None
                and tth_from_q is not None
                and r.ndim == 1
                and tth_from_q.ndim == 1
                and r.shape[0] == tth_from_q.shape[0]
            ):
                rtth = self._ensure_group(dataroot, "R_of_tth")
                self._nxclass(rtth, "NXdata")
                rtth.attrs["signal"] = "R"
                rtth.attrs["axes"] = "tth"
                self._write_1d(rtth, "R", r, units="")  # dimensionless
                self._write_1d(
                    rtth, "tth", tth_from_q, units="deg", long_name="two-theta"
                )
                default_nxdata = "R_of_tth"

            # R(q)
            if (
                r is not None
                and q is not None
                and r.ndim == 1
                and q.ndim == 1
                and r.shape[0] == q.shape[0]
            ):
                rq = self._ensure_group(dataroot, "R_of_q")
                self._nxclass(rq, "NXdata")
                rq.attrs["signal"] = "R"
                rq.attrs["axes"] = "q"
                self._write_1d(rq, "R", r, units="")
                self._write_1d(rq, "q", q, units="1/Angstrom")
                if default_nxdata is None:
                    default_nxdata = "R_of_q"

            # Rraw(qraw)
            if (
                rraw is not None
                and qraw is not None
                and rraw.ndim == 1
                and qraw.ndim == 1
                and rraw.shape[0] == qraw.shape[0]
            ):
                rr = self._ensure_group(dataroot, "Rraw_of_qraw")
                self._nxclass(rr, "NXdata")
                rr.attrs["signal"] = "Rraw"
                rr.attrs["axes"] = "qraw"
                self._write_1d(rr, "Rraw", rraw, units="")
                self._write_1d(rr, "qraw", qraw, units="1/Angstrom")
                if default_nxdata is None:
                    default_nxdata = "Rraw_of_qraw"

            # delta_rho(z)
            if (
                z is not None
                and delta_rho is not None
                and z.ndim == 1
                and delta_rho.ndim == 1
                and z.shape[0] == delta_rho.shape[0]
            ):
                dr = self._ensure_group(dataroot, "delta_rho")
                self._nxclass(dr, "NXdata")
                dr.attrs["signal"] = "delta_rho"
                dr.attrs["axes"] = "z"
                self._write_1d(dr, "z", z, units="Angstrom")
                self._write_1d(dr, "delta_rho", delta_rho, units="e/Angstrom^3")
                if default_nxdata is None:
                    default_nxdata = "delta_rho"

            # delta_rho_bl(z)
            if (
                z is not None
                and delta_rho_bl is not None
                and z.ndim == 1
                and delta_rho_bl.ndim == 1
                and z.shape[0] == delta_rho_bl.shape[0]
            ):
                dr_bl = self._ensure_group(dataroot, "delta_rho_bl")
                self._nxclass(dr_bl, "NXdata")
                dr_bl.attrs["signal"] = "delta_rho_bl"
                dr_bl.attrs["axes"] = "z"
                self._write_1d(dr_bl, "z", z, units="Angstrom")
                self._write_1d(
                    dr_bl, "delta_rho_bl", delta_rho_bl, units="e/Angstrom^3"
                )
                default_nxdata = "delta_rho_bl"

            # delta_rho_nc(z)
            if (
                z is not None
                and delta_rho_nc is not None
                and z.ndim == 1
                and delta_rho_nc.ndim == 1
                and z.shape[0] == delta_rho_nc.shape[0]
            ):
                dr_nc = self._ensure_group(dataroot, "delta_rho_nc")
                self._nxclass(dr_nc, "NXdata")
                dr_nc.attrs["signal"] = "delta_rho_nc"
                dr_nc.attrs["axes"] = "z"
                self._write_1d(dr_nc, "z", z, units="Angstrom")
                self._write_1d(
                    dr_nc, "delta_rho_nc", delta_rho_nc, units="e/Angstrom^3"
                )
                if default_nxdata is None:
                    default_nxdata = "delta_rho_nc"

            # /entry/process (NXprocess) for intermediates/derived arrays
            proc = self._ensure_group(entry, "process")
            self._nxclass(proc, "NXprocess")

            # Angles (corrected)
            self._write_1d(proc, "tthc", inp.tthc, units="deg")
            self._write_1d(proc, "tthcr", inp.tthcr, units="deg")
            self._write_1d(proc, "tthc_bg", inp.tthc_bg, units="deg")
            self._write_1d(proc, "tthcr_bg", inp.tthcr_bg, units="deg")

            # Preprocessed detector signals
            self._write_1d(proc, "detm", inp.detm, units="counts")
            self._write_1d(proc, "detmn", inp.detmn, units="counts")
            self._write_1d(proc, "detmnp", inp.detmnp, units="counts")
            self._write_1d(proc, "detmnpr", inp.detmnpr, units="counts")
            self._write_1d(proc, "detm_bg", inp.detm_bg, units="counts")
            self._write_1d(proc, "detmn_bg", inp.detmn_bg, units="counts")
            self._write_1d(proc, "detmnp_bg", inp.detmnp_bg, units="counts")
            self._write_1d(proc, "detmnpr_bg", inp.detmnpr_bg, units="counts")
            self._write_1d(proc, "detmnpri_bg", inp.detmnpri_bg, units="counts")

            # Corrections & analysis
            self._write_1d(proc, "fpc", inp.fpc, units="")
            self._write_1d(proc, "baseline", inp.baseline, units="")
            self._write_1d(proc, "minima", inp.minima, units="")
            self._write_1d(proc, "phase", inp.phase, units="rad")
            self._write_1d(proc, "phase_bl", inp.phase_bl, units="rad")
            self._write_1d(proc, "Rq4_bl", inp.Rq4_bl, units="")
            self._write_1d(proc, "Rq4_sm", inp.Rq4_sm, units="")

            # Profiles / model outputs
            self._write_1d(proc, "z", inp.z, units="Angstrom")
            self._write_1d(proc, "delta_rho", inp.delta_rho, units="e/Angstrom^3")
            self._write_1d(proc, "delta_rho_bl", inp.delta_rho_bl, units="e/Angstrom^3")
            self._write_1d(proc, "delta_rho_nc", inp.delta_rho_nc, units="e/Angstrom^3")

            if default_nxdata is not None:
                entry.attrs["default"] = "data"
                dataroot.attrs["default"] = default_nxdata
            else:
                entry.attrs["default"] = "process"

        self.outputs.filename = out
