# -----------------------------------------------------------------------------
# SW Corner Stair, Levels 1 -> 4  -  Dynamo 3.x (CPython3 / PythonNet3)
#
# A switchback (U-shaped) stair in the southwest corner of the building,
# built as one stair per story (1-2, 2-3, 3-4) stacked on the same footprint,
# so each flight fits that story's actual floor-to-floor height.
#
#   Plan (SW corner; x east, y north; ~ values for 10' floor-to-floor):
#
#     y=15.3 ---------------------------[door]--   corridor wall (Levels 2-3)
#            |                                 |
#            |   floor landing / arrival       |
#     y~8.7  +-------+----------------+        |
#            |       |  run 2  --->   |  floor |
#            |landing|----------------|  ...   |   x = 20  (demising wall)
#            |       |  run 1  <---   |  start |
#     y=1    +-------+----------------+        |
#          x=1     x~4.7            x~12
#      (curtain wall at x=-1, y=-1; corner column at 0,0 is clear)
#
#   * Risers: max 7" (from the stair type), equal risers per story,
#     11" treads, 44" wide flights, 4" between flights
#   * Run 1 climbs west from the east end, landing at the west end,
#     run 2 returns east and arrives at the floor near the corridor door
#   * Railings added on both sides of every flight
#   * Shaft opening through Level 2, Level 3 and the Level 4 roof over the
#     flights and intermediate landing (floor kept at the arrival end)
#   * On Levels 2-3 the SW corner office becomes the stair room: its room
#     is renamed "Stair 1" and the furniture inside it is removed (the
#     corridor door stays as the stair door)
#
# Safe to re-run: skips any story that already has a stair starting on that
# level, and doesn't add a second shaft.
#
# HOW TO USE
#   Dynamo 3.x > Python Script node (CPython3 or PythonNet3) > paste > Run.
#   Ctrl+Z in Revit undoes it (one undo step per story + one for openings).
# -----------------------------------------------------------------------------

import clr
import math
import uuid

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

from Autodesk.Revit.DB import (
    XYZ, Line, CurveArray, Level, Wall, WallKind, LocationCurve, LocationPoint,
    Transaction, ElementId, FilteredElementCollector, BuiltInCategory, BuiltInParameter,
    IFailuresPreprocessor, FailureProcessingResult, Opening,
)
from Autodesk.Revit.DB.Architecture import (
    Stairs, StairsRun, StairsLanding, StairsEditScope, StairsRunJustification,
    Railing, RailingType, RailingPlacementPosition,
)
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager

doc = DocumentManager.Instance.CurrentDBDocument

# ---------------------------------------------------------------- settings ---
STORIES = [("Level 1", "Level 2"), ("Level 2", "Level 3"), ("Level 3", "Level 4")]
CLEAR_FROM_CURTAIN_WALL = 2.0   # ft from curtain wall line to stair edge (clears corner column)
FLIGHT_WIDTH = 44.0 / 12
FLIGHT_GAP = 4.0 / 12
TREAD = 11.0 / 12
MAX_RISER = 7.0 / 12            # used if the stair type doesn't say
ARRIVAL_LANDING = 44.0 / 12     # floor needed beyond the runs at each level
EAST_LIMIT_FROM_SW = 21.0       # ft from curtain wall to the corner office's east wall
ROOM_NAME = "Stair 1"
SHAFT_BASE_OFFSET = 1.0         # keep the Level 1 slab
SHAFT_TOP_OFFSET = 2.0          # cut through the Level 4 roof

report = []

# ------------------------------------------------------------- levels -------
levels = {l.Name: l for l in FilteredElementCollector(doc).OfClass(Level)}
for a, b in STORIES:
    for n in (a, b):
        if n not in levels:
            raise Exception("Level '{}' not found.".format(n))

# ------------------------------------------------ SW corner reference -------
cw_pts = []
for w in FilteredElementCollector(doc).OfClass(Wall):
    if w.WallType.Kind == WallKind.Curtain and isinstance(w.Location, LocationCurve):
        c = w.Location.Curve
        cw_pts += [c.GetEndPoint(0), c.GetEndPoint(1)]
if cw_pts:
    sx, sy = min(p.X for p in cw_pts), min(p.Y for p in cw_pts)
else:
    sx, sy = -1.0, -1.0
    report.append("No curtain walls found - assuming the SW corner is at (-1', -1')")

x_land0 = sx + CLEAR_FROM_CURTAIN_WALL            # west edge of intermediate landing
x_land1 = x_land0 + FLIGHT_WIDTH                   # where the runs start/end
y0 = sy + CLEAR_FROM_CURTAIN_WALL
lane1_y = y0 + FLIGHT_WIDTH / 2.0                  # run 1 (south lane)
lane2_y = y0 + FLIGHT_WIDTH + FLIGHT_GAP + FLIGHT_WIDTH / 2.0   # run 2 (north lane)
y1 = y0 + 2 * FLIGHT_WIDTH + FLIGHT_GAP
east_limit = sx + EAST_LIMIT_FROM_SW


# ------------------------------------------------ failure handling ----------
class _StairFailures(IFailuresPreprocessor):
    __namespace__ = "SWStair_" + uuid.uuid4().hex   # unique per run (re-runs in Dynamo)

    def PreprocessFailures(self, failuresAccessor):
        return FailureProcessingResult.Continue


def stair_exists(base_level):
    for s in FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_Stairs).WhereElementIsNotElementType():
        if isinstance(s, Stairs):
            p = s.get_Parameter(BuiltInParameter.STAIRS_BASE_LEVEL_PARAM)
            if p and p.AsElementId() == base_level.Id:
                bb = s.get_BoundingBox(None)
                if bb and bb.Min.X < east_limit and bb.Min.Y < y1 + 1 and bb.Max.X > x_land0 - 1:
                    return True
    return False


# Dynamo holds an open transaction; stairs need their own edit scope.
TransactionManager.Instance.ForceCloseTransaction()

max_east = x_land1
for base_name, top_name in STORIES:
    base, top = levels[base_name], levels[top_name]
    if stair_exists(base):
        report.append("Skipped {} -> {}: a stair already starts on {} here".format(base_name, top_name, base_name))
        continue

    height = top.Elevation - base.Elevation
    scope = StairsEditScope(doc, "SW Stair {} - {}".format(base_name, top_name))
    try:
        stairs_id = scope.Start(base.Id, top.Id)
        t = Transaction(doc, "Stair runs")
        t.Start()
        stairs = doc.GetElement(stairs_id)
        stype = doc.GetElement(stairs.GetTypeId())
        max_riser = MAX_RISER
        try:
            if stype.MaxRiserHeight > 0:
                max_riser = stype.MaxRiserHeight
        except Exception:
            pass
        risers = int(math.ceil(height / max_riser - 1e-6))
        stairs.DesiredRisersNumber = risers
        try:
            stairs.ActualTreadDepth = TREAD
        except Exception:
            pass
        tread = stairs.ActualTreadDepth

        r1 = int(math.ceil(risers / 2.0))
        r2 = risers - r1

        # Run 1: east -> west in the south lane, starting at the base level
        run1_len = (r1 - 1) * tread
        run1 = StairsRun.CreateStraightRun(
            doc, stairs_id,
            Line.CreateBound(XYZ(x_land1 + run1_len, lane1_y, 0), XYZ(x_land1, lane1_y, 0)),
            StairsRunJustification.Center)
        run1.ActualRunWidth = FLIGHT_WIDTH

        # Revit adds an end riser depending on the stair type; adapt run 2 to it
        end_riser = run1.ActualRisersNumber - (r1 - 1)
        z1 = run1.TopElevation
        run2_len = (r2 - end_riser) * tread
        run2 = StairsRun.CreateStraightRun(
            doc, stairs_id,
            Line.CreateBound(XYZ(x_land1, lane2_y, z1), XYZ(x_land1 + run2_len, lane2_y, z1)),
            StairsRunJustification.Center)
        run2.ActualRunWidth = FLIGHT_WIDTH

        StairsLanding.CreateAutomaticLanding(doc, run1.Id, run2.Id)
        t.Commit()
        scope.Commit(_StairFailures())

        stairs = doc.GetElement(stairs_id)
        max_east = max(max_east, x_land1 + max(run1_len, run2_len))
        report.append("{} -> {}: {:.2f}' rise, {} risers @ {:.2f}\", {:.0f}\" treads ({} + {} per flight)".format(
            base_name, top_name, height, stairs.ActualRisersNumber,
            stairs.ActualRiserHeight * 12, stairs.ActualTreadDepth * 12, r1, r2))

        # Railings
        t = Transaction(doc, "Stair railings")
        t.Start()
        try:
            rtypes = list(FilteredElementCollector(doc).OfClass(RailingType))
            if rtypes:
                Railing.Create(doc, stairs_id, rtypes[0].Id, RailingPlacementPosition.Treads)
        except Exception as ex:
            report.append("  Railings not added ({}) - add them with Railing > Place on Stair".format(ex))
        t.Commit()
    except Exception as ex:
        if scope.IsActive:
            scope.Cancel()
        raise Exception("Stair {} -> {} failed: {}".format(base_name, top_name, ex))
    finally:
        scope.Dispose()

if max_east + ARRIVAL_LANDING > east_limit:
    report.append("Warning: flights reach x = {:.1f}'; with a {:.1f}' arrival landing that passes the "
                  "corner office wall at x = {:.1f}'. Check the stair in plan.".format(
                      max_east, ARRIVAL_LANDING, east_limit))

# ---------------------------------------- shaft opening + room cleanup -------
t = Transaction(doc, "Stair shaft and room")
t.Start()
try:
    shaft_x1 = max_east
    existing_shaft = False
    for o in FilteredElementCollector(doc).OfClass(Opening):
        bb = o.get_BoundingBox(None)
        if bb and abs(bb.Min.X - x_land0) < 0.1 and abs(bb.Min.Y - y0) < 0.1:
            existing_shaft = True
    if not existing_shaft:
        bottom, top_level = levels[STORIES[0][0]], levels[STORIES[-1][1]]
        z = bottom.ProjectElevation
        pts = [XYZ(x_land0, y0, z), XYZ(shaft_x1, y0, z), XYZ(shaft_x1, y1, z), XYZ(x_land0, y1, z)]
        profile = CurveArray()
        for i in range(4):
            profile.Append(Line.CreateBound(pts[i], pts[(i + 1) % 4]))
        shaft = doc.Create.NewOpening(bottom, top_level, profile)
        for name, val in (("Base Offset", SHAFT_BASE_OFFSET), ("Top Offset", SHAFT_TOP_OFFSET)):
            p = shaft.LookupParameter(name)
            if p and not p.IsReadOnly:
                p.Set(val)
        report.append("Shaft opening {:.1f}' x {:.1f}' through Levels 2-3 and the roof".format(
            shaft_x1 - x_land0, y1 - y0))
    else:
        report.append("Shaft opening already present - skipped")

    # Levels 2-3: corner office -> stair room, remove its furniture
    cx, cy = (x_land0 + shaft_x1) / 2.0, (y0 + y1) / 2.0
    furn_cats = [BuiltInCategory.OST_Furniture, BuiltInCategory.OST_Casework,
                 BuiltInCategory.OST_FurnitureSystems]
    for base_name, _ in STORIES[1:]:
        lvl = levels[base_name]
        z = lvl.ProjectElevation + 1.0
        room = doc.GetRoomAtPoint(XYZ(cx, cy, z))
        if room is None:
            continue
        old = room.get_Parameter(BuiltInParameter.ROOM_NAME).AsString()
        room.Name = ROOM_NAME
        removed = 0
        for cat in furn_cats:
            for f in list(FilteredElementCollector(doc).OfCategory(cat).WhereElementIsNotElementType()):
                if f.LevelId != lvl.Id or not isinstance(f.Location, LocationPoint):
                    continue
                p = f.Location.Point
                if room.IsPointInRoom(XYZ(p.X, p.Y, z)):
                    doc.Delete(f.Id)
                    removed += 1
        report.append("{}: room '{}' {} renamed '{}', {} furniture items removed".format(
            base_name, old, room.Number, ROOM_NAME, removed))
    t.Commit()
except Exception:
    if t.HasStarted() and not t.HasEnded():
        t.RollBack()
    raise

OUT = report
