from datetime import datetime
from datetime import timezone
from pathlib import Path

import h5py
import numpy
import pytest

from ewoksid01.fit import gauss
from ewoksid01.tasks.cdi_plot import CdiPlot

N_FRAMES = 30
AXIS_NAME = "lab_piy"
RC_AXIS_NAME = "phi"
COUNTER_NAME = "mpx1x4_target"

SCANS = {
    "1.1": {"center": N_FRAMES // 2 - 2, "scale": 10, "rc": 8.1},
    "2.1": {"center": N_FRAMES // 2 + 2, "scale": 8, "rc": 10.2},
}


def create_scan_file(root_dir: Path) -> Path:
    with h5py.File(root_dir / "sample.h5", "w") as h5file:
        for scan_nb, scan_params in SCANS.items():
            scan = h5file.create_group(scan_nb)
            scan.create_dataset("title", data="ascan")
            scan.create_dataset(
                "start_time", data=datetime.now(tz=timezone.utc).isoformat()
            )
            instr = scan.create_group("instrument")
            positioners = instr.create_group("positioners")
            positioners.create_dataset(RC_AXIS_NAME, data=scan_params["rc"])

            meas = scan.create_group("measurement")

            X = numpy.arange(N_FRAMES)
            group = instr.create_group(AXIS_NAME)
            dset = group.create_dataset(AXIS_NAME, data=X)
            meas[AXIS_NAME] = h5py.SoftLink(dset.name)

            group = instr.create_group(COUNTER_NAME)
            dset = group.create_dataset(
                COUNTER_NAME,
                data=gauss(
                    X,
                    H=2,
                    A=scan_params["scale"],
                    x0=scan_params["center"],
                    sigma=N_FRAMES / 3,
                ),
            )
            meas[COUNTER_NAME] = h5py.SoftLink(dset.name)

    return root_dir / "sample.h5"


@pytest.fixture
def scan_file(tmp_path: Path) -> Path:
    return create_scan_file(tmp_path)


def test_task(tmp_path: Path, scan_file: Path):
    task = CdiPlot(
        inputs={
            "filename": str(scan_file),
            "scan_nos": [1, 2],
            "axis_name": AXIS_NAME,
            "rc_axis_name": RC_AXIS_NAME,
            "counter": COUNTER_NAME,
            "output_folder": str(tmp_path),
        }
    )

    task.execute()
    for fig_path in task.outputs.saved_figs:
        Path(fig_path).exists()

    with h5py.File(scan_file, "r") as output:
        cdi_fits = output["1-2/cdi_fits"]
        assert isinstance(cdi_fits, h5py.Group)
        assert "intensity_fits" in cdi_fits
        assert "offset_fit" in cdi_fits
