from numpy import ndarray
from typing import Any
import numpy as np
import pyclipper

# Scaling factor for pyclipper (it works with integers)
SCALE = 1_000_000

def slice_at_range(models, z_values) -> list[list[np.ndarray]]:
    if not models:
        return []

    mmps: list[list[np.ndarray]] = []

    for model in models:
        model_edges: np.ndarray = edges_from_model(model)
        layered_edges: np.ndarray = layered_model(model_edges, z_values)
        sliced_edges: list[np.ndarray] = slice_at_z(layered_edges)
        
        all_model_paths: list[np.ndarray] = []

        for z, e in zip(z_values, sliced_edges):
            if e is None or len(e) == 0:
                continue

            zf = float(z)

            # Build 3D segments directly
            seg3d = e.copy()
            seg3d[..., 2] = zf

            # Build geometry for this slice
            slice_paths = segments_to_slice(seg3d)
            all_model_paths.extend(slice_paths)

        mmps.append(all_model_paths)
    return mmps

def edges_from_model(model): # returns a numpy: (facet, 3, 2, xyz)
    pairs_indices = np.array([[0, 1], [1, 2], [2, 0]]) # Create an array of indices to select the pairs
    return model[:, pairs_indices, :] # Use advanced indexing to create the array

def layered_model(models: np.ndarray, z_values: np.ndarray) -> np.ndarray:
    models_array: np.ndarray = np.array(models)[None, :]  # (1, F, 3, 2, 3)
    m = len(z_values)
    # array of zeros with shape (m, 1, 1, 1, 3)
    z_offsets_full = np.zeros((m, 1, 1, 1, 3))
    z_offsets_full[..., 2] = z_values[:, None, None, None]
    
    layered = models_array - z_offsets_full  # (m, F, 3, 2, 3)
    return layered

def slice_at_z(edges, normal_vector=np.array([0,0,1]), tolerance=1e-15):
    edges_offs = edges  # Shape (m, n, 3, 2, 3)

    # Obtain intersection points
    line_vector = edges_offs[..., 1, :] - edges_offs[..., 0, :]

    numerator = np.sum(normal_vector * -edges_offs[..., 0, :], axis=-1)
    denominator = np.sum(normal_vector[np.newaxis, np.newaxis, :] * line_vector, axis=-1)

    # Avoid division by zero
    with np.errstate(divide='ignore', invalid='ignore'):
        t = numerator / denominator
        t[denominator == 0] = np.nan  # Assign NaN where there's no valid intersection
    t[(t < 0) | (t > 1)] = np.nan

    intersections = edges_offs[..., 0, :] + t[..., np.newaxis] * line_vector
    intersections[np.abs(intersections) < tolerance] = 0  # Zero out small values

    all_segments = []
    for itx in intersections:
        mask = ~np.any(np.isnan(itx), axis=-1)
        filtered_arr = itx[mask]

        segment_count = filtered_arr.shape[-2] // 2
        segments = filtered_arr.reshape(segment_count, 2, 3)
        segments = np.around(segments, 6)
        all_segments.append(segments)

    return all_segments

def link_segments(segments: np.ndarray, tolerance=1e-6) -> list[np.ndarray]:
    """Links segments (N, 2, 3) into continuous paths."""
    if len(segments) == 0: return []
    z = segments[0, 0, 2]
    pts = segments[:, :, :2]
    
    def to_key(p): return tuple((p / tolerance).round().astype(np.int64))
    
    adj = {}
    for i in range(len(segments)):
        p1, p2 = pts[i, 0], pts[i, 1]
        k1, k2 = to_key(p1), to_key(p2)
        if k1 == k2: continue
        adj.setdefault(k1, []).append((k2, i))
        adj.setdefault(k2, []).append((k1, i))
        
    paths = []
    used = set()
    for i in range(len(segments)):
        if i in used: continue
        p1, p2 = pts[i, 0], pts[i, 1]
        k1, k2 = to_key(p1), to_key(p2)
        if k1 == k2: continue
        
        path_pts = [p1, p2]
        used.add(i)
        
        # Forward
        curr_k = k2
        while True:
            found = False
            for next_k, idx in adj.get(curr_k, []):
                if idx not in used:
                    pk1, pk2 = to_key(pts[idx, 0]), to_key(pts[idx, 1])
                    other_pt = pts[idx, 1] if pk1 == curr_k else pts[idx, 0]
                    path_pts.append(other_pt)
                    used.add(idx)
                    curr_k = to_key(other_pt)
                    found = True
                    break
            if not found: break
            
        # Backward
        curr_k = k1
        while True:
            found = False
            for next_k, idx in adj.get(curr_k, []):
                if idx not in used:
                    pk1, pk2 = to_key(pts[idx, 0]), to_key(pts[idx, 1])
                    other_pt = pts[idx, 1] if pk1 == curr_k else pts[idx, 0]
                    path_pts.insert(0, other_pt)
                    used.add(idx)
                    curr_k = to_key(other_pt)
                    found = True
                    break
            if not found: break
            
        res = np.column_stack([np.array(path_pts), np.full(len(path_pts), z)])
        paths.append(res)
    return paths

def segments_to_slice(segments: np.ndarray) -> list[np.ndarray]:

    linked = link_segments(segments)
    if not linked: return []
    
    pc = pyclipper.Pyclipper()
    for p in linked:
        # Pyclipper expects open paths (last point != first point) for ET_CLOSEDPOLYGON
        path_pts = p[:, :2]
        if len(path_pts) > 2 and np.allclose(path_pts[0], path_pts[-1]):
            path_pts = path_pts[:-1]
            
        scaled = (path_pts * SCALE).astype(np.int64).tolist()
        pc.AddPath(scaled, pyclipper.PT_SUBJECT, closed = True)
    
    # Use EVENODD for robust hole detection.
    solution = pc.Execute(pyclipper.CT_UNION, pyclipper.PFT_EVENODD, pyclipper.PFT_EVENODD)
    
    z = segments[0, 0, 2]
    res = []
    for path in solution:
        if not path: continue
        fpath = np.array(path) / SCALE
        fpath_3d = np.column_stack([fpath, np.full(len(fpath), z)])
        # Close loop by appending the first point for internal consistency
        fpath_3d = np.vstack([fpath_3d, fpath_3d[0]])
        res.append(fpath_3d)
    return res

def boolean_slice(subject: list[np.ndarray], object: list[np.ndarray], operation: str = 'UNION') -> list[np.ndarray]:
    all_z_keys: set[int] = set()
    subject_dict: dict[int, list[ndarray]] = {}
    object_dict: dict[int, list[ndarray]] = {}

    for s in subject:
        z = int(round(s[0, 2] * 1000))
        all_z_keys.add(z)
        subject_dict.setdefault(z, []).append(s)

    for s in object:
        z = int(round(s[0, 2] * 1000))
        all_z_keys.add(z)
        object_dict.setdefault(z, []).append(s)

    results: list[np.ndarray] = []

    for k in all_z_keys:
        subject_slices = subject_dict.get(k)
        object_slices = object_dict.get(k)

        if subject_slices is None:
            assert object_slices
            if operation == 'UNION':
                for s in object_slices: results.append(s)                
            continue

        if object_slices is None:
            assert subject_slices
            if operation == 'INTERSECTION': continue
            for s in subject_slices: results.append(s)
            continue

        z = subject_slices[0][0, 2]

        pc = pyclipper.Pyclipper()

        # Subject
        for s in subject_slices:
            path_pts = s[:, :2]
            scaled = (path_pts * SCALE).astype(np.int64).tolist()
            pc.AddPath(scaled, pyclipper.PT_SUBJECT, True)

        # Object
        for s in object_slices:
            path_pts = s[:, :2]
            scaled = (path_pts * SCALE).astype(np.int64).tolist()
            pc.AddPath(scaled, pyclipper.PT_CLIP, True)
    
        # Determine operation type. For now only UNION is supported.
        op_type = pyclipper.CT_UNION
        if operation == 'UNION': op_type = pyclipper.CT_UNION
        elif operation == 'INTERSECTION': op_type = pyclipper.CT_INTERSECTION
        elif operation == 'DIFFERENCE': op_type = pyclipper.CT_DIFFERENCE
        elif operation == 'XOR': op_type = pyclipper.CT_XOR

        solution = pc.Execute(op_type, pyclipper.PFT_EVENODD, pyclipper.PFT_EVENODD)
        
        for path in solution:
            if not path: continue
            fpath = np.array(path) / SCALE
            fpath_3d = np.column_stack([fpath, np.full(len(fpath), z)])
            fpath_3d = np.vstack([fpath_3d, fpath_3d[0]])
            results.append(fpath_3d)
            
    return results

def buffer_slice(slice_paths: list[np.ndarray], distance: float, miter='round', tolerance=.01) -> list[np.ndarray]:

    if not slice_paths: return []
    
    # Group paths by Z height to avoid flattening the entire model into 2D
    groups: dict[int, list[np.ndarray]] = {}
    for p in slice_paths:
        z = p[0, 2]
        # Use scaled integer for stable dictionary grouping of floats
        z_key = int(round(z * 1000))
        groups.setdefault(z_key, []).append(p)
    
    join_types = {
        'round': pyclipper.JT_ROUND,
        'square': pyclipper.JT_SQUARE,
        'flat': pyclipper.JT_SQUARE,
        'miter': pyclipper.JT_MITER
    }
    jt = join_types.get(miter, pyclipper.JT_ROUND)
    
    res = []
    pco = pyclipper.PyclipperOffset(arc_tolerance=tolerance * SCALE)
    
    # Buffer each layer independently
    for z_key in sorted(groups.keys()):
        layer_paths = groups[z_key]
        z = layer_paths[0][0, 2]
        
        pco.Clear()
        scaled_paths = []
        for p in layer_paths:
            path_pts = p[:, :2]
            # Strip redundant closing points
            if len(path_pts) > 2 and np.allclose(path_pts[0], path_pts[-1]):
                path_pts = path_pts[:-1]
            
            if len(path_pts) < 3: continue
            scaled_paths.append((path_pts * SCALE).astype(np.int64).tolist())
        
        if not scaled_paths: continue

        # Add all paths for this specific Z layer
        pco.AddPaths(scaled_paths, jt, pyclipper.ET_CLOSEDPOLYGON)
        solution = pco.Execute(distance * SCALE)
        
        for path in solution:
            if not path: continue
            fpath = np.array(path) / SCALE
            fpath_3d = np.column_stack([fpath, np.full(len(fpath), z)])
            # Re-close loop
            fpath_3d = np.vstack([fpath_3d, fpath_3d[0]])
            res.append(fpath_3d)
            
    return res

def group_segments_by_z(segments: np.ndarray) -> dict[int, np.ndarray]:
    z_values: np.ndarray = segments[:, 0, 2]
    unique_z: np.ndarray = np.unique(z_values)
    groups: dict[int, np.ndarray] = {int(round(z*1000)): segments[z_values == z] for z in unique_z}
    return groups

############################################################

class Pather(object):

    def __init__(self, paths: list[np.ndarray], **kwargs) -> None:
        self.paths = paths
        self.bed_center_x: float = 250.
        self.bed_center_y: float = 250.
        self.feedrate: float = 1200.
        self.extrusion_rate: float = 0.017

    def to_gcode(self) -> str:
        if not self.paths:
            return ""

        def segment_length_3d(p1, p2) -> float:
            dx = p2[0] - p1[0]
            dy = p2[1] - p1[1]
            dz = p2[2] - p1[2]
            return (dx * dx + dy * dy + dz * dz) ** 0.5

        gcode: list[str] = [
            "; Generated from NodeSlicer",
            f";TYPE:Custom",
            f";FEEDRATE={self.feedrate}",
            f";EXTRUSION_RATE={self.extrusion_rate}",
            "G21 ; set units to mm",
            "G90 ; absolute XYZ positioning",
            "M83 ; relative extrusion mode",
            "G92 E0",
        ]

        last_z = -1.0
        max_z = 0

        for line in self.paths:
            coords = line.tolist()
            if len(coords) < 2:
                continue

            x0, y0, z0 = coords[0]

            if abs(z0 - last_z) > 0.001:
                gcode.append(f";BEFORE_LAYER_CHANGE")
                gcode.append(f"G0 Z{z0:.3f} F720")
                gcode.append(f";AFTER_LAYER_CHANGE")
                last_z = z0

            gcode.append(f"G0 X{x0:.3f} Y{y0:.3f} Z{z0:.3f}")
            gcode.append(";TYPE:External perimeter")

            for p1, p2 in zip(coords, coords[1:]):
                seg_len: float = segment_length_3d(p1, p2)
                e_segment = seg_len * self.extrusion_rate
                x, y, z = p2
                gcode.append(
                    f"G1 X{x:.3f} Y{y:.3f} Z{z:.3f} E{e_segment:.5f} F{self.feedrate}"
                )
                max_z = max(max_z, z)
            
        return "\n".join(gcode)