from pathlib import Path

import h5py
import numpy

from ...xrffit.batch._save import XrfResultWriter
from ...xrffit.fit.types import QuantificationParameters
from ...xrffit.fit.types import XRFBatchFitResult
from ...xrffit.pymca_config import PyMcaXrfConfiguration
from ..utils import h5content


def test_xrf_result_writer(tmp_path: Path):
    result = _generate_example_xrf_batch_fit()
    filename = tmp_path / "test.h5"
    configuration = PyMcaXrfConfiguration()

    with XrfResultWriter(f"{filename}::/1.1", default_group="fit") as writer:
        writer.save_configuration(configuration)

        for i in range(2):
            start = i * result.num_spectra
            stop = start + result.num_spectra
            writer.save_fit_result(result, start, stop)

    with h5py.File(filename) as f:
        actual = h5content(f)

    expected = {
        "1.1": {
            "@attrs": {"NX_class", "default"},
            "fit": {
                "@attrs": {"NX_class", "default"},
                "configuration": {
                    "@attrs": {"NX_class"},
                    "data@shape": (),
                    "date@shape": (),
                    "type@shape": (),
                },
                "program@shape": (),
                "quantification": {
                    "@attrs": {"NX_class"},
                    "DetectorArea@attrs": {"units"},
                    "DetectorArea@shape": (),
                    "DetectorDistance@attrs": {"units"},
                    "DetectorDistance@shape": (),
                    "Flux@attrs": {"units"},
                    "Flux@shape": (),
                    "I0@shape": (),
                    "ReferenceElement@shape": (),
                    "ReferenceTransitions@shape": (),
                    "SolidAngle@attrs": {"units"},
                    "SolidAngle@shape": (),
                    "Time@attrs": {"units"},
                    "Time@shape": (),
                },
                "results": {
                    "@attrs": {"NX_class", "default"},
                    "derivatives": {
                        "@attrs": {"NX_class", "axes"},
                        "Cu_Ka@attrs": {"interpretation"},
                        "Cu_Ka@shape": (1024,),
                        "Fe_Ka@attrs": {"interpretation"},
                        "Fe_Ka@shape": (1024,),
                        "Zn_Ka@attrs": {"interpretation"},
                        "Zn_Ka@shape": (1024,),
                        "energy@shape": (1024,),
                        "gain@attrs": {"interpretation"},
                        "gain@shape": (1024,),
                        "zero@attrs": {"interpretation"},
                        "zero@shape": (1024,),
                    },
                    "fit": {
                        "@attrs": {"NX_class", "auxiliary_signals", "axes", "signal"},
                        "energy@shape": (1024,),
                        "data@shape": (4, 1024),
                        "data@attrs": {"interpretation"},
                        "model@shape": (4, 1024),
                        "model@attrs": {"interpretation"},
                        "residuals@shape": (4, 1024),
                        "residuals@attrs": {"interpretation"},
                    },
                    "mass_fractions": {
                        "@attrs": {"NX_class", "auxiliary_signals", "signal"},
                        "Cu@shape": (4,),
                        "Fe@shape": (4,),
                        "Zn@shape": (4,),
                    },
                    "mass_fractions_layer1": {
                        "@attrs": {"NX_class", "auxiliary_signals", "signal"},
                        "Cu@shape": (4,),
                        "Fe@shape": (4,),
                        "Zn@shape": (4,),
                    },
                    "parameters": {
                        "@attrs": {"NX_class", "auxiliary_signals", "signal"},
                        "Cu_Ka@shape": (4,),
                        "Fe_Ka@shape": (4,),
                        "Zn_Ka@shape": (4,),
                        "gain@shape": (4,),
                        "zero@shape": (4,),
                    },
                    "uncertainties": {
                        "@attrs": {"NX_class", "auxiliary_signals", "signal"},
                        "Cu_Ka@shape": (4,),
                        "Fe_Ka@shape": (4,),
                        "Zn_Ka@shape": (4,),
                        "gain@shape": (4,),
                        "zero@shape": (4,),
                    },
                },
                "version@shape": (),
            },
        },
        "@attrs": {"NX_class", "default"},
    }
    assert actual == expected


def _generate_example_xrf_batch_fit() -> XRFBatchFitResult:
    rng = numpy.random.default_rng(seed=42)

    Nspectra = 2
    Nchan = 1024
    Nfit = 800

    raw_channels = numpy.arange(Nchan)
    raw_energies = raw_channels * 0.010
    channels = numpy.linspace(50, 850, Nfit).astype(int)
    energies = raw_energies[channels]

    spectra = rng.poisson(
        lam=50 + 500 * numpy.exp(-0.5 * ((raw_energies - 6.4) / 0.15) ** 2),
        size=(Nspectra, Nchan),
    )
    model = numpy.zeros_like(spectra, dtype=float)
    for i in range(Nspectra):
        model[i] = (
            45
            + 480 * numpy.exp(-0.5 * ((raw_energies - 6.4) / 0.14) ** 2)
            + 200 * numpy.exp(-0.5 * ((raw_energies - 8.0) / 0.20) ** 2)
        )
    residuals = spectra - model

    parameter_names = ["Fe_Ka", "Cu_Ka", "Zn_Ka", "zero", "gain"]
    parameters = numpy.array(
        [[1200.0, 1150.0], [800.0, 820.0], [450.0, 430.0], [40.0, 42.0], [0.010, 0.011]]
    )
    uncertainties = parameters * 0.05

    mass_fraction_names = ["Fe", "Cu", "Zn"]
    mass_fractions = numpy.array([[0.55, 0.53], [0.30, 0.32], [0.15, 0.15]])
    mass_fractions_per_layer = [mass_fractions.copy()]

    Nparameters = len(parameter_names)
    derivatives = rng.normal(scale=1e-3, size=(Nparameters, Nchan))

    quant_params = QuantificationParameters(
        I0=1.0e10,
        Time=10.0,
        Flux=1.0e9,
        DetectorDistance=5.0,
        DetectorArea=0.5,
        SolidAngle=0.02,
        ReferenceElement="Fe",
        ReferenceTransitions="K",
    )

    return XRFBatchFitResult(
        parameters=parameters,
        uncertainties=uncertainties,
        derivatives=derivatives,
        parameter_names=parameter_names,
        channels=channels,
        energies=energies,
        raw_channels=raw_channels,
        raw_energies=raw_energies,
        spectra=spectra,
        model=model,
        residuals=residuals,
        mass_fractions=mass_fractions,
        mass_fractions_per_layer=mass_fractions_per_layer,
        mass_fraction_names=mass_fraction_names,
        quantification_parameters=quant_params,
    )
