from typing import List, Union
import itertools

import numpy as np
from numpy.typing import NDArray
from shapely import Point, LineString, MultiLineString, Polygon, MultiPolygon, buffer, prepare, intersection

from .raster import _raster_sequential
from ..beams import BeamType
from ..mill import Mill


def serpentine_horizontal(shape: Union[Polygon, MultiPolygon], mill: Mill, beam: BeamType,
                          horizontal: bool = True) -> List[NDArray]:
    streamlines = _raster_sequential(shape, mill, beam, horizontal=True)

    tolerance = beam.get_radius("nm")
    streamlines.sort(key=lambda s: (int(s[0, 2] / tolerance), s[0, 1]))

    # Group the streamlines by the bucketed y coordinate.
    grouped_streamlines = []
    for group_key, group in itertools.groupby(streamlines, key=lambda s: int(s[0, 2] / tolerance)):
        grouped_streamlines.append(list(group))

    # Now, for every other pass (group), reverse the order of points in each scan line.
    final_streamlines = []
    for i, group in enumerate(grouped_streamlines):
        if i % 2 == 1:  # odd-indexed pass: reverse the scan direction
            group = [line[::-1] for line in group][::-1]
        final_streamlines.extend(group)

    return final_streamlines


def serpentine_vertical(shape: Union[Polygon, MultiPolygon], mill: Mill, beam: BeamType,
                        horizontal: bool = True) -> List[NDArray]:

    streamlines = _raster_sequential(shape, mill, beam, horizontal=False)

    tolerance = beam.get_radius("nm")
    streamlines.sort(key=lambda s: (int(s[0, 1] / tolerance), s[0, 2]))

    # Group the streamlines by the bucketed y coordinate.
    grouped_streamlines = []
    for group_key, group in itertools.groupby(streamlines, key=lambda s: int(s[0, 1] / tolerance)):
        grouped_streamlines.append(list(group))

    # Now, for every other pass (group), reverse the order of points in each scan line.
    final_streamlines = []
    for i, group in enumerate(grouped_streamlines):
        if i % 2 == 1:  # odd-indexed pass: reverse the scan direction
            group = [line[::-1] for line in group][::-1]
        final_streamlines.extend(group)

    return final_streamlines
