"""Tests for the smart and fast stacking."""

import logging
from random import randint
from typing import Optional

import numpy as np
import pytest
from hypothesis import given
from hypothesis import strategies as st

from labthings_fastapi.testing import create_thing_without_server

from openflexure_microscope_server.scan_directories import IMAGE_REGEX
from openflexure_microscope_server.things.autofocus import (
    EXTRA_STACK_CAPTURES,
    MAX_TEST_IMAGE_COUNT,
    MIN_TEST_IMAGE_COUNT,
    AutofocusThing,
    CaptureInfo,
    NotAPeakError,
    SmartStackParams,
    _count_turning_points,
    _get_capture_by_id,
    _get_capture_index_by_id,
    _get_peak_turning_point,
)

RANDOM_GENERATOR = np.random.default_rng()


def odd_integers(min_value=0, max_value=1000):
    """Return a hypothesis strategy for odd integers."""
    min_base = (min_value) // 2
    max_base = (max_value - 1) // 2
    # Ensure the range allows at least one odd number
    if min_base > max_base:
        return st.nothing()
    return st.integers(min_value=min_base, max_value=max_base).map(lambda x: 2 * x + 1)


def even_integers(min_value=0, max_value=1000):
    """Return a hypothesis strategy for even integers."""
    min_base = (min_value + 1) // 2
    max_base = (max_value) // 2
    # Ensure the range allows at least one even number
    if min_base > max_base:
        return st.nothing()
    return st.integers(min_value=min_base, max_value=max_base).map(lambda x: 2 * x)


@given(
    save_ims=odd_integers(min_value=1, max_value=9),
    extra_ims=even_integers(min_value=0, max_value=10),
)
def test_stack_params_validation(save_ims, extra_ims):
    """Test the validation of the image numbers for a stack for valid combinations.

    save_ims is the number to save (must be odd and positive)
    extra_ims is how many more images there are in min_images_to_test than
    images_to_save. (even so that, min_images_to_test is odd and larger than
    images_to_save
    """
    # Coerce min_images_to_test as the max extra ims depends on save_ims so is hard
    # to do automatically in hypothesis. This clamps the number between 3 and 9.
    min_images_to_test = max(min(save_ims + extra_ims, 9), 3)
    SmartStackParams(
        stack_dz=50,
        images_to_save=save_ims,
        min_images_to_test=min_images_to_test,
        autofocus_dz=2000,
        images_dir="/this/is/fake",
        save_resolution=(1640, 1232),
    )


@given(
    save_ims=odd_integers(min_value=0, max_value=10),
    extra_ims=even_integers(min_value=-10, max_value=-1),
)
def test_stack_params_not_enough_test_images(save_ims, extra_ims):
    """Test error is raised if min_images_to_test is smaller than images_to_save.

    ``extra_ims`` negative so that min_images_to_test is smaller than images_to_save.

    For arguments see test_stack_params_validation
    """
    # Depending on the values multiple messages are possible
    match = (
        "(Can't test for focus with fewer than 3 images|"
        "Can't save more images than the minimum number tested)"
    )
    with pytest.raises(ValueError, match=match):
        SmartStackParams(
            stack_dz=50,
            images_to_save=save_ims,
            min_images_to_test=save_ims + extra_ims,
            autofocus_dz=2000,
            images_dir="/this/is/fake",
            save_resolution=(1640, 1232),
        )


@given(
    save_ims=odd_integers(min_value=-10, max_value=-1),
    extra_ims=even_integers(min_value=0, max_value=10),
)
def test_stack_params_negative_images_to_save(save_ims, extra_ims):
    """save_ims is negative so images_to_save is negative, failing validation.

    For arguments see test_stack_params_validation
    """
    # Depending on the values multiple messages are possible
    match = (
        "(Can't test for focus with fewer than 3 images|"
        "Images to save must be positive and odd)"
    )
    with pytest.raises(ValueError, match=match):
        SmartStackParams(
            stack_dz=50,
            images_to_save=save_ims,
            min_images_to_test=save_ims + extra_ims,
            autofocus_dz=2000,
            images_dir="/this/is/fake",
            save_resolution=(1640, 1232),
        )


@given(
    save_ims=odd_integers(min_value=0, max_value=10),
    extra_ims=odd_integers(min_value=0, max_value=10),
)
def test_even_min_images_to_test(save_ims, extra_ims):
    """extra_ims is odd so min_images_to_test is even, failing validation.

    For arguments see test_stack_params_validation
    """
    # Depending on the values multiple messages are possible
    match = (
        "(Can't test for focus with fewer than 3 images|"
        "Testing with more than 9 images is|"  # may give more than 9 which errors first
        "Minimum number of images to test should be positive and odd)"
    )
    with pytest.raises(ValueError, match=match):
        SmartStackParams(
            stack_dz=50,
            images_to_save=save_ims,
            min_images_to_test=save_ims + extra_ims,
            autofocus_dz=2000,
            images_dir="/this/is/fake",
            save_resolution=(1640, 1232),
        )


@given(
    save_ims=even_integers(min_value=0, max_value=10),
    extra_ims=odd_integers(min_value=0, max_value=10),
)
def test_even_images_to_save(save_ims, extra_ims):
    """save_ims is even so images_to_save is even, failing validation.

    For arguments see test_stack_params_validation
    """
    match = (
        "(Can't test for focus with fewer than 3 images|"
        "Images to save must be positive and odd)"
    )
    with pytest.raises(ValueError, match=match):
        SmartStackParams(
            stack_dz=50,
            images_to_save=save_ims,
            min_images_to_test=save_ims + extra_ims,
            autofocus_dz=2000,
            images_dir="/this/is/fake",
            save_resolution=(1640, 1232),
        )


def test_computed_stack_params():
    """Test SmartStackParams computed properties are as expected.

    Not using hypothesis or we will just copy in the same formulas.
    """
    stack_parameters = SmartStackParams(
        stack_dz=50,
        images_to_save=5,
        min_images_to_test=9,
        autofocus_dz=2000,
        images_dir="/this/is/fake",
        save_resolution=(1640, 1232),
    )

    assert stack_parameters.stack_z_range == 8 * 50

    assert stack_parameters.steps_undershoot == stack_parameters.img_undershoot * 50

    assert stack_parameters.max_images_to_test == 9 + 15

    sharpnesses = [50, 77, 234, 324, 390, 496, 569, 454, 333, 222, 178, 70]
    max_ind = np.argmax(sharpnesses)
    slice_to_save = stack_parameters.slice_to_save(max_ind)
    # Check the slice corresponds to the index for the 5 images centred on the sharpest
    assert sharpnesses[slice_to_save] == [390, 496, 569, 454, 333]

    # For a failed smart stack the slice may be truncated. Check it truncates correctly
    sharpnesses = [496, 569, 454, 333, 222, 178, 70, 69, 66, 50, 45, 40]
    max_ind = np.argmax(sharpnesses)
    slice_to_save = stack_parameters.slice_to_save(max_ind)
    # Check the slice corresponds to the index for the first 4 images
    assert sharpnesses[slice_to_save] == [496, 569, 454, 333]

    # And again for sharpest at the end
    sharpnesses = [13, 21, 26, 31, 39, 49, 50, 77, 234, 324, 390, 496, 569]
    max_ind = np.argmax(sharpnesses)
    slice_to_save = stack_parameters.slice_to_save(max_ind)
    # Check the slice corresponds to the index for the final 3 images
    assert sharpnesses[slice_to_save] == [390, 496, 569]


def random_capture(set_id: Optional[int] = None):
    """Create a capture with random values.

    :param set_id: Optional, use to set a fixed id rather than a random one
    """
    buffer_id = set_id if set_id is not None else randint(0, 1000)
    return CaptureInfo(
        buffer_id=buffer_id,
        position={
            "x": randint(-100000, 100000),
            "y": randint(-100000, 100000),
            "z": randint(-100000, 100000),
        },
        sharpness=randint(0, 100000),
    )


def test_capture_filename_matches_regex():
    """For 100 random captures check the image always matches the regex."""
    for _ in range(100):
        assert IMAGE_REGEX.search(random_capture().filename)


@given(st.integers(min_value=0, max_value=5000))
def test_retrieval_of_captures(start):
    """For 20 random captures, check each can be retrieved correctly by id."""
    captures = [random_capture(start + i) for i in range(20)]

    for i, capture in enumerate(captures):
        buffer_id = capture.buffer_id
        assert _get_capture_index_by_id(captures, buffer_id) == i
        assert _get_capture_by_id(captures, buffer_id) is capture

    # Check errors are raised when supplying ids that aren't in the list
    with pytest.raises(ValueError, match="No capture has a buffer id of"):
        _get_capture_index_by_id(captures, start - 1)
    with pytest.raises(ValueError, match="No capture has a buffer id of"):
        _get_capture_index_by_id(captures, start + 21)
    with pytest.raises(ValueError, match="No capture has a buffer id of"):
        _get_capture_by_id(captures, start - 1)
    with pytest.raises(ValueError, match="No capture has a buffer id of"):
        _get_capture_by_id(captures, start + 21)


@pytest.fixture
def autofocus_thing():
    """Return an autofocus thing connected to a server."""
    return create_thing_without_server(AutofocusThing, mock_all_slots=True)


def test_create_stack(autofocus_thing, caplog):
    """Run create stack with default values and check there is no coercion or logging."""
    initial_min_images_to_test = autofocus_thing.stack_min_images_to_test
    initial_images_to_save = autofocus_thing.stack_images_to_save
    with caplog.at_level(logging.INFO):
        stack_params = autofocus_thing.create_stack_params(
            autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232)
        )

    assert len(caplog.records) == 0
    assert autofocus_thing.stack_min_images_to_test == initial_min_images_to_test
    assert autofocus_thing.stack_images_to_save == initial_images_to_save
    assert stack_params.min_images_to_test == autofocus_thing.stack_min_images_to_test
    assert stack_params.images_to_save == autofocus_thing.stack_images_to_save


@pytest.mark.parametrize(
    ("initial_test_ims", "coerced_test_ims", "expected_log_start"),
    [
        (0, MIN_TEST_IMAGE_COUNT, "Cannot test only 0"),
        (-1, MIN_TEST_IMAGE_COUNT, "Cannot test only -1"),
        (10, MAX_TEST_IMAGE_COUNT, "Testing 10 images will"),
        (4, 5, "Minimum number of images to test should be odd"),
    ],
)
def test_coercing_stack_test_ims(
    initial_test_ims, coerced_test_ims, expected_log_start, autofocus_thing, caplog
):
    """Run create stack with images to test set to values requiring coercion, and check result."""
    autofocus_thing.stack_min_images_to_test = initial_test_ims

    with caplog.at_level(logging.WARNING):
        stack_params = autofocus_thing.create_stack_params(
            autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232)
        )

    assert len(caplog.records) == 1
    assert str(caplog.records[0].msg).startswith(expected_log_start)
    # Check the value is coerced in the stack_params
    assert stack_params.min_images_to_test == coerced_test_ims
    # Check that the setting in the Thing was updated to the coerced value
    assert stack_params.min_images_to_test == autofocus_thing.stack_min_images_to_test


@pytest.mark.parametrize(
    ("initial_save_ims", "coerced_save_ims", "expected_log_start"),
    [
        (0, 1, "At least 1 images must be saved"),
        (-1, 1, "At least 1 images must be saved"),
        (10, 9, "Cannot save 10 images"),
        (4, 5, "Images to save should be odd, setting to 5"),
    ],
)
def test_coercing_stack_save_ims(
    initial_save_ims, coerced_save_ims, expected_log_start, autofocus_thing, caplog
):
    """Run create stack with images to save set to values requiring coercion, and check result."""
    autofocus_thing.stack_images_to_save = initial_save_ims

    with caplog.at_level(logging.WARNING):
        stack_params = autofocus_thing.create_stack_params(
            autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232)
        )

    assert len(caplog.records) == 1
    assert str(caplog.records[0].msg).startswith(expected_log_start)
    # Check the value is coerced in the stack_params
    assert stack_params.images_to_save == coerced_save_ims
    # Check that the setting in the Thing was updated to the coerced value
    assert stack_params.images_to_save == autofocus_thing.stack_images_to_save


@pytest.mark.parametrize("pass_on", [1, 2, 3, 4])
def test_run_smart_stack(pass_on, autofocus_thing, mocker):
    """Test Running smart stack with the stack passing on different attempts."""
    stack_params = autofocus_thing.create_stack_params(
        autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232)
    )
    assert stack_params.max_attempts == 3

    # Set up returns from z-stack
    fake_captures = [
        CaptureInfo(
            buffer_id="first", position={"x": 0, "y": 0, "z": -99}, sharpness=123
        ),
        CaptureInfo(
            buffer_id="pick_me", position={"x": 0, "y": 0, "z": 555}, sharpness=456
        ),
        CaptureInfo(
            buffer_id="last", position={"x": 0, "y": 0, "z": 999}, sharpness=123
        ),
    ]

    successful_return = (True, fake_captures, "pick_me")
    failed_return = (False, fake_captures, "pick_me")
    return_list = [failed_return] * (pass_on - 1) + [successful_return]

    # Mock z_stack and looping_autofocus
    autofocus_thing.z_stack = mocker.Mock(side_effect=return_list)
    autofocus_thing.looping_autofocus = mocker.Mock()

    # Run it
    success, final_z = autofocus_thing.run_smart_stack(
        stack_parameters=stack_params,
        save_on_failure=False,
        check_turning_points=True,
    )

    # Only passes if the attempt it passes on is less than max attempts
    assert success == (pass_on <= stack_params.max_attempts)
    # Final z is the one from the id returned by the stack "pick_me"
    assert final_z == 555

    # z_stack should run up until the time it passes. Running no more than max_attempts
    n_stacks = min(pass_on, stack_params.max_attempts)
    assert autofocus_thing.z_stack.call_count == n_stacks
    # Move absolute should be 1 less time that the number of times z_stack_run
    assert autofocus_thing._stage.move_absolute.call_count == n_stacks - 1
    # As should looping autofocus
    assert autofocus_thing.looping_autofocus.call_count == n_stacks - 1

    # Check rest stack is moving to the first image in the stack.
    if n_stacks > 1:
        assert autofocus_thing._stage.move_absolute.call_args.kwargs["z"] == -99

    # Mock called to save image
    assert autofocus_thing._cam.save_from_memory.call_count == (1 if success else 0)


def setup_and_run_z_stack(check_returns, check_turning_points, autofocus_thing, mocker):
    """Set up a z_stack, run it, and return the result.

    :param check_returns: The return values from check_stack_result. Note that if this
        is a list, it will be set as a side effect (and should be a list of tuples of
        results). If it a tuple (or anything else), it is set as a return value.
    """
    stack_params = autofocus_thing.create_stack_params(
        autofocus_dz=2000, images_dir="/this/is/fake", save_resolution=(1640, 1232)
    )
    stack_params.settling_time = 0  # Don't settle or tests take forever.

    autofocus_thing.capture_stack_image = mocker.Mock()
    if isinstance(check_returns, list):
        autofocus_thing.check_stack_result = mocker.Mock(side_effect=check_returns)
    else:
        autofocus_thing.check_stack_result = mocker.Mock(return_value=check_returns)
    return autofocus_thing.z_stack(
        stack_parameters=stack_params,
        check_turning_points=check_turning_points,
    )


def test_z_stack_turning_toggle_passed(autofocus_thing, mocker):
    """Check that the toggling of turning points is passed to the check."""
    check_returns = ("success", "mock_id")
    for check_turning in [True, False]:
        setup_and_run_z_stack(check_returns, check_turning, autofocus_thing, mocker)
        check_kwargs = autofocus_thing.check_stack_result.call_args.kwargs
        assert check_kwargs["check_turning_points"] == check_turning


def test_z_stack_returns_on_success_and_restart(autofocus_thing, mocker):
    """Check that if the check returns success or restart then the stack exits with correct return value."""
    for result in ["success", "restart"]:
        check_returns = (result, "mock_id")
        ret = setup_and_run_z_stack(check_returns, True, autofocus_thing, mocker)
        assert autofocus_thing.check_stack_result.call_count == 1
        # Check the number of images taken is exactly the call count.
        ims_taken = autofocus_thing.capture_stack_image.call_count
        assert ims_taken == autofocus_thing.stack_min_images_to_test
        # And the result is as expected.
        assert ret[0] == (result == "success")


def test_z_stack_exits_if_focus_never_found(autofocus_thing, mocker):
    """Check that if the check returns continue the stack exits eventually with a failure."""
    check_returns = ("continue", "mock_id")
    ret = setup_and_run_z_stack(check_returns, True, autofocus_thing, mocker)

    assert autofocus_thing.check_stack_result.call_count == EXTRA_STACK_CAPTURES + 1
    # Check the number of images taken is the maximum possible, set by the min images to
    # test and the number of extra images that can be taken
    ims_taken = autofocus_thing.capture_stack_image.call_count
    max_ims = autofocus_thing.stack_min_images_to_test + EXTRA_STACK_CAPTURES
    assert ims_taken == max_ims
    # And the result is as expected.
    assert not ret[0]


def test_z_stack_return(autofocus_thing, mocker):
    """Check z-stack returns as expected for more complex cases the fixed results above."""
    for i in range(2, EXTRA_STACK_CAPTURES):
        check_returns = [
            ("restart" if j == i - 1 else "continue", f"id_{j}") for j in range(i)
        ]
        ret = setup_and_run_z_stack(check_returns, True, autofocus_thing, mocker)
        # Calculate images taken
        images_taken = autofocus_thing.stack_min_images_to_test + i - 1
        assert autofocus_thing.capture_stack_image.call_count == images_taken
        # Check it reports a failure
        assert not ret[0]

        # Repeat ending with a success rather than a failure
        check_returns = [
            ("success" if j == i - 1 else "continue", f"id_{j}") for j in range(i)
        ]
        ret = setup_and_run_z_stack(check_returns, True, autofocus_thing, mocker)
        # Calculate images taken
        assert autofocus_thing.capture_stack_image.call_count == images_taken
        # Check it reports a success
        assert ret[0]


def test_capture_stack_image(autofocus_thing):
    """Check that capture stack image calls the expected functions and returns the expected data."""
    autofocus_thing._stage.position = {"x": 123, "y": 456, "z": 789}
    autofocus_thing._cam.capture_to_memory.return_value = "fake_buffer_id"
    autofocus_thing._cam.grab_jpeg_size.return_value = 54321
    buffer_max = 11

    info = autofocus_thing.capture_stack_image(buffer_max=buffer_max)
    assert autofocus_thing._cam.capture_to_memory.call_count == 1
    assert autofocus_thing._cam.grab_jpeg_size.call_count == 1
    assert info.buffer_id == "fake_buffer_id"
    assert info.position == {"x": 123, "y": 456, "z": 789}
    assert info.sharpness == 54321


def mock_capture(buffer_id: int, sharpness: int) -> CaptureInfo:
    """Create a CaptureInfo instance with a dummy position."""
    return CaptureInfo(
        buffer_id=buffer_id,
        position={"x": 0, "y": 0, "z": buffer_id},
        sharpness=sharpness,
    )


def test_check_stack_single_image_returns_success(autofocus_thing):
    """A single image is always successful."""
    captures = [mock_capture("mock-id", 10)]
    result, cap_id = autofocus_thing.check_stack_result(
        captures, check_turning_points=False
    )
    assert result == "success"
    assert cap_id == "mock-id"


@pytest.mark.parametrize(
    ("sharpnesses", "expected"),
    [
        ([5, 10, 3], "success"),
        ([10, 4, 2], "restart"),
        ([1, 2, 10], "continue"),
    ],
)
def test_check_stack_three_image_logic(sharpnesses, expected, autofocus_thing):
    """For 3 images, success is the highest one is central."""
    captures = [mock_capture(i, s) for i, s in enumerate(sharpnesses)]
    result, _ = autofocus_thing.check_stack_result(captures, check_turning_points=False)
    assert result == expected


def _run_check_stack_with_good_peak(autofocus_thing, count_turnings=False):
    """Run check stack on a good peak that should pass, and return the result.

    This can be used to check how other mocked results of subfunctions affects the
    result.
    """
    # Create an obvious peak that would normally pass.
    sharpnesses = [1, 2, 4, 7, 12, 7, 4, 2, 1]
    captures = [mock_capture(i, s) for i, s in enumerate(sharpnesses)]

    result, cap_id = autofocus_thing.check_stack_result(
        captures, check_turning_points=count_turnings
    )
    # Nothing a mocked function does should change which is the sharpest image.
    assert cap_id == 4
    return result


def test_check_stack_continues_if_no_tuning_point(autofocus_thing, mocker):
    """Check that continue is returned if no turning point is found."""
    # Mock to simulate not finding a peak
    mocker.patch(
        "openflexure_microscope_server.things.autofocus._get_peak_turning_point",
        side_effect=NotAPeakError,
    )
    result = _run_check_stack_with_good_peak(autofocus_thing)
    # Check that the NotAPeakError causes it to continue instead.
    assert result == "continue"


@pytest.mark.parametrize(
    ("location", "expected"),
    [
        (-10, "restart"),  # Restart if lower than 1.5 (halfway between im 2 and 3)
        (-1, "restart"),
        (0, "restart"),
        (1, "restart"),
        (1.49, "restart"),
        (1.5, "success"),  # Success up to 6.5 (as we have 9 images, final index is 8)
        (2.5, "success"),
        (4.5, "success"),
        (6.5, "success"),
        (6.51, "continue"),  # Continue if thrung point is after 6.5
        (7, "continue"),
        (8.1, "continue"),
        (123, "continue"),
    ],
)
def test_check_stack_affected_by_turning_point_location(
    location, expected, autofocus_thing, mocker
):
    """Check that the turning point location affects the return as expected."""
    # Mock to give the turning point location specified
    mocker.patch(
        "openflexure_microscope_server.things.autofocus._get_peak_turning_point",
        return_value=location,
    )
    result = _run_check_stack_with_good_peak(autofocus_thing)
    assert result == expected


def test_check_stack_affected_by_number_of_turning_points(autofocus_thing, mocker):
    """Check that the turning point location affects the return as expected."""
    # Set the turning point to the centre
    mocker.patch(
        "openflexure_microscope_server.things.autofocus._get_peak_turning_point",
        return_value=5,
    )
    mocker.patch(
        "openflexure_microscope_server.things.autofocus._count_turning_points",
        return_value=1,
    )
    result = _run_check_stack_with_good_peak(autofocus_thing, count_turnings=True)
    # Successful with 1 peak
    assert result == "success"

    # Change return to be 2 peaks
    mocker.patch(
        "openflexure_microscope_server.things.autofocus._count_turning_points",
        return_value=2,
    )
    result = _run_check_stack_with_good_peak(autofocus_thing, count_turnings=True)
    # Continue with 2 peaks
    assert result == "continue"
    # Unless this check is turned off
    result = _run_check_stack_with_good_peak(autofocus_thing, count_turnings=False)
    assert result == "success"


def test_get_peak_turning_point():
    """Check that the peak fitting returns expected value (or error)."""
    with pytest.raises(NotAPeakError):
        _get_peak_turning_point(np.ones(9))

    linear = np.arange(9)
    u_shape = 2 * (linear - 4) ** 2 + 17
    peak = -2 * (linear - 4) ** 2 + 55

    with pytest.raises(NotAPeakError):
        _get_peak_turning_point(linear)

    with pytest.raises(NotAPeakError):
        _get_peak_turning_point(u_shape)

    # Should be 4 to within a fitting error
    assert abs(_get_peak_turning_point(peak) - 4) < 1e-7


def test_count_turning_points():
    """Check the turing point count works as expected."""
    linear = np.arange(9)
    u_shape = 2 * (linear - 4) ** 2 + 17
    peak = -2 * (linear - 4) ** 2 + 55
    assert _count_turning_points(np.ones(9)) == 0
    assert _count_turning_points(linear) == 0
    assert _count_turning_points(u_shape) == 1
    assert _count_turning_points(peak) == 1

    assert _count_turning_points(np.array([1, 2, 3, 4, 5, 4, 3, 2, 1])) == 1
    # Double peak is 3 points
    assert _count_turning_points(np.array([1, 2, 3, 4, 2, 4, 3, 2, 1])) == 3
    # But only one if the dip isn't prominent
    assert _count_turning_points(np.array([1, 2, 3, 4, 3.8, 4, 3, 2, 1])) == 1