"""Integration tests for the generic XRPD processor."""

import os
import time
from contextlib import contextmanager

from ...bliss_globals import current_session
from ...bliss_globals import setup_globals
from ...xrpd.plots import Xrpd2dIntegrationPlot
from ...xrpd.plots import XrpdCurvePlot
from ...xrpd.plots import XrpdImagePlot
from .. import testing
from ..processors.xrpd import DemoXrpdProcessor


def xrpd_demo_1d(nrepeats: int = None, expo=0.2, npoints=10):
    test_xrpd_ct_with_1d_integration(nrepeats=nrepeats, expo=expo)
    test_xrpd_scan_with_1d_integration(nrepeats=nrepeats, expo=expo, npoints=npoints)
    test_xrpd_scan_with_two_1d_integrations(
        nrepeats=nrepeats, expo=expo, npoints=npoints
    )


def xrpd_demo_2d(nrepeats: int = None, expo=0.2, npoints=10):
    test_xrpd_ct_with_2d_integration(nrepeats=nrepeats, expo=expo)
    test_xrpd_scan_with_2d_integration(nrepeats=nrepeats, expo=expo, npoints=npoints)


@testing.integration_fixture
def difflab6_processor_1d(gc_collect):
    xrpd_processor = DemoXrpdProcessor()
    xrpd_processor.enable(setup_globals.difflab6)
    xrpd_processor._plotter.clear_lima_plots("difflab6")
    yield xrpd_processor
    xrpd_processor.disable()


@testing.integration_fixture
def diffcam_processor_1d(gc_collect):
    # Ensure this processor can exist in parallel to difflab6_processor_1d
    config = {
        "name": "diffcam_processor",
        "plugin": "generic",
        "class": "DemoXrpdProcessor",
        "package": "blissoda.demo.processors.xrpd.processor",
        "shared": False,
        "singleton": False,
    }
    xrpd_processor = DemoXrpdProcessor(config=config)
    xrpd_processor.enable(setup_globals.diffcam)
    # TODO: produce different data or fix PONI
    # xrpd_processor.integration_options["mask"] = numpy.zeros((540, 516), dtype=int).tolist()
    xrpd_processor._plotter.clear_lima_plots("diffcam")
    yield xrpd_processor
    xrpd_processor.disable()


@testing.integration_fixture
def difflab6_processor_2d(gc_collect):
    xrpd_processor = DemoXrpdProcessor()
    xrpd_processor.enable(setup_globals.difflab6)
    xrpd_processor.integration_options["nbpt_azim"] = 360
    xrpd_processor._plotter.clear_lima_plots("difflab6")
    yield xrpd_processor
    xrpd_processor.integration_options.pop("nbpt_azim", None)
    xrpd_processor.disable()


@testing.integration_test
def test_xrpd_ct_with_1d_integration(
    difflab6_processor_1d, nrepeats: int = None, expo=0.2
):
    _test_xrpd_ct_with_1d_integration(
        difflab6_processor_1d,
        [setup_globals.difflab6, setup_globals.diode1, setup_globals.diode2],
        nrepeats=nrepeats,
        expo=expo,
    )


@testing.integration_test
def test_xrpd_ct_with_2d_integration(
    difflab6_processor_2d, nrepeats: int = None, expo=0.2
):
    _test_xrpd_ct_with_2d_integration(
        difflab6_processor_2d,
        [setup_globals.difflab6, setup_globals.diode1, setup_globals.diode2],
        nrepeats=nrepeats,
        expo=expo,
    )


@testing.integration_test
def test_xrpd_scan_with_1d_integration(
    difflab6_processor_1d, nrepeats: int = None, expo=0.2, npoints=10
):
    _test_xrpd_scan_with_1d_integration(
        [difflab6_processor_1d],
        [
            setup_globals.difflab6,
            setup_globals.diode1,
            setup_globals.diode2,
        ],
        nrepeats=nrepeats,
        expo=expo,
        npoints=npoints,
    )


@testing.integration_test
def test_xrpd_scan_with_2d_integration(
    difflab6_processor_2d, nrepeats: int = None, expo=0.2, npoints=10
):
    _test_xrpd_scan_with_2d_integration(
        [difflab6_processor_2d],
        [
            setup_globals.difflab6,
            setup_globals.diode1,
            setup_globals.diode2,
        ],
        nrepeats=nrepeats,
        expo=expo,
        npoints=npoints,
    )


@testing.integration_test
def test_xrpd_scan_with_two_1d_integrations(
    difflab6_processor_1d,
    diffcam_processor_1d,
    nrepeats: int = None,
    expo=0.2,
    npoints=10,
):
    _test_xrpd_scan_with_1d_integration(
        [difflab6_processor_1d, diffcam_processor_1d],
        [
            setup_globals.difflab6,
            setup_globals.diffcam,
            setup_globals.diode1,
            setup_globals.diode2,
        ],
        nrepeats=nrepeats,
        expo=expo,
        npoints=npoints,
    )


def _test_xrpd_ct_with_1d_integration(
    processor, detectors, nrepeats: int = None, expo=0.2
):
    if nrepeats is None:
        nrepeats = processor.number_of_scans + 1

    with _scan_context() as scans:
        for _ in range(nrepeats):
            scan = processor.pct(expo, *detectors)
            scans.append(scan)
            # Workflow is getting data from Lima directly so we cannot start
            # the next scan until the workflow is done.
            testing.wait_workflows()

    for lima_name in processor.lima_names:
        scan_legends = [_get_scan_legend(scans[-1], lima_name)]
        _assert_1d_plot_accum(processor._plotter, scan_legends, lima_name)

    nremoved = max(len(scans) - processor.number_of_scans, 0)
    scan_legends = [
        _get_scan_legend(scan, lima_name)
        for scan in scans[nremoved:]
        for lima_name in processor.lima_names
    ]
    _assert_1d_plot_last(processor._plotter, scan_legends)


def _test_xrpd_ct_with_2d_integration(
    processor, detectors, nrepeats: int = None, expo=0.2
):
    if nrepeats is None:
        nrepeats = processor.number_of_scans + 1

    with _scan_context() as scans:
        for _ in range(nrepeats):
            scan = processor.pct(expo, *detectors)
            scans.append(scan)
            # Workflow is getting data from Lima directly so we cannot start
            # the next scan until the workflow is done.
            testing.wait_workflows()

    for lima_name in processor.lima_names:
        scan_legends = [_get_scan_legend(scans[-1], lima_name)]
        _assert_2d_plot_last(processor._plotter, scan_legends, lima_name)


def _test_xrpd_scan_with_1d_integration(
    processors,
    detectors,
    nrepeats: int = None,
    expo=0.2,
    npoints=10,
):
    if nrepeats is None:
        nrepeats = min(p.number_of_scans for p in processors) + 1

    with _scan_context() as scans:
        for _ in range(nrepeats):
            scan = setup_globals.loopscan(npoints, expo, *detectors)
            scans.append(scan)

    testing.wait_workflows()

    for processor in processors:
        for lima_name in processor.lima_names:
            scan_legends = [_get_scan_legend(scans[-1], lima_name)]
            _assert_1d_plot_accum(processor._plotter, scan_legends, lima_name)

        for scan in scans:
            for lima_name in processor.lima_names:
                _assert_1d_data(lima_name, processor, scan, npoints)

    scan_legends = [
        _get_scan_legend(scan, lima_name)
        for scan in scans[max(len(scans) - processor.number_of_scans, 0) :]
        for processor in processors
        for lima_name in processor.lima_names
    ]
    _assert_1d_plot_last(processor._plotter, scan_legends)


def _test_xrpd_scan_with_2d_integration(
    processors, detectors, nrepeats: int = None, expo=0.2, npoints=10
):
    if nrepeats is None:
        nrepeats = min(p.number_of_scans for p in processors) + 1

    with _scan_context() as scans:
        for _ in range(nrepeats):
            scan = setup_globals.loopscan(npoints, expo, *detectors)
            scans.append(scan)

    testing.wait_workflows()

    for processor in processors:
        for lima_name in processor.lima_names:
            scan_legends = [_get_scan_legend(scans[-1], lima_name)]
            _assert_2d_plot_last(processor._plotter, scan_legends, lima_name)

        for scan in scans:
            _assert_2d_data(processor, scan, npoints)


def _assert_1d_plot_last(plotter, scan_legends):
    """The last 1D diffractogram of each scan is plotted as a curve."""
    _assert_plot(plotter, scan_legends, "Integrated (Last)", XrpdCurvePlot)


def _assert_1d_plot_accum(plotter, scan_legends, lima_name):
    """The 1D diffractograms of the last scan are plotted as an image."""
    _assert_plot(plotter, scan_legends, f"Integrated {lima_name}", XrpdImagePlot)


def _assert_2d_plot_last(plotter, scan_legends, lima_name):
    """The last 2D cake of the last scan is plotted as an image."""
    _assert_plot(
        plotter,
        scan_legends,
        f"2D Integrated {lima_name} (last)",
        Xrpd2dIntegrationPlot,
    )


@testing.demo_assert("Check {plot_id} plot for {scan_legends}")
def _assert_plot(
    plotter, scan_legends, plot_id, plot_cls, retry_timeout=10, retry_period=0.2
):
    image_plot = plotter._get_plot(plot_id, plot_cls)
    t0 = time.time()
    while True:
        labels = image_plot.get_labels()
        if set(labels) == set(scan_legends):
            return
        dt = time.time() - t0
        if dt > retry_timeout:
            break
        time.sleep(retry_period)
    assert set(labels) == set(scan_legends), labels


def _get_scan_legend(scan, lima_name):
    filename = scan.scan_info["filename"]
    if not filename:
        filename = current_session.scan_saving.filename
    scan_nb = scan.scan_info["scan_nb"]
    scan_name = os.path.splitext(os.path.basename(filename))[0]
    return f"{scan_name}: {scan_nb}.1 {scan.name} ({lima_name})"


@testing.demo_assert("Check 1D integration data")
def _assert_1d_data(lima_name, processor, scan, npoints):
    data_keys = processor.get_data_keys(
        scan, lima_name, retry_timeout=10, retry_period=0.2
    )

    shapes = {
        name: processor.get_data(scan, name, retry_timeout=10, retry_period=0.2).shape
        for name in data_keys
    }
    expected = {
        f"{lima_name}:q": (4096,),
        f"{lima_name}:intensity": (npoints, 4096),
        f"{lima_name}:points": (npoints,),
    }
    assert shapes == expected, shapes


@testing.demo_assert("Check 2D integration data")
def _assert_2d_data(processor, scan, npoints):
    data_keys = processor.get_data_keys(
        scan, "difflab6", retry_timeout=10, retry_period=0.2
    )

    shapes = {
        name: processor.get_data(scan, name, retry_timeout=10, retry_period=0.2).shape
        for name in data_keys
    }
    expected = {
        "difflab6:chi": (360,),
        "difflab6:q": (4096,),
        "difflab6:intensity": (npoints, 360, 4096),
        "difflab6:points": (npoints,),
    }
    assert shapes == expected, shapes


@contextmanager
def _scan_context():
    scans = []
    try:
        yield scans
    finally:
        # Holds references to processor instances
        for scan in scans:
            scan._Scan__user_scan_meta = None
