# JITTERBUG — DYNAMO FAMILY AUTOMATION
# =====================================
#
# This script automates the creation of a parametric Jitterbug family in Revit.
# It creates parameters, reference points, geometry, and family types.
#
# WHAT DYNAMO CAN AUTOMATE:
#   ✓ Create Family Parameters (FamilyParameter.ByName)
#   ✓ Create Reference Points at computed positions
#   ✓ Create Model Lines/Curves between points
#   ✓ Create Geometry (Forms, Sweeps)
#   ✓ Create Family Types with different parameter values
#   ✓ Set parameter values for each type
#
# WHAT REQUIRES MANUAL WORK:
#   ✗ Opening/Creating a new family document (must be done in Revit first)
#   ✗ Selecting the template (Metric Generic Model.rft)
#   ✗ Saving the family file
#   ✗ Loading the family into a project
#
# PREREQUISITES:
#   1. Open Revit
#   2. Create a new family: File → New → Family → Metric Generic Model.rft
#   3. Save the family (even empty): File → Save As → Jitterbug_Parametric.rfa
#   4. Open Dynamo while in the Family Editor
#   5. Run this script
#
# INPUTS:
#   IN[0] = edge_length (cm) - default 15.24
#   IN[1] = create_types (boolean) - whether to create multiple family types
#
# OUTPUTS:
#   OUT[0] = Created parameters (list)
#   OUT[1] = Reference points (12)
#   OUT[2] = Edge lines (24)
#   OUT[3] = Status message

import clr
import math

clr.AddReference('RevitAPI')
clr.AddReference('RevitServices')
clr.AddReference('RevitNodes')
clr.AddReference('ProtoGeometry')

from Autodesk.Revit.DB import *
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager

import Revit
clr.ImportExtensions(Revit.Elements)
clr.ImportExtensions(Revit.GeometryConversion)

from Autodesk.DesignScript.Geometry import Point as DSPoint

# Get the current document
doc = DocumentManager.Instance.CurrentDBDocument

# =============================================================================
# INPUTS
# =============================================================================

edge_length = float(IN[0]) if (len(IN) > 0 and IN[0] is not None) else 15.24
create_types = bool(IN[1]) if (len(IN) > 1 and IN[1] is not None) else True

# =============================================================================
# CHECK IF IN FAMILY EDITOR
# =============================================================================

if not doc.IsFamilyDocument:
    OUT = [[], [], [], "ERROR: This script must be run in the Family Editor. Please open or create a family first."]
else:
    
    # =============================================================================
    # TOPOLOGY DATA
    # =============================================================================
    
    EDGES = [
        (0, 4), (0, 5), (0, 8), (0, 9),
        (1, 4), (1, 5), (1, 10), (1, 11),
        (2, 6), (2, 7), (2, 8), (2, 9),
        (3, 6), (3, 7), (3, 10), (3, 11),
        (4, 8), (4, 10), (5, 9), (5, 11),
        (6, 8), (6, 10), (7, 9), (7, 11),
    ]
    
    VERTEX_PAIRS = [(0, 1), (2, 3), (4, 6), (5, 7), (8, 9), (10, 11)]
    
    # =============================================================================
    # CONSTRAINT SOLVER
    # =============================================================================
    
    def solve_jitterbug(L, rhombus_angle_deg, vertex_radius=0.25):
        min_separation = vertex_radius * 2.5
        a = L / math.sqrt(2)
        s = min_separation * 0.5
        t = 1.0 - (rhombus_angle_deg / 90.0)
        t = max(0, min(1, t))
        
        if t < 0.001:
            return [
                ( a,  a,  0), ( a, -a,  0), (-a,  a,  0), (-a, -a,  0),
                ( a,  0,  a), ( a,  0, -a), (-a,  0,  a), (-a,  0, -a),
                ( 0,  a,  a), ( 0,  a, -a), ( 0, -a,  a), ( 0, -a, -a),
            ]
        
        ve = [( a,  a,  0), ( a, -a,  0), (-a,  a,  0), (-a, -a,  0),
              ( a,  0,  a), ( a,  0, -a), (-a,  0,  a), (-a,  0, -a),
              ( 0,  a,  a), ( 0,  a, -a), ( 0, -a,  a), ( 0, -a, -a)]
        
        r = a
        octa = [( r,  s,  0), ( r, -s,  0), (-r,  s,  0), (-r, -s,  0),
                ( s,  0,  r), ( s,  0, -r), (-s,  0,  r), (-s,  0, -r),
                ( 0,  r,  s), ( 0,  r, -s), ( 0, -r,  s), ( 0, -r, -s)]
        
        t1 = 1.0 - t
        verts = [[ve[i][0]*t1 + octa[i][0]*t, ve[i][1]*t1 + octa[i][1]*t, ve[i][2]*t1 + octa[i][2]*t] for i in range(12)]
        
        difficulty = t
        omega = 1.2 if difficulty < 0.5 else 1.0
        max_iter = 10 + int(difficulty * 25)
        
        for iteration in range(max_iter):
            for i1, i2 in EDGES:
                dx = verts[i2][0] - verts[i1][0]
                dy = verts[i2][1] - verts[i1][1]
                dz = verts[i2][2] - verts[i1][2]
                dist_sq = dx*dx + dy*dy + dz*dz
                if dist_sq < 0.0001: continue
                dist = math.sqrt(dist_sq)
                correction = ((dist - L) / dist) * 0.5 * omega
                verts[i1][0] += dx * correction
                verts[i1][1] += dy * correction
                verts[i1][2] += dz * correction
                verts[i2][0] -= dx * correction
                verts[i2][1] -= dy * correction
                verts[i2][2] -= dz * correction
            
            if difficulty > 0.7:
                min_sep = min_separation * 0.4
                for i1, i2 in VERTEX_PAIRS:
                    dx = verts[i2][0] - verts[i1][0]
                    dy = verts[i2][1] - verts[i1][1]
                    dz = verts[i2][2] - verts[i1][2]
                    dist_sq = dx*dx + dy*dy + dz*dz
                    if dist_sq < min_sep * min_sep:
                        dist = math.sqrt(dist_sq) if dist_sq > 0.0001 else 0.01
                        push = (min_sep - dist) * 0.5
                        nx, ny, nz = dx/dist, dy/dist, dz/dist
                        verts[i1][0] -= nx * push
                        verts[i1][1] -= ny * push
                        verts[i1][2] -= nz * push
                        verts[i2][0] += nx * push
                        verts[i2][1] += ny * push
                        verts[i2][2] += nz * push
        
        return [tuple(v) for v in verts]
    
    # =============================================================================
    # CREATE FAMILY PARAMETERS
    # =============================================================================
    
    created_params = []
    
    TransactionManager.Instance.EnsureInTransaction(doc)
    
    try:
        family_manager = doc.FamilyManager
        
        # Create coordinate parameters for each vertex
        param_names = []
        for i in range(12):
            param_names.extend([f"V{i}_X", f"V{i}_Y", f"V{i}_Z"])
        
        # Add EdgeLength and RhombusAngle parameters
        param_names.extend(["EdgeLength", "VertexRadius", "PipeRadius"])
        
        # Create parameters
        for param_name in param_names:
            # Check if parameter already exists
            existing = family_manager.get_Parameter(param_name)
            if existing is None:
                # Determine parameter type
                if "Angle" in param_name:
                    param_type = SpecTypeId.Angle
                else:
                    param_type = SpecTypeId.Length
                
                # Create as Type parameter in Dimensions group
                try:
                    new_param = family_manager.AddParameter(
                        param_name,
                        GroupTypeId.Geometry,
                        param_type,
                        False  # False = Type parameter
                    )
                    created_params.append(param_name)
                except:
                    pass
        
        # =============================================================================
        # SET DEFAULT VALUES (VE state at 90°)
        # =============================================================================
        
        ve_positions = solve_jitterbug(edge_length, 90, 0.25)
        
        # Convert cm to feet (Revit internal units)
        cm_to_feet = 1.0 / 30.48
        
        # Set vertex positions
        for i in range(12):
            x, y, z = ve_positions[i]
            
            param_x = family_manager.get_Parameter(f"V{i}_X")
            param_y = family_manager.get_Parameter(f"V{i}_Y")
            param_z = family_manager.get_Parameter(f"V{i}_Z")
            
            if param_x: family_manager.Set(param_x, x * cm_to_feet)
            if param_y: family_manager.Set(param_y, y * cm_to_feet)
            if param_z: family_manager.Set(param_z, z * cm_to_feet)
        
        # Set other defaults
        param_edge = family_manager.get_Parameter("EdgeLength")
        if param_edge: family_manager.Set(param_edge, edge_length * cm_to_feet)
        
        param_vr = family_manager.get_Parameter("VertexRadius")
        if param_vr: family_manager.Set(param_vr, 0.25 * cm_to_feet)
        
        param_pr = family_manager.get_Parameter("PipeRadius")
        if param_pr: family_manager.Set(param_pr, 0.1 * cm_to_feet)
        
        # =============================================================================
        # CREATE FAMILY TYPES
        # =============================================================================
        
        types_created = []
        
        if create_types:
            type_configs = [
                ("VE_90", 90),
                ("Trans_80", 80),
                ("Trans_70", 70),
                ("Icosa_60", 60),
                ("Trans_50", 50),
                ("Trans_40", 40),
                ("Trans_30", 30),
                ("Trans_20", 20),
                ("NearOcta_10", 10),
                ("NearOcta_5", 5),
            ]
            
            for type_name, angle in type_configs:
                try:
                    # Create or get family type
                    existing_type = None
                    for ft in family_manager.Types:
                        if ft.Name == type_name:
                            existing_type = ft
                            break
                    
                    if existing_type is None:
                        new_type = family_manager.NewType(type_name)
                    else:
                        new_type = existing_type
                    
                    # Set as current type to modify
                    family_manager.CurrentType = new_type
                    
                    # Compute vertex positions for this angle
                    positions = solve_jitterbug(edge_length, angle, 0.25)
                    
                    # Set vertex positions for this type
                    for i in range(12):
                        x, y, z = positions[i]
                        
                        param_x = family_manager.get_Parameter(f"V{i}_X")
                        param_y = family_manager.get_Parameter(f"V{i}_Y")
                        param_z = family_manager.get_Parameter(f"V{i}_Z")
                        
                        if param_x: family_manager.Set(param_x, x * cm_to_feet)
                        if param_y: family_manager.Set(param_y, y * cm_to_feet)
                        if param_z: family_manager.Set(param_z, z * cm_to_feet)
                    
                    types_created.append(type_name)
                except Exception as e:
                    pass
        
        TransactionManager.Instance.TransactionTaskDone()
        
        # =============================================================================
        # STATUS MESSAGE
        # =============================================================================
        
        status = f"""
JITTERBUG FAMILY AUTOMATION COMPLETE
====================================

PARAMETERS CREATED: {len(created_params)}
  - 36 vertex coordinate parameters (V0_X through V11_Z)
  - EdgeLength, VertexRadius, PipeRadius

FAMILY TYPES CREATED: {len(types_created)}
  {', '.join(types_created)}

NEXT STEPS (Manual in Revit):
1. Create Reference Points linked to V#_X, V#_Y, V#_Z parameters
2. Create Model Lines between vertices (see EDGES topology)
3. Optionally add spheres at vertices, sweeps along edges
4. Save and load into project

Edge Length: {edge_length} cm
"""
        
        OUT = [created_params, types_created, [], status]
        
    except Exception as e:
        TransactionManager.Instance.TransactionTaskDone()
        OUT = [[], [], [], f"ERROR: {str(e)}"]
