from __future__ import annotations

from typing import TYPE_CHECKING, cast

import numpy as np
from numpy.testing import assert_allclose, assert_array_equal
import pytest

from contourpy import FillType, LineType, contour_generator, dechunk_filled, dechunk_lines

from . import util_test

if TYPE_CHECKING:
    import contourpy._contourpy as cpy


@pytest.fixture
def z() -> cpy.CoordinateArray:
    return np.array([[0, 1, 0, 0, 0],
                     [0, 1, 0, 0, 0],
                     [0, 1, 0, 0, 0],
                     [0, 2, 0, 1, 0],
                     [0, 0, 0, 0, 0]], dtype=np.float64)


@pytest.mark.parametrize("fill_type", FillType.__members__.values())
@pytest.mark.parametrize("chunk_size", (0, 2))
def test_dechunk_filled(z: cpy.CoordinateArray, fill_type: FillType, chunk_size: int) -> None:
    cont_gen = contour_generator(z=z, fill_type=fill_type, chunk_size=chunk_size)
    assert cont_gen.chunk_count == ((2, 2) if chunk_size==2 else (1, 1))
    filled = cont_gen.filled(0.5, 1.5)

    dechunked = dechunk_filled(filled, fill_type)
    util_test.assert_filled(dechunked, fill_type)

    if chunk_size == 0 or fill_type in (FillType.OuterCode, FillType.OuterOffset):
        # Dechunking is a no-op for non-chunked fill types.
        assert dechunked == filled
    else:
        nchunks = len(dechunked[0])
        assert nchunks == 1

        # Note it is important that this contains a polygon with a hole,
        # and using chunk_size=2 there is an empty chunk.
        expected_points = np.array([
            [0.5, 0], [1, 0], [1.5, 0], [1.5, 1], [1.5, 2], [1, 2], [0.5, 2], [0.5, 1], [0.5, 0],
            [0.5, 2], [1, 2], [1.5, 2], [1.75, 3], [1, 3.75], [0.25, 3], [0.5, 2], [1, 2.5],
            [0.75, 3], [1, 3.25], [1.25, 3], [1, 2.5], [2.5, 3], [3, 2.5], [3.5, 3], [3, 3.5],
            [2.5, 3]])
        expected_codes = np.array([1, 2, 2, 2, 2, 2, 2, 2, 79, 1, 2, 2, 2, 2, 2, 79, 1, 2, 2, 2, 79,
                                   1, 2, 2, 2, 79])
        expected_offsets = np.array([0, 9, 16, 21, 26])

        assert dechunked[0][0] is not None
        assert_allclose(dechunked[0][0], expected_points)

        if fill_type == FillType.ChunkCombinedCode:
            if TYPE_CHECKING:
                dechunked = cast(cpy.FillReturn_ChunkCombinedCode, dechunked)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_codes)
        elif fill_type == FillType.ChunkCombinedOffset:
            if TYPE_CHECKING:
                dechunked = cast(cpy.FillReturn_ChunkCombinedOffset, dechunked)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_offsets)
        elif fill_type == FillType.ChunkCombinedCodeOffset:
            if TYPE_CHECKING:
                dechunked = cast(cpy.FillReturn_ChunkCombinedCodeOffset, dechunked)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_codes)
            assert dechunked[2][0] is not None
            assert_array_equal(dechunked[2][0], [0, 9, 21, 26])
        elif fill_type == FillType.ChunkCombinedOffsetOffset:
            if TYPE_CHECKING:
                dechunked = cast(cpy.FillReturn_ChunkCombinedOffsetOffset, dechunked)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_offsets)
            assert dechunked[2][0] is not None
            assert_array_equal(dechunked[2][0], [0, 1, 3, 4])
        else:
            raise RuntimeError(f"Unexpected fill_type {fill_type}")


@pytest.mark.parametrize("fill_type", FillType.__members__.values())
@pytest.mark.parametrize("chunk_size", (0, 1))
def test_dechunk_filled_empty(z: cpy.CoordinateArray, fill_type: FillType, chunk_size: int) -> None:
    cont_gen = contour_generator(z=z, fill_type=fill_type, chunk_size=chunk_size)
    filled = cont_gen.filled(3, 4)

    dechunked = dechunk_filled(filled, fill_type)
    util_test.assert_filled(dechunked, fill_type)

    if fill_type in (FillType.OuterCode, FillType.OuterOffset):
        assert dechunked == ([], [])
    elif fill_type in (FillType.ChunkCombinedCode, FillType.ChunkCombinedOffset):
        assert dechunked == ([None], [None])
    elif fill_type in (FillType.ChunkCombinedCodeOffset, FillType.ChunkCombinedOffsetOffset):
        assert dechunked == ([None], [None], [None])
    else:
        raise RuntimeError(f"Unexpected fill_type {fill_type}")


@pytest.mark.parametrize("line_type", LineType.__members__.values())
@pytest.mark.parametrize("chunk_size", (0, 2))
def test_dechunk_lines(z: cpy.CoordinateArray, line_type: LineType, chunk_size: int) -> None:
    cont_gen = contour_generator(z=z, line_type=line_type, chunk_size=chunk_size)
    assert cont_gen.chunk_count == ((2, 2) if chunk_size==2 else (1, 1))
    lines = cont_gen.lines(0.5)

    dechunked = dechunk_lines(lines, line_type)
    util_test.assert_lines(dechunked, line_type)

    if chunk_size == 0 or line_type in (LineType.Separate, LineType.SeparateCode):
        # Dechunking is a no-op for non-chunked line types.
        assert dechunked == lines
    else:
        nchunks = len(dechunked[0])
        assert nchunks == 1

        # Note it is important that this contains a closed line loop as well as open line strips,
        # and using chunk_size=2 there is an empty chunk.
        expected_points = np.array([
            [1.5, 0], [1.5, 1], [1.5, 2], [0.5, 2], [0.5, 1], [0.5, 0], [1.5, 2], [1.75, 3],
            [1, 3.75], [0.25, 3], [0.5, 2], [2.5, 3], [3, 2.5], [3.5, 3], [3, 3.5], [2.5, 3]])
        expected_codes = np.array([1, 2, 2, 1, 2, 2, 1, 2, 2, 2, 2, 1, 2, 2, 2, 79])
        expected_offsets = np.array([0, 3, 6, 11, 16])

        if line_type == LineType.ChunkCombinedCode:
            if TYPE_CHECKING:
                dechunked = cast(cpy.LineReturn_ChunkCombinedCode, dechunked)
            assert dechunked[0][0] is not None
            assert_allclose(dechunked[0][0], expected_points)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_codes)
        elif line_type == LineType.ChunkCombinedOffset:
            if TYPE_CHECKING:
                dechunked = cast(cpy.LineReturn_ChunkCombinedOffset, dechunked)
            assert dechunked[0][0] is not None
            assert_allclose(dechunked[0][0], expected_points)
            assert dechunked[1][0] is not None
            assert_array_equal(dechunked[1][0], expected_offsets)
        elif line_type == LineType.ChunkCombinedNan:
            if TYPE_CHECKING:
                dechunked = cast(cpy.LineReturn_ChunkCombinedNan, dechunked)
            assert dechunked[0][0] is not None
            # Convert offsets to int64 to avoid numpy error when mixing signed and unsigned ints.
            expected = np.insert(expected_points, expected_offsets[1:-1].astype(np.int64),
                                 [np.nan, np.nan], axis=0)
            assert_allclose(dechunked[0][0], expected)
        else:
            raise RuntimeError(f"Unexpected line_type {line_type}")


@pytest.mark.parametrize("line_type", LineType.__members__.values())
@pytest.mark.parametrize("chunk_size", (0, 1))
def test_dechunk_lines_empty(z: cpy.CoordinateArray, line_type: LineType, chunk_size: int) -> None:
    cont_gen = contour_generator(z=z, line_type=line_type, chunk_size=chunk_size)
    lines = cont_gen.lines(4)

    dechunked = dechunk_lines(lines, line_type)
    util_test.assert_lines(dechunked, line_type)

    if line_type == LineType.Separate:
        assert dechunked == []
    elif line_type == LineType.SeparateCode:
        assert dechunked == ([], [])
    elif line_type in (LineType.ChunkCombinedCode, LineType.ChunkCombinedOffset):
        assert dechunked == ([None], [None])
    elif line_type == LineType.ChunkCombinedNan:
        assert dechunked == ([None],)
    else:
        raise RuntimeError(f"Unexpected line_type {line_type}")
