# ilda/utils.py
from dataclasses import dataclass
from typing import List, Tuple, Optional
import math

Point = Tuple[float, float]
PathPoints = List[Point]


@dataclass
class PathOrderResult:
    """
    Container for ordered paths result.

    Attributes:
      paths: List of ordered polylines (each polyline is a list of (x,y) tuples).
      reversed_flags: Parallel list of bools, True if corresponding original path was reversed.
    """
    paths: List[PathPoints]
    reversed_flags: List[bool]


def _dist(a: Point, b: Point) -> float:
    """Euclidean distance between two points."""
    return math.hypot(a[0] - b[0], a[1] - b[1])


def order_paths_greedy_with_2opt(paths: List[PathPoints],
                                 start_point: Optional[Point] = None,
                                 do_2opt: bool = True) -> PathOrderResult:
    """
    Order polylines to reduce dead (blanked) travel between them.

    Parameters
    ----------
    paths :
        List of polylines (each polyline must have >= 2 points).
    start_point :
        Optional coordinate (x,y) to start the sequence from. If None, starts from first path.
    do_2opt :
        Whether to run a local 2-opt improvement after greedy construction.

    Returns
    -------
    PathOrderResult
        Ordered paths and flags indicating if each path was reversed relative to original.
    """
    n = len(paths)
    if n == 0:
        return PathOrderResult([], [])

    # Precompute endpoints (start, end) for each path
    endpoints: List[Tuple[Point, Point]] = [(p[0], p[-1]) for p in paths]

    used = [False] * n
    order: List[int] = []
    rev_flags: List[bool] = []

    # choose initial index
    if start_point is None:
        cur_idx = 0
        cur_rev = False
        used[cur_idx] = True
        order.append(cur_idx)
        rev_flags.append(cur_rev)
        cur_pt = endpoints[cur_idx][1] if not cur_rev else endpoints[cur_idx][0]
    else:
        # pick closest endpoint among all paths
        best_idx = -1
        best_rev = False
        best_d = float('inf')
        for i, (s, e) in enumerate(endpoints):
            d_s = _dist(start_point, s)
            if d_s < best_d:
                best_d = d_s; best_idx = i; best_rev = False
            d_e = _dist(start_point, e)
            if d_e < best_d:
                best_d = d_e; best_idx = i; best_rev = True
        cur_idx = best_idx
        cur_rev = best_rev
        used[cur_idx] = True
        order.append(cur_idx)
        rev_flags.append(cur_rev)
        cur_pt = endpoints[cur_idx][1] if not cur_rev else endpoints[cur_idx][0]

    # greedy nearest neighbor choosing orientation
    for _ in range(1, n):
        best_j = -1
        best_j_rev = False
        best_d = float('inf')
        for j in range(n):
            if used[j]:
                continue
            s, e = endpoints[j]
            d1 = _dist(cur_pt, s)   # not reversed
            if d1 < best_d:
                best_d = d1; best_j = j; best_j_rev = False
            d2 = _dist(cur_pt, e)   # reversed
            if d2 < best_d:
                best_d = d2; best_j = j; best_j_rev = True
        if best_j == -1:
            break
        used[best_j] = True
        order.append(best_j)
        rev_flags.append(best_j_rev)
        s, e = endpoints[best_j]
        cur_pt = e if not best_j_rev else s

    # Build ordered path list (apply reversal where needed)
    ordered_paths: List[PathPoints] = []
    for idx, rev in zip(order, rev_flags):
        p = paths[idx]
        ordered_paths.append(list(reversed(p)) if rev else list(p))

    # Optional 2-opt improvement (operates on ordered_paths)
    if do_2opt and len(ordered_paths) > 2:
        def connection_cost(a: PathPoints, b: PathPoints) -> float:
            return _dist(a[-1], b[0])

        improved = True
        seq = ordered_paths
        while improved:
            improved = False
            m = len(seq)
            for i in range(0, m - 2):
                for j in range(i + 2, m):
                    # cost affected at edges (i -> i+1) and (j -> j+1)
                    old_cost = connection_cost(seq[i], seq[i+1])
                    if j + 1 < m:
                        old_cost += connection_cost(seq[j], seq[j+1])
                    # form new sequence by reversing subsequence (i+1 .. j) and reversing each subpath
                    new_seq = seq[:i+1] + [list(reversed(seg)) for seg in seq[i+1:j+1]][::-1] + seq[j+1:]
                    new_cost = connection_cost(new_seq[i], new_seq[i+1])
                    if j + 1 < m:
                        new_cost += connection_cost(new_seq[j], new_seq[j+1])
                    if new_cost + 1e-12 < old_cost:
                        seq = new_seq
                        improved = True
                        break
                if improved:
                    break
        ordered_paths = seq

    # Reconstruct reversed_flags relative to originals
    final_flags: List[bool] = []
    # Create mapping of original endpoints for approximate identification
    for p in ordered_paths:
        if not p:
            final_flags.append(False)
            continue
        first = p[0]
        found = False
        for orig in paths:
            if abs(orig[0][0] - first[0]) < 1e-9 and abs(orig[0][1] - first[1]) < 1e-9:
                final_flags.append(False)
                found = True
                break
            if abs(orig[-1][0] - first[0]) < 1e-9 and abs(orig[-1][1] - first[1]) < 1e-9:
                final_flags.append(True)
                found = True
                break
        if not found:
            final_flags.append(False)

    return PathOrderResult(paths=ordered_paths, reversed_flags=final_flags)