# -----------------------------------------------------------------------------
# Copy Level 2 Steel Frame to Upper Levels  -  Dynamo 3.x (CPython3 / PythonNet3)
#
# Source  : structural columns running Level 1 -> Level 2
#           structural beams with Reference Level = Level 2
# Copies  : columns Level 2 -> Level 3, beams at Level 3
#           columns Level 3 -> Level 4, beams at Level 4
#
# Elements are copied (not redrawn), so type, rotation and other instance
# settings carry over. Each copy is then re-hosted so it is tied to its own
# levels with zero offsets - columns follow if levels move later.
#
# Safe to re-run: skips any copy whose position already has a matching
# column (same base level) or beam (same reference level).
#
# HOW TO USE
#   Dynamo 3.x > Python Script node (CPython3 or PythonNet3) > paste > Run.
#   Ctrl+Z in Revit undoes the whole run.
# -----------------------------------------------------------------------------

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

from System.Collections.Generic import List
from Autodesk.Revit.DB import (
    XYZ, Level, ElementId, ElementTransformUtils, LocationPoint, LocationCurve,
    FilteredElementCollector, BuiltInCategory, BuiltInParameter,
)
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager

doc = DocumentManager.Instance.CurrentDBDocument

# ---------------------------------------------------------------- settings ---
SOURCE_COLUMN_BASE = "Level 1"
SOURCE_COLUMN_TOP = "Level 2"
SOURCE_BEAM_LEVEL = "Level 2"
# (column base level, column top level / beam level)
STORIES = [("Level 2", "Level 3"), ("Level 3", "Level 4")]

TOL = 0.01

levels = {l.Name: l for l in FilteredElementCollector(doc).OfClass(Level)}
needed = {SOURCE_COLUMN_BASE, SOURCE_COLUMN_TOP, SOURCE_BEAM_LEVEL}
for b, t in STORIES:
    needed.update([b, t])
missing = sorted(n for n in needed if n not in levels)
if missing:
    raise Exception("Level(s) not found: {}. Run create_levels_above.py first.".format(", ".join(missing)))


def elev(name):
    return levels[name].ProjectElevation


def pid(elem, bip):
    p = elem.get_Parameter(bip)
    return p.AsElementId() if p else ElementId.InvalidElementId


def flat(p):
    return XYZ(p.X, p.Y, 0)


# ------------------------------------------------------------- collect ------
all_cols = list(FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_StructuralColumns)
                .WhereElementIsNotElementType())
all_beams = list(FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_StructuralFraming)
                 .WhereElementIsNotElementType())

src_cols = [c for c in all_cols
            if pid(c, BuiltInParameter.FAMILY_BASE_LEVEL_PARAM) == levels[SOURCE_COLUMN_BASE].Id
            and pid(c, BuiltInParameter.FAMILY_TOP_LEVEL_PARAM) == levels[SOURCE_COLUMN_TOP].Id
            and isinstance(c.Location, LocationPoint)]
src_beams = [b for b in all_beams
             if pid(b, BuiltInParameter.INSTANCE_REFERENCE_LEVEL_PARAM) == levels[SOURCE_BEAM_LEVEL].Id
             and isinstance(b.Location, LocationCurve)]

if not src_cols and not src_beams:
    raise Exception("No columns ({} -> {}) or beams at {} found to copy.".format(
        SOURCE_COLUMN_BASE, SOURCE_COLUMN_TOP, SOURCE_BEAM_LEVEL))


def column_at(base_name, p):
    for c in all_cols:
        if pid(c, BuiltInParameter.FAMILY_BASE_LEVEL_PARAM) == levels[base_name].Id \
                and isinstance(c.Location, LocationPoint) \
                and flat(c.Location.Point).DistanceTo(p) < TOL:
            return True
    return False


def beam_at(level_name, p, q):
    for b in all_beams:
        if pid(b, BuiltInParameter.INSTANCE_REFERENCE_LEVEL_PARAM) != levels[level_name].Id:
            continue
        if not isinstance(b.Location, LocationCurve):
            continue
        a = flat(b.Location.Curve.GetEndPoint(0))
        z = flat(b.Location.Curve.GetEndPoint(1))
        if (a.DistanceTo(p) < TOL and z.DistanceTo(q) < TOL) or \
           (a.DistanceTo(q) < TOL and z.DistanceTo(p) < TOL):
            return True
    return False


def set_param(elem, bip, value):
    p = elem.get_Parameter(bip)
    if p and not p.IsReadOnly:
        p.Set(value)


def copy(ids, dz):
    id_list = List[ElementId](ids)
    return list(ElementTransformUtils.CopyElements(doc, id_list, XYZ(0, 0, dz)))


# ---------------------------------------------------------------- copy ------
report = []

TransactionManager.Instance.EnsureInTransaction(doc)
try:
    for base_name, top_name in STORIES:
        # Columns
        col_ids = [c.Id for c in src_cols if not column_at(base_name, flat(c.Location.Point))]
        skipped_c = len(src_cols) - len(col_ids)
        new_cols = []
        if col_ids:
            dz = elev(base_name) - elev(SOURCE_COLUMN_BASE)
            for nid in copy(col_ids, dz):
                c = doc.GetElement(nid)
                # Raise the top first so top never sits below base mid-edit.
                set_param(c, BuiltInParameter.FAMILY_TOP_LEVEL_PARAM, levels[top_name].Id)
                set_param(c, BuiltInParameter.FAMILY_TOP_LEVEL_OFFSET_PARAM, 0.0)
                set_param(c, BuiltInParameter.FAMILY_BASE_LEVEL_PARAM, levels[base_name].Id)
                set_param(c, BuiltInParameter.FAMILY_BASE_LEVEL_OFFSET_PARAM, 0.0)
                new_cols.append(c)
            all_cols.extend(new_cols)

        # Beams
        beam_ids = [b.Id for b in src_beams
                    if not beam_at(top_name,
                                   flat(b.Location.Curve.GetEndPoint(0)),
                                   flat(b.Location.Curve.GetEndPoint(1)))]
        skipped_b = len(src_beams) - len(beam_ids)
        new_beams = []
        if beam_ids:
            dz = elev(top_name) - elev(SOURCE_BEAM_LEVEL)
            for nid in copy(beam_ids, dz):
                b = doc.GetElement(nid)
                set_param(b, BuiltInParameter.INSTANCE_REFERENCE_LEVEL_PARAM, levels[top_name].Id)
                set_param(b, BuiltInParameter.STRUCTURAL_BEAM_END0_ELEVATION, 0.0)
                set_param(b, BuiltInParameter.STRUCTURAL_BEAM_END1_ELEVATION, 0.0)
                new_beams.append(b)
            all_beams.extend(new_beams)

        line = "{} -> {}: {} columns, {} beams at {}".format(
            base_name, top_name, len(new_cols), len(new_beams), top_name)
        if skipped_c or skipped_b:
            line += " (skipped {} columns, {} beams already there)".format(skipped_c, skipped_b)
        report.append(line)

    doc.Regenerate()

    # Verify every copy ended up on the intended levels.
    problems = []
    for base_name, top_name in STORIES:
        for c in all_cols:
            if pid(c, BuiltInParameter.FAMILY_BASE_LEVEL_PARAM) == levels[base_name].Id and \
               pid(c, BuiltInParameter.FAMILY_TOP_LEVEL_PARAM) != levels[top_name].Id:
                problems.append("Column {} base {} but top is not {}".format(c.Id, base_name, top_name))
    if problems:
        report.append("Check these:")
        report.extend(problems)
finally:
    TransactionManager.Instance.TransactionTaskDone()

report.insert(0, "Source: {} columns ({} -> {}), {} beams at {}".format(
    len(src_cols), SOURCE_COLUMN_BASE, SOURCE_COLUMN_TOP, len(src_beams), SOURCE_BEAM_LEVEL))

OUT = report