import time
import pytest

from ophyd import (
    Component as Cpt,
    SimDetector,
    SimDetectorCam,
    StatsPlugin,
    EpicsMotor,
    ROIPlugin,
    ProcessPlugin,
    Device,
)
from ophyd.areadetector.base import EpicsSignalWithRBV
from ophyd.flyers import (
    AreaDetectorTimeseriesCollector,
    WaveformCollector,
    MonitorFlyerMixin,
)
from ophyd.status import wait


@pytest.fixture(scope="function")
def ts_sim_detector(ad_prefix):
    class Detector(SimDetector):
        acquire = Cpt(EpicsSignalWithRBV, "cam1:Acquire", trigger_value=1)
        cam = Cpt(SimDetectorCam, "cam1:")
        ts_col = Cpt(AreaDetectorTimeseriesCollector, "Stats1:")
        stats = Cpt(StatsPlugin, "Stats1:")
        roi1 = Cpt(ROIPlugin, "ROI1:")
        proc1 = Cpt(ProcessPlugin, "Proc1:")

    try:
        det = Detector(ad_prefix, name="sim")
        det.wait_for_connection(timeout=1.0)
    except TimeoutError:
        pytest.skip("IOC unavailable")
    return det


@pytest.fixture
def tscollector(ts_sim_detector):
    return ts_sim_detector.ts_col


@pytest.mark.adsim
def test_ad_time_series(ts_sim_detector, tscollector):
    sim_detector = ts_sim_detector

    num_points = 3

    cam = sim_detector.cam
    cam.stage_sigs[cam.acquire_time] = 0.001
    cam.stage_sigs[cam.acquire_period] = 0.001
    cam.stage_sigs[cam.image_mode] = "Single"
    cam.stage_sigs[cam.trigger_mode] = "Internal"

    print("tscollector desc", tscollector.describe())
    print("tscollector flyer desc", tscollector.describe_collect())
    print("tscollector repr", repr(tscollector))
    print("simdet stage sigs", sim_detector.stage_sigs)
    print("tscoll stage sigs", tscollector.stage_sigs)
    print("cam stage sigs", cam.stage_sigs)
    print("stats stage sigs", sim_detector.stats.stage_sigs)

    try:
        # In case the collector is currently running...
        wait(tscollector.complete())
    except RuntimeError:
        ...

    tscollector.stage_sigs[tscollector.num_points] = num_points

    sim_detector.stage()
    tscollector.kickoff()

    for i in range(num_points):
        st = sim_detector.trigger()
        wait(st)
        print(st)
        time.sleep(0.1)

    tscollector.pause()
    tscollector.resume()

    st = tscollector.complete()
    wait(st)

    collected = list(tscollector.collect())
    print("collected", collected)
    sim_detector.unstage()
    assert len(collected) == num_points
    # TODO any more validation here?


@pytest.fixture(scope="function")
def wf_sim_detector(ad_prefix):
    suffix = "??TODO??"

    class Detector(SimDetector):
        wfcol = Cpt(WaveformCollector, suffix)

    try:
        det = Detector(ad_prefix, name="det")
        det.wait_for_connection(timeout=1.0)
    except TimeoutError:
        pytest.skip("IOC unavailable")
    return det


@pytest.fixture
def wfcol(wf_sim_detector):
    return wf_sim_detector.wfcol


@pytest.mark.adsim
def test_waveform(wf_sim_detector, wfcol):
    print("waveform collector", wfcol)
    print("wfcol flyer desc", wfcol.describe_collect())


@pytest.mark.motorsim
@pytest.mark.parametrize("pivot", [True, False])
def test_monitor_flyer(pivot):
    class BasicDevice(Device):
        mtr1 = Cpt(EpicsMotor, "XF:31IDA-OP{Tbl-Ax:X2}Mtr")
        mtr2 = Cpt(EpicsMotor, "XF:31IDA-OP{Tbl-Ax:X3}Mtr")

    class FlyerDevice(MonitorFlyerMixin, BasicDevice):
        pass

    fdev = FlyerDevice(
        "", name="fdev", stream_names={"mtr1.user_readback": "oranges"}, pivot=pivot
    )
    fdev.wait_for_connection()

    fdev.monitor_attrs = ["mtr1.user_readback", "mtr2.user_readback"]
    fdev.describe()

    st = fdev.kickoff()
    wait(st)

    mtr1, mtr2 = fdev.mtr1, fdev.mtr2
    rbv1, rbv2 = mtr1.position, mtr2.position
    fdev.mtr1.move(rbv1 + 0.2, wait=True)
    fdev.mtr2.move(rbv2 + 0.2, wait=True)

    fdev.pause()

    fdev.mtr1.move(rbv1 - 0.2, wait=True)
    fdev.mtr2.move(rbv2 - 0.2, wait=True)

    fdev.resume()
    st = fdev.complete()
    wait(st)

    print(fdev.describe_collect())

    desc1 = fdev.mtr1.user_readback.describe()
    desc2 = fdev.mtr2.user_readback.describe()

    if not pivot:
        desc1[fdev.mtr1.user_readback.name]["dtype"] = "array"
        desc2[fdev.mtr2.user_readback.name]["dtype"] = "array"

    assert fdev.describe_collect() == {
        "oranges": desc1,
        "fdev_mtr2": desc2,
    }
    data = list(fdev.collect())
    print("collected data", data)
    # data from both motors

    if not pivot:
        assert len(data) == 2
        d1 = data[0]["data"]["fdev_mtr1"]
        d2 = data[1]["data"]["fdev_mtr2"]

    else:
        assert len(data) >= 2

        d1 = [d["data"]["fdev_mtr1"] for d in data if "fdev_mtr1" in d["data"]]
        d2 = [d["data"]["fdev_mtr2"] for d in data if "fdev_mtr2" in d["data"]]

    # and at least more than one data point...
    assert len(d1) > 1
    assert len(d2) > 1

    print("data1", d1)
    print("data2", d2)
    # raise ValueError()
