import numpy as np
import pytest

from starfish import ImageStack
from starfish.core.imagestack.test.factories import unique_tiles_imagestack
from starfish.core.test.factories import (
    two_spot_informative_blank_coded_data_factory,
    two_spot_one_hot_coded_data_factory,
    two_spot_sparse_coded_data_factory,
)
from starfish.types import Axes, FunctionSource
from .._base import FindSpotsAlgorithm
from ..blob import BlobDetector
from ..local_max_peak_finder import LocalMaxPeakFinder
from ..trackpy_local_max_peak_finder import TrackpyLocalMaxPeakFinder

# verify all spot finders handle different coding types
_, ONE_HOT_IMAGESTACK, ONE_HOT_MAX_INTENSITY = two_spot_one_hot_coded_data_factory()
_, SPARSE_IMAGESTACK, SPARSE_MAX_INTENSITY = two_spot_sparse_coded_data_factory()
_, BLANK_IMAGESTACK, BLANK_MAX_INTENSITY = two_spot_informative_blank_coded_data_factory()

# make sure that all spot finders handle empty arrays
EMPTY_IMAGESTACK = ImageStack.from_numpy(np.zeros((4, 2, 10, 100, 100), dtype=np.float32))


def simple_gaussian_spot_detector() -> BlobDetector:
    """create a basic gaussian spot detector"""
    return BlobDetector(
        min_sigma=1,
        max_sigma=4,
        num_sigma=5,
        threshold=0,
        measurement_type='max')


def simple_trackpy_local_max_spot_detector() -> TrackpyLocalMaxPeakFinder:
    """create a basic local max peak finder"""
    return TrackpyLocalMaxPeakFinder(
        spot_diameter=3,
        min_mass=0.01,
        max_size=10,
        separation=2,
        radius_is_gyration=True
    )


def simple_local_max_spot_detector_3d() -> LocalMaxPeakFinder:
    return LocalMaxPeakFinder(
        min_distance=6,
        stringency=0,
        min_obj_area=0,
        max_obj_area=np.iinfo(int).max,
        threshold=0,
        is_volume=True,
    )


def simple_local_max_spot_detector_2d() -> LocalMaxPeakFinder:
    return LocalMaxPeakFinder(
        min_distance=6,
        stringency=0,
        min_obj_area=0,
        max_obj_area=np.iinfo(int).max,
        threshold=0,
        is_volume=False
    )


# initialize spot detectors
gaussian_spot_detector = simple_gaussian_spot_detector()
trackpy_local_max_spot_detector = simple_trackpy_local_max_spot_detector()
local_max_spot_detector_2d = simple_local_max_spot_detector_2d()
local_max_spot_detector_3d = simple_local_max_spot_detector_3d()

# test parameterization
test_parameters = (
    'data_stack, spot_detector, is_2d',
    [
        (ONE_HOT_IMAGESTACK, gaussian_spot_detector, False),
        (ONE_HOT_IMAGESTACK, trackpy_local_max_spot_detector, False),
        (ONE_HOT_IMAGESTACK, local_max_spot_detector_2d, True),
        (ONE_HOT_IMAGESTACK, local_max_spot_detector_3d, False),
        (SPARSE_IMAGESTACK, gaussian_spot_detector, False),
        (SPARSE_IMAGESTACK, trackpy_local_max_spot_detector, False),
        (SPARSE_IMAGESTACK, local_max_spot_detector_2d, True),
        (SPARSE_IMAGESTACK, local_max_spot_detector_3d, False),
        (BLANK_IMAGESTACK, gaussian_spot_detector, False),
        (BLANK_IMAGESTACK, trackpy_local_max_spot_detector, False),
        (BLANK_IMAGESTACK, local_max_spot_detector_2d, True),
        (BLANK_IMAGESTACK, local_max_spot_detector_3d, False),
    ]
)


@pytest.mark.parametrize(*test_parameters)
def test_spot_detection_with_reference_image(
        data_stack: ImageStack,
        spot_detector: FindSpotsAlgorithm,
        is_2d: bool,
):
    """This testing method uses a reference image to identify spot locations. Each method should
    find 2 spots in the reference image then measure the intensity of those locations in each
    r/ch pair. The final spot results should represent 2 spots for each r/ch totalling
    2*num_rounds*num_ch spots """

    # 2D spot detection will find spots on 8 zlayers, so it produces 7x as many spots.
    multiplier = 7 if is_2d else 1
    reference_image = data_stack.reduce((Axes.ROUND, Axes.CH), func=FunctionSource.np("max"))
    spots_results = spot_detector.run(image_stack=data_stack, reference_image=reference_image)
    assert spots_results.count_total_spots() == (
        2 * data_stack.num_chs * data_stack.num_rounds * multiplier), \
        "wrong number of spots detected"


@pytest.mark.parametrize(*test_parameters)
def test_spot_detection_no_reference_image(
        data_stack: ImageStack,
        spot_detector: FindSpotsAlgorithm,
        is_2d: bool,
):
    """
    This testing method does not provide a reference image, and should therefore check for spots
    in each (round, ch) combination in sequence. The final spot results should represent the
    total number of original spots in the ImageStack.
    """

    # 2D spot detection will find spots on 8 zlayers, so it produces 7x as many spots.
    multiplier = 7 if is_2d else 1
    spots_results = spot_detector.run(image_stack=data_stack)
    assert spots_results.count_total_spots() == 4 * multiplier, "wrong number of spots detected"

    spots = spot_detector.run(image_stack=EMPTY_IMAGESTACK)
    assert spots.count_total_spots() == 0


def _make_labeled_image() -> ImageStack:
    ROUND_LABELS = (1, 4, 6)
    CH_LABELS = (2, 4, 6, 8)
    ZPLANE_LABELS = (3, 4)
    HEIGHT = 2
    WIDTH = 4

    return unique_tiles_imagestack(
        ROUND_LABELS, CH_LABELS, ZPLANE_LABELS, HEIGHT, WIDTH)


def test_reference_image_spot_detection_with_image_with_labeled_axes():
    """This testing method uses a reference image to identify spot locations."""
    data_stack = _make_labeled_image()
    reference_image = data_stack.reduce((Axes.ROUND, Axes.CH), func=FunctionSource.np("max"))
    spot_results = gaussian_spot_detector.run(image_stack=data_stack,
                                              reference_image=reference_image)
    return spot_results


def test_spot_detection_with_image_with_labeled_axes():
    """This testing method uses no reference image to identify spot locations."""
    data_stack = _make_labeled_image()
    spot_results = gaussian_spot_detector.run(image_stack=data_stack)
    return spot_results


def test_blob_detector_2d_spot_coordinates():
    """Test that BlobDetector with is_volume=False produces correct y, x coordinates and radius.

    This is a regression test for the bug where y-values were all 0 and radius was incorrect
    due to incorrect indexing of 2D images.
    """
    # Create a simple 2D image with a bright spot
    image_2d = np.zeros((100, 100), dtype=np.float32)
    # Add a 5x5 bright spot centered approximately at y=50, x=60
    image_2d[48:53, 58:63] = 1.0

    # Create an ImageStack with this 2D image
    # Shape: (rounds=1, channels=1, z-planes=1, height=100, width=100)
    image_stack = ImageStack.from_numpy(image_2d.reshape(1, 1, 1, 100, 100))

    # Create a BlobDetector with is_volume=False
    detector_2d = BlobDetector(
        min_sigma=1,
        max_sigma=3,
        num_sigma=5,
        threshold=0.01,
        is_volume=False,
        measurement_type='max'
    )

    # Run detection
    spot_results = detector_2d.run(image_stack=image_stack)

    # Get the spot attributes for round 0, channel 0
    spots = spot_results[{Axes.ROUND: 0, Axes.CH: 0}].spot_attrs

    # Verify we found at least one spot
    assert len(spots.data) > 0, "No spots detected"

    # Check that y-values are not all 0 (the bug symptom)
    y_values = spots.data['y'].values
    assert not np.all(y_values == 0), "All y-values are 0, indicating the bug is present"

    # Check that the detected spot is near the expected location (y=50, x=60)
    # Allow 5 pixel tolerance since blob detection may not find exact center
    assert np.any(np.abs(y_values - 50) < 5), f"No spot found near y=50, found: {y_values}"

    x_values = spots.data['x'].values
    assert np.any(np.abs(x_values - 60) < 5), f"No spot found near x=60, found: {x_values}"

    # Check that radius is reasonable (not extremely large like 843.0 in the bug)
    radius_values = spots.data['radius'].values
    assert np.all(radius_values < 100), f"Radius values too large: {radius_values}"
    assert np.all(radius_values > 0), f"Radius values should be positive: {radius_values}"


def test_blob_detector_2d_with_reference_image():
    """Test BlobDetector with is_volume=False and a reference_image.

    This tests the case where reference_image has multiple z-planes, which results in
    a 3D data_image after squeezing ROUND and CH dimensions.
    """
    # Create a 3D reference image with multiple z-planes
    reference_3d = np.zeros((3, 100, 100), dtype=np.float32)
    # Add a bright spot at z=0, y=30, x=40
    reference_3d[0, 28:33, 38:43] = 1.0

    # Create ImageStacks with the reference image (3 z-planes, 1 round, 1 channel)
    reference_stack = ImageStack.from_numpy(reference_3d.reshape(1, 1, 3, 100, 100))
    # Create a primary image stack (same dimensions)
    primary_stack = ImageStack.from_numpy(reference_3d.reshape(1, 1, 3, 100, 100))

    # Create a BlobDetector with is_volume=False
    detector_2d = BlobDetector(
        min_sigma=1,
        max_sigma=3,
        num_sigma=5,
        threshold=0.01,
        is_volume=False,
        measurement_type='max'
    )

    # Run detection with reference_image
    spot_results = detector_2d.run(image_stack=primary_stack, reference_image=reference_stack)

    # The reference image approach produces spots for all (round, ch) combinations
    # For our case: 1 round x 1 channel = 1 combination
    total_spots = spot_results.count_total_spots()
    assert total_spots > 0, "No spots detected with reference image"

    # Check a specific round/channel
    spots = spot_results[{Axes.ROUND: 0, Axes.CH: 0}].spot_attrs

    # Verify spot coordinates and radius are correct
    y_values = spots.data['y'].values
    x_values = spots.data['x'].values
    radius_values = spots.data['radius'].values

    # Check that y-values are not all 0
    assert not np.all(y_values == 0), "All y-values are 0, bug is present"

    # Check that the spot is near the expected location (y=30, x=40)
    assert np.any(np.abs(y_values - 30) < 5), f"No spot found near y=30, found: {y_values}"
    assert np.any(np.abs(x_values - 40) < 5), f"No spot found near x=40, found: {x_values}"

    # Check that radius is reasonable
    assert np.all(radius_values < 100), f"Radius values too large: {radius_values}"
    assert np.all(radius_values > 0), f"Radius values should be positive: {radius_values}"


def test_blob_detector_round_channel_assignment():
    """Test that BlobDetector with is_volume=False assigns spots to correct (round, channel) pairs.

    This is a regression test for a bug where spots were assigned to incorrect round/channel
    combinations due to a mismatch between the sorted (r, ch) tuples and the iteration order
    from _iter_axes({Axes.ROUND, Axes.CH}).
    """
    # Create a test ImageStack with 2 rounds, 3 channels, 1 z-plane
    # Each (round, channel) pair gets a spot at a unique location
    data = np.zeros((2, 3, 1, 100, 100), dtype=np.float32)

    # Add spots at different y-coordinates for each (round, channel)
    # Round 0, Channel 0: y=10
    data[0, 0, 0, 8:13, 8:13] = 1.0
    # Round 0, Channel 1: y=20
    data[0, 1, 0, 18:23, 18:23] = 1.0
    # Round 0, Channel 2: y=30
    data[0, 2, 0, 28:33, 28:33] = 1.0
    # Round 1, Channel 0: y=40
    data[1, 0, 0, 38:43, 38:43] = 1.0
    # Round 1, Channel 1: y=50
    data[1, 1, 0, 48:53, 48:53] = 1.0
    # Round 1, Channel 2: y=60
    data[1, 2, 0, 58:63, 58:63] = 1.0

    image_stack = ImageStack.from_numpy(data)

    # Test with is_volume=False
    detector_2d = BlobDetector(
        min_sigma=1,
        max_sigma=3,
        num_sigma=5,
        threshold=0.01,
        is_volume=False,
        measurement_type='mean'
    )

    spots_2d = detector_2d.run(image_stack=image_stack)

    # Test with is_volume=True for comparison
    detector_3d = BlobDetector(
        min_sigma=1,
        max_sigma=3,
        num_sigma=5,
        threshold=0.01,
        is_volume=True,
        measurement_type='mean'
    )

    spots_3d = detector_3d.run(image_stack=image_stack)

    # Define expected y-coordinates for each (round, channel) pair
    expected_y = {
        (0, 0): 10,
        (0, 1): 20,
        (0, 2): 30,
        (1, 0): 40,
        (1, 1): 50,
        (1, 2): 60,
    }

    # Check that both is_volume=False and is_volume=True produce the same results
    for r in range(2):
        for ch in range(3):
            # Get spots for this (round, channel)
            spots_2d_data = spots_2d[{Axes.ROUND: r, Axes.CH: ch}].spot_attrs.data
            spots_3d_data = spots_3d[{Axes.ROUND: r, Axes.CH: ch}].spot_attrs.data

            # Both should have found the spot
            assert len(spots_2d_data) > 0, \
                f"No spots found for Round={r}, CH={ch} with is_volume=False"
            assert len(spots_3d_data) > 0, \
                f"No spots found for Round={r}, CH={ch} with is_volume=True"

            # Check that the y-coordinate matches the expected value
            y_2d = spots_2d_data['y'].values[0]
            y_3d = spots_3d_data['y'].values[0]
            expected = expected_y[(r, ch)]

            assert np.abs(y_2d - expected) < 5, \
                f"is_volume=False: Round={r}, CH={ch} has y={y_2d:.0f}, expected ~{expected}"
            assert np.abs(y_3d - expected) < 5, \
                f"is_volume=True: Round={r}, CH={ch} has y={y_3d:.0f}, expected ~{expected}"

            # Both should produce the same y-coordinate
            assert np.abs(y_2d - y_3d) < 2, \
                f"Round={r}, CH={ch}: is_volume=False has y={y_2d:.0f} " \
                f"but is_volume=True has y={y_3d:.0f}"
