# AverageAxis - Fusion 360 Script
# Finds the average axis direction of N selected edges (any type)
# and creates a construction axis + perpendicular construction plane.
#
# - Supports linear edges (uses start->end direction)
# - Supports curved edges (uses tangent at the midpoint parameter)
# - Select as many edges as you like; press Esc when done (minimum 2)
# - All construction geometry is created in the *active* component
#
# Author: Claude / Anthropic
# Tested against: Fusion 360 API - January 2026

import adsk.core
import adsk.fusion
import traceback
import math


def get_edge_direction_and_midpoint(edge):
    """Return (direction_tuple, midpoint_tuple) for any BRepEdge.

    For linear edges the direction is simply start -> end.
    For curved edges the tangent vector at the parametric midpoint is used.
    All direction vectors are normalised to unit length.
    Returns None if the edge is degenerate (zero-length tangent).
    """
    geom = edge.geometry  # Curve3D subclass

    if geom.curveType == adsk.core.Curve3DTypes.Line3DCurveType:
        # ---- Linear edge: straightforward start -> end ----
        line = adsk.core.Line3D.cast(geom)
        sp = line.startPoint
        ep = line.endPoint
        dx = ep.x - sp.x
        dy = ep.y - sp.y
        dz = ep.z - sp.z
        length = math.sqrt(dx * dx + dy * dy + dz * dz)
        if length < 1e-10:
            return None
        direction = (dx / length, dy / length, dz / length)
        midpoint = (
            (sp.x + ep.x) / 2.0,
            (sp.y + ep.y) / 2.0,
            (sp.z + ep.z) / 2.0,
        )
        return direction, midpoint

    else:
        # ---- Curved edge: tangent at the parametric midpoint ----
        evaluator = edge.evaluator

        # Get the parameter range of the edge
        ok, pStart, pEnd = evaluator.getParameterExtents()
        if not ok:
            return None
        pMid = (pStart + pEnd) / 2.0

        # Tangent vector at the midpoint parameter
        ok, tangents = evaluator.getTangents([pMid])
        if not ok or len(tangents) == 0:
            return None
        tangent = tangents[0]  # Vector3D

        length = tangent.length
        if length < 1e-10:
            return None
        tangent.normalize()
        direction = (tangent.x, tangent.y, tangent.z)

        # 3-D point at the midpoint parameter (used for axis origin)
        ok, points = evaluator.getPointsAtParameters([pMid])
        if not ok or len(points) == 0:
            return None
        pt = points[0]
        midpoint = (pt.x, pt.y, pt.z)

        return direction, midpoint


def curve_type_label(geom):
    """Return a human-readable label for a Curve3D type."""
    lookup = {
        adsk.core.Curve3DTypes.Line3DCurveType:           'linear',
        adsk.core.Curve3DTypes.Arc3DCurveType:             'arc',
        adsk.core.Curve3DTypes.Circle3DCurveType:          'circle',
        adsk.core.Curve3DTypes.EllipticalArc3DCurveType:   'elliptical arc',
        adsk.core.Curve3DTypes.Ellipse3DCurveType:         'ellipse',
        adsk.core.Curve3DTypes.NurbsCurve3DCurveType:      'spline',
    }
    return lookup.get(geom.curveType, 'curve')


def run(context):
    ui = None
    try:
        app = adsk.core.Application.get()
        ui = app.userInterface
        design = adsk.fusion.Design.cast(app.activeProduct)

        if not design:
            ui.messageBox('No active Fusion design. Please open a design first.')
            return

        # Use the currently active (edited) component, not the root
        comp = design.activeComponent

        # ------------------------------------------------------------------
        # 1. Collect edges from the user (any number, Esc to finish)
        # ------------------------------------------------------------------
        edges = []
        while True:
            count = len(edges)
            if count < 2:
                hint = 'Esc to cancel'
            else:
                hint = f'{count} selected \u2014 Esc to finish'
            prompt = f'Select edge {count + 1}  ({hint})'

            try:
                sel = ui.selectEntity(prompt, 'Edges')
            except:
                # selectEntity raises an exception when the user presses Esc
                break

            if sel is None:
                break

            edge = adsk.fusion.BRepEdge.cast(sel.entity)
            if edge is None:
                ui.messageBox('Selected entity is not a BRepEdge \u2014 try again.')
                continue
            edges.append(edge)

        if len(edges) < 2:
            ui.messageBox(
                f'Need at least 2 edges to compute an average '
                f'({len(edges)} selected). Script cancelled.'
            )
            return

        # ------------------------------------------------------------------
        # 2. Extract direction vectors from each edge
        # ------------------------------------------------------------------
        directions = []
        midpoints = []
        edge_labels = []

        for i, edge in enumerate(edges):
            result = get_edge_direction_and_midpoint(edge)
            if result is None:
                ui.messageBox(
                    f'Edge {i + 1} has a degenerate tangent and was skipped.'
                )
                continue
            d, m = result
            directions.append(d)
            midpoints.append(m)
            edge_labels.append(f'Edge {i + 1} ({curve_type_label(edge.geometry)})')

        if len(directions) < 2:
            ui.messageBox('Not enough valid edges to compute an average.')
            return

        # ------------------------------------------------------------------
        # 3. Align all vectors so they point in a consistent direction
        # ------------------------------------------------------------------
        ref = directions[0]
        aligned = [ref]
        for d in directions[1:]:
            dot = ref[0] * d[0] + ref[1] * d[1] + ref[2] * d[2]
            if dot < 0:
                aligned.append((-d[0], -d[1], -d[2]))
            else:
                aligned.append(d)

        # ------------------------------------------------------------------
        # 4. Compute the average direction
        # ------------------------------------------------------------------
        n = len(aligned)
        avg_x = sum(d[0] for d in aligned) / n
        avg_y = sum(d[1] for d in aligned) / n
        avg_z = sum(d[2] for d in aligned) / n

        avg_len = math.sqrt(avg_x ** 2 + avg_y ** 2 + avg_z ** 2)
        if avg_len < 1e-10:
            ui.messageBox(
                'The selected edges cancel each other out \u2014 '
                'no meaningful average axis exists.'
            )
            return

        avg_x /= avg_len
        avg_y /= avg_len
        avg_z /= avg_len

        # Centroid of all edge midpoints (anchor for the axis)
        cx = sum(m[0] for m in midpoints) / len(midpoints)
        cy = sum(m[1] for m in midpoints) / len(midpoints)
        cz = sum(m[2] for m in midpoints) / len(midpoints)

        # ------------------------------------------------------------------
        # 5. Create construction axis + perpendicular plane
        # ------------------------------------------------------------------
        origin = adsk.core.Point3D.create(cx, cy, cz)
        direction = adsk.core.Vector3D.create(avg_x, avg_y, avg_z)

        axes = comp.constructionAxes
        planes = comp.constructionPlanes

        if design.designType == adsk.fusion.DesignTypes.DirectDesignType:
            # --- Direct modelling mode ---

            # Construction axis via InfiniteLine3D
            axisInput = axes.createInput()
            infLine = adsk.core.InfiniteLine3D.create(origin, direction)
            axisInput.setByLine(infLine)
            axis = axes.add(axisInput)

            # Perpendicular plane via Plane object (normal = average direction)
            planeInput = planes.createInput()
            perpPlane = adsk.core.Plane.create(origin, direction)
            planeInput.setByPlane(perpPlane)
            plane = planes.add(planeInput)

        else:
            # --- Parametric (timeline) mode ---
            # setByLine and setByPlane are direct-mode only, so we use a
            # helper sketch line and setByDistanceOnPath to get a plane
            # perpendicular to that line at its midpoint.

            sketches = comp.sketches
            sketch = sketches.add(comp.xYConstructionPlane)
            sketch.name = 'AverageAxis_Helper'
            sketch.isVisible = False

            # Draw a construction line along the average direction (10 cm)
            half = 5.0  # cm
            p1 = adsk.core.Point3D.create(
                cx - direction.x * half,
                cy - direction.y * half,
                cz - direction.z * half
            )
            p2 = adsk.core.Point3D.create(
                cx + direction.x * half,
                cy + direction.y * half,
                cz + direction.z * half
            )
            skLine = sketch.sketchCurves.sketchLines.addByTwoPoints(p1, p2)
            skLine.isConstruction = True

            # Construction axis from the sketch line
            axisInput = axes.createInput()
            axisInput.setByEdge(skLine)
            axis = axes.add(axisInput)

            # Perpendicular plane at the midpoint of the sketch line (0.5 = 50%)
            planeInput = planes.createInput()
            planeInput.setByDistanceOnPath(
                skLine,
                adsk.core.ValueInput.createByReal(0.5)
            )
            plane = planes.add(planeInput)

        axis.name = 'AverageAxis'
        plane.name = 'AverageAxis_Perp'

        # ------------------------------------------------------------------
        # 6. Report result
        # ------------------------------------------------------------------
        msg = (
            f'Average axis + perpendicular plane created from {n} edges!\n\n'
            f'Direction:  ({avg_x:.6f}, {avg_y:.6f}, {avg_z:.6f})\n'
            f'Origin:       ({cx:.4f}, {cy:.4f}, {cz:.4f}) cm\n\n'
            f'Component: {comp.name}\n\n'
            f'Individual edge directions (aligned):\n'
        )
        for label, d in zip(edge_labels, aligned):
            msg += f'  {label}: ({d[0]:.6f}, {d[1]:.6f}, {d[2]:.6f})\n'

        ui.messageBox(msg, 'AverageAxis')

    except:
        if ui:
            ui.messageBox('Failed:\n{}'.format(traceback.format_exc()))