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 .helper_functions import interpolate_line
from ..beams import BeamType
from ..mill import Mill


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

    beam_radius = beam.get_radius(units="nm")
    overlap = mill.get_overlap()
    step = 2 * beam_radius * (1 - overlap)
    dwelltime = int(mill.get_dwell_time("ns")) / 100

    # Ensure we have a MultiPolygon to loop through
    if isinstance(shape, Polygon):
        shape = MultiPolygon([shape])

    streamlines_per_polygon = []  # Store scan lines per polygon

    for polygon in shape.geoms:  # Process each polygon separately
        buffered: Polygon = buffer(polygon, -beam_radius, join_style="mitre", quad_segs=32)
        if buffered.is_empty:
            continue

        prepare(buffered)
        streamlines = []
        minx, miny, maxx, maxy = buffered.bounds

        buffered: Polygon = buffer(buffered, + beam_radius * 0.1, join_style="mitre", quad_segs=32)
        prepare(buffered)

        # Fetch some scan direction dependent values
        if horizontal:
            nr_steps = np.rint((maxy - miny) / step)
            span = (maxy - miny)
            scan_positions = np.linspace(miny, maxy, int(nr_steps))
        else:
            nr_steps = np.rint((maxx - minx) / step)
            span = (maxx - minx)
            scan_positions = np.linspace(minx, maxx, int(nr_steps))

        # Find the stepsize in the direction normal to the scan direction.
        normal_stepsize = span / nr_steps

        for current in scan_positions:
            if horizontal:
                scan_line = LineString([(minx - 1000, current), (maxx + 1000, current)])
            else:
                scan_line = LineString([(current, miny - 1000), (current, maxy + 1000)])

            inters = intersection(buffered, scan_line)

            if inters.is_empty:
                current += normal_stepsize
                continue

            segments = []
            if isinstance(inters, LineString):
                segments = [inters]
            elif isinstance(inters, MultiLineString):
                segments = list(inters.geoms)
            elif isinstance(inters, Point):
                segments = [inters]

            for seg in segments:

                if isinstance(seg, Point):
                    points = np.array(seg.coords)
                    codes = np.full(len(points), 2, dtype=int)
                    codes[0] = 1  # Start code
                    dwelltimes = np.full(len(points), dwelltime)
                    streamline = np.column_stack((dwelltimes, points[:, 0], points[:, 1], codes))
                    streamlines.append(streamline)
                    continue

                coords = np.array(seg.coords)
                if coords.size == 0:
                    continue
                sorted_indices = np.argsort(coords[:, 0])
                sorted_coords = coords[sorted_indices]
                sorted_seg = LineString(sorted_coords)

                line_length = sorted_seg.length
                if line_length == 0:
                    continue
                nr_steps = int(np.rint(line_length / step))
                distances = np.linspace(0, line_length, nr_steps)
                points = np.array([sorted_coords[0]]) if nr_steps == 0 else interpolate_line(sorted_seg, distances)

                if points.size == 0:
                    continue

                codes = np.full(len(points), 2, dtype=int)
                codes[0] = 1  # Start code
                dwelltimes = np.full(len(points), dwelltime)
                streamline = np.column_stack((dwelltimes, points[:, 0], points[:, 1], codes))
                streamlines.append(streamline)

            current += normal_stepsize

        streamlines_per_polygon.extend(streamlines)

    return streamlines_per_polygon


def horizontal_left_to_right(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]))
    return streamlines


def horizontal_right_to_left(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]))

    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):
        group = [line[::-1] for line in group][::-1]
        final_streamlines.extend(group)

    return final_streamlines


def vertical_bottom_to_top(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]))
    return streamlines


def vertical_top_to_bottom(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):
        group = [line[::-1] for line in group][::-1]
        final_streamlines.extend(group)

    return final_streamlines
