# -----------------------------------------------------------------------------
# Place Steel Columns + Beams on Grid  -  Dynamo 3.x (CPython3 / PythonNet3)
#
#   * W10x49 structural column at every grid intersection, Level 1 -> Level 2
#   * W12x26 beam along every grid line between adjacent columns, at Level 2
#     (top of steel at Level 2, beams frame into the column tops)
#
# Works from the grids actually in the model, so it follows whatever grid
# layout exists. With the atrium layout (grids 1,2,3,5 x A,C,D,F) that is
# 16 columns and 24 beams, with the 40' x 40' atrium bay left open.
#
# Family types: uses W10X49 / W12X26 if already loaded. If not, it tries to
# load just those types from the standard Revit library:
#   Structural Columns\Steel\W Shapes-Column.rfa
#   Structural Framing\Steel\W Shapes.rfa
# If the library isn't installed, load them via Insert > Load Family and re-run.
#
# Safe to re-run: skips locations that already have a column, and beams that
# already exist between the same two points.
#
# HOW TO USE
#   Dynamo 3.x > Python Script node (CPython3 or PythonNet3) > paste > Run.
#   Ctrl+Z in Revit undoes the whole run.
# -----------------------------------------------------------------------------

import clr
import glob
import os

clr.AddReference("RevitAPI")
clr.AddReference("RevitServices")

from Autodesk.Revit.DB import (
    XYZ, Line, Grid, Level, FamilySymbol, FamilyInstance, LocationPoint, LocationCurve,
    FilteredElementCollector, BuiltInCategory, BuiltInParameter,
)
from Autodesk.Revit.DB.Structure import StructuralType
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager

doc = DocumentManager.Instance.CurrentDBDocument

# ---------------------------------------------------------------- settings ---
COLUMN_TYPE = "W10x49"
BEAM_TYPE = "W12x26"
BASE_LEVEL = "Level 1"
TOP_LEVEL = "Level 2"

COLUMN_FAMILY_FILE = os.path.join("Structural Columns", "Steel", "W Shapes-Column.rfa")
BEAM_FAMILY_FILE = os.path.join("Structural Framing", "Steel", "W Shapes.rfa")

TOL = 0.01  # feet


def norm(s):
    return s.replace(" ", "").upper()


# ------------------------------------------------------------- levels -------
levels = {l.Name: l for l in FilteredElementCollector(doc).OfClass(Level)}
for lv in (BASE_LEVEL, TOP_LEVEL):
    if lv not in levels:
        raise Exception("Level '{}' not found. Levels in project: {}".format(
            lv, ", ".join(sorted(levels))))
base_level, top_level = levels[BASE_LEVEL], levels[TOP_LEVEL]
if top_level.ProjectElevation <= base_level.ProjectElevation:
    raise Exception("{} must be above {}.".format(TOP_LEVEL, BASE_LEVEL))


# ------------------------------------------------------- family types -------
def find_symbol(category, type_name):
    for s in FilteredElementCollector(doc).OfCategory(category).OfClass(FamilySymbol):
        if norm(s.Name) == norm(type_name):
            return s
    return None


def library_files(relative_path):
    patterns = [
        os.path.join(r"C:\ProgramData\Autodesk", "RVT *", "Libraries", "*", relative_path),
        os.path.join(r"C:\ProgramData\Autodesk", "RVT *", "Libraries", "*", "*", relative_path),
    ]
    found = []
    for p in patterns:
        found.extend(glob.glob(p))
    return sorted(found, reverse=True)  # newest Revit version first


def get_symbol(category, type_name, relative_path, report):
    sym = find_symbol(category, type_name)
    if sym:
        return sym
    for path in library_files(relative_path):
        for candidate in (type_name.upper(), type_name):
            try:
                if doc.LoadFamilySymbol(path, candidate):
                    sym = find_symbol(category, type_name)
                    if sym:
                        report.append("Loaded {} from {}".format(sym.Name, path))
                        return sym
            except Exception:
                pass
    raise Exception(
        "Type '{}' isn't loaded and couldn't be found in the Revit library. "
        "Load it with Insert > Load Family ({}) and run again.".format(type_name, relative_path))


# ---------------------------------------------------- grid intersections ----
def flat(p):
    return XYZ(p.X, p.Y, 0)


grids = []
for g in FilteredElementCollector(doc).OfClass(Grid):
    if isinstance(g.Curve, Line):
        grids.append((g.Name, flat(g.Curve.GetEndPoint(0)), flat(g.Curve.GetEndPoint(1))))


def intersect(a0, a1, b0, b1):
    """2D segment intersection; returns XYZ or None."""
    d1x, d1y = a1.X - a0.X, a1.Y - a0.Y
    d2x, d2y = b1.X - b0.X, b1.Y - b0.Y
    den = d1x * d2y - d1y * d2x
    if abs(den) < 1e-9:
        return None  # parallel
    t = ((b0.X - a0.X) * d2y - (b0.Y - a0.Y) * d2x) / den
    u = ((b0.X - a0.X) * d1y - (b0.Y - a0.Y) * d1x) / den
    if -1e-6 <= t <= 1 + 1e-6 and -1e-6 <= u <= 1 + 1e-6:
        return XYZ(a0.X + t * d1x, a0.Y + t * d1y, 0)
    return None


points = []                                 # unique intersection points
on_grid = {name: [] for name, _, _ in grids}  # grid name -> [(param, point)]


def add_unique(p):
    for q in points:
        if q.DistanceTo(p) < TOL:
            return q
    points.append(p)
    return p


for i in range(len(grids)):
    for j in range(i + 1, len(grids)):
        n1, a0, a1 = grids[i]
        n2, b0, b1 = grids[j]
        p = intersect(a0, a1, b0, b1)
        if p is None:
            continue
        p = add_unique(p)
        on_grid[n1].append((p.DistanceTo(a0), p))
        on_grid[n2].append((p.DistanceTo(b0), p))

if not points:
    raise Exception("No grid intersections found.")


def label(p):
    names = [n for n, lst in on_grid.items() if any(q.DistanceTo(p) < TOL for _, q in lst)]
    return "/".join(sorted(names, key=lambda s: (s.isalpha(), s)))


# ----------------------------------------------------- existing elements ----
existing_cols = []
for c in (FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_StructuralColumns)
          .WhereElementIsNotElementType()):
    loc = c.Location
    if isinstance(loc, LocationPoint):
        existing_cols.append(flat(loc.Point))

existing_beams = []
for b in (FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_StructuralFraming)
          .WhereElementIsNotElementType()):
    loc = b.Location
    if isinstance(loc, LocationCurve):
        existing_beams.append((flat(loc.Curve.GetEndPoint(0)), flat(loc.Curve.GetEndPoint(1))))


def column_exists(p):
    return any(q.DistanceTo(p) < TOL for q in existing_cols)


def beam_exists(p, q):
    for a, b in existing_beams:
        if (a.DistanceTo(p) < TOL and b.DistanceTo(q) < TOL) or \
           (a.DistanceTo(q) < TOL and b.DistanceTo(p) < TOL):
            return True
    return False


# ---------------------------------------------------------------- place -----
report = []
cols_made, cols_skipped, beams_made, beams_skipped = [], [], 0, 0
z_top = top_level.ProjectElevation

TransactionManager.Instance.EnsureInTransaction(doc)
try:
    col_sym = get_symbol(BuiltInCategory.OST_StructuralColumns, COLUMN_TYPE, COLUMN_FAMILY_FILE, report)
    beam_sym = get_symbol(BuiltInCategory.OST_StructuralFraming, BEAM_TYPE, BEAM_FAMILY_FILE, report)
    if not col_sym.IsActive:
        col_sym.Activate()
    if not beam_sym.IsActive:
        beam_sym.Activate()
    doc.Regenerate()

    # Columns
    for p in points:
        if column_exists(p):
            cols_skipped.append(label(p))
            continue
        col = doc.Create.NewFamilyInstance(p, col_sym, base_level, StructuralType.Column)
        col.get_Parameter(BuiltInParameter.FAMILY_BASE_LEVEL_PARAM).Set(base_level.Id)
        col.get_Parameter(BuiltInParameter.FAMILY_BASE_LEVEL_OFFSET_PARAM).Set(0.0)
        col.get_Parameter(BuiltInParameter.FAMILY_TOP_LEVEL_PARAM).Set(top_level.Id)
        col.get_Parameter(BuiltInParameter.FAMILY_TOP_LEVEL_OFFSET_PARAM).Set(0.0)
        existing_cols.append(p)
        cols_made.append(label(p))

    # Beams: along each grid, between consecutive intersections
    for name, lst in on_grid.items():
        lst.sort(key=lambda t: t[0])
        for (_, p), (_, q) in zip(lst, lst[1:]):
            if p.DistanceTo(q) < TOL:
                continue
            if beam_exists(p, q):
                beams_skipped += 1
                continue
            line = Line.CreateBound(XYZ(p.X, p.Y, z_top), XYZ(q.X, q.Y, z_top))
            doc.Create.NewFamilyInstance(line, beam_sym, top_level, StructuralType.Beam)
            existing_beams.append((p, q))
            beams_made += 1
finally:
    TransactionManager.Instance.TransactionTaskDone()

report.append("Columns ({}) placed: {} at {}".format(col_sym.Name, len(cols_made), ", ".join(cols_made) or "-"))
if cols_skipped:
    report.append("Columns skipped (already there): {}".format(", ".join(cols_skipped)))
report.append("Beams ({}) placed: {}".format(beam_sym.Name, beams_made))
if beams_skipped:
    report.append("Beams skipped (already there): {}".format(beams_skipped))
report.append("{} -> {}, beams at {}".format(BASE_LEVEL, TOP_LEVEL, TOP_LEVEL))

OUT = report
