# -----------------------------------------------------------------------------
# Open Atrium Stair, Levels 1 -> 2 -> 3  -  Dynamo 3.x (CPython3 / PythonNet3)
#
# Two open switchback stairs that hang into the 40' x 40' atrium from its
# slab edges (no shaft needed - the atrium is already open at Levels 2-3):
#
#   * Level 1 -> 2 : on the SOUTH atrium edge. Starts on the atrium floor,
#                    climbs north, turns at a landing, and arrives back on
#                    the Level 2 slab edge (corridor).
#   * Level 2 -> 3 : on the NORTH atrium edge. Starts at the Level 2 slab
#                    edge, climbs south, turns, arrives on the Level 3 edge.
#   (Rotated 180 degrees from each other so the flights never stack - full
#    headroom everywhere. Level 4 is the roof, open over the atrium, so the
#    stair stops at Level 3.)
#
#   Atrium plan (# = slab edge, ^ v = flight direction):
#
#                     north edge (y = 60')
#        #############[ v  |  ^ ]#############   <- 2->3 starts/ends here
#                     [ v  |  ^ ]
#                     [  landing ]   y ~ 47.7'
#
#                        open atrium
#
#                     [  landing ]   y ~ 32.3'
#                     [ ^  |  v ]
#        #############[ ^  |  v ]#############   <- 1->2 starts on L1 floor,
#                     south edge (y = 20')            arrives on L2 edge
#
#   * 5'-0" wide flights, 4" between flights, 11" treads, equal risers
#     (max 7" from the stair type; even riser count so both flights match)
#   * Railings on every flight
#
# NOTE: an open stair through an atrium is a "communicating" stair - it does
# not count as a required exit. The building still needs a second exit stair.
#
# Safe to re-run: skips a story that already has a stair starting on that
# level within this footprint.
#
# HOW TO USE
#   Dynamo 3.x > Python Script node (CPython3 or PythonNet3) > paste > Run.
#   Ctrl+Z in Revit undoes it (one undo step per stair, plus railings).
# -----------------------------------------------------------------------------

import clr
import math
import uuid

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

from Autodesk.Revit.DB import (
    XYZ, Line, Grid, Level, Transaction, FilteredElementCollector,
    BuiltInCategory, BuiltInParameter, IFailuresPreprocessor, FailureProcessingResult,
)
import Autodesk.Revit.DB as _DB
import Autodesk.Revit.DB.Architecture as _ARCH
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager


def _revit_type(name):
    for ns in (_ARCH, _DB):
        t = getattr(ns, name, None)
        if t is not None:
            return t
    raise ImportError("Revit API type '{}' not found".format(name))


Stairs = _revit_type("Stairs")
StairsRun = _revit_type("StairsRun")
StairsLanding = _revit_type("StairsLanding")
StairsEditScope = _revit_type("StairsEditScope")
StairsRunJustification = _revit_type("StairsRunJustification")
Railing = _revit_type("Railing")
RailingType = _revit_type("RailingType")
RailingPlacementPosition = _revit_type("RailingPlacementPosition")

doc = DocumentManager.Instance.CurrentDBDocument

# ---------------------------------------------------------------- settings ---
# (base level, top level, atrium edge the stair hangs from)
STORIES = [("Level 1", "Level 2", "S"), ("Level 2", "Level 3", "N")]
FLIGHT_WIDTH = 5.0
FLIGHT_GAP = 4.0 / 12
TREAD = 11.0 / 12
MAX_RISER = 7.0 / 12          # used if the stair type doesn't say
EDGE_SETBACK = 0.0            # ft from the atrium slab edge to the end of the flights
TOL = 0.01

report = []

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))

# ------------------------------------------------------ atrium from grids ---
xs, ys = [], []
for g in FilteredElementCollector(doc).OfClass(Grid):
    if isinstance(g.Curve, Line):
        p0, p1 = g.Curve.GetEndPoint(0), g.Curve.GetEndPoint(1)
        if abs(p1.Y - p0.Y) > abs(p1.X - p0.X):
            xs.append(p0.X)
        else:
            ys.append(p0.Y)


def uniq(vals):
    out = []
    for v in sorted(vals):
        if not out or abs(v - out[-1]) > TOL:
            out.append(v)
    return out


def bay(vals, c):
    for lo, hi in zip(vals, vals[1:]):
        if lo - TOL <= c <= hi + TOL:
            return lo, hi
    raise Exception("No grid bay at the building center.")


xs, ys = uniq(xs), uniq(ys)
if len(xs) < 3 or len(ys) < 3:
    raise Exception("Need at least 3 grids each way to find the atrium.")
ax0, ax1 = bay(xs, (xs[0] + xs[-1]) / 2.0)
ay0, ay1 = bay(ys, (ys[0] + ys[-1]) / 2.0)
acx, acy = (ax0 + ax1) / 2.0, (ay0 + ay1) / 2.0

# Edge frames: point on edge (center), n = into the atrium, a = along the edge
EDGES = {
    "S": (XYZ(acx, ay0, 0), XYZ(0, 1, 0), XYZ(1, 0, 0)),
    "N": (XYZ(acx, ay1, 0), XYZ(0, -1, 0), XYZ(-1, 0, 0)),
    "W": (XYZ(ax0, acy, 0), XYZ(1, 0, 0), XYZ(0, -1, 0)),
    "E": (XYZ(ax1, acy, 0), XYZ(-1, 0, 0), XYZ(0, 1, 0)),
}


class _StairFailures(IFailuresPreprocessor):
    __namespace__ = "AtriumStair_" + uuid.uuid4().hex

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


def footprint(edge, depth):
    c, n, a = EDGES[edge]
    half = FLIGHT_WIDTH + FLIGHT_GAP / 2.0
    pts = [c + a.Multiply(s * half) + n.Multiply(d) for s in (-1, 1) for d in (EDGE_SETBACK, depth)]
    return (min(p.X for p in pts), min(p.Y for p in pts), max(p.X for p in pts), max(p.Y for p in pts))


def stair_exists(base_level, edge):
    fx0, fy0, fx1, fy1 = footprint(edge, 20.0)
    for s in FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_Stairs).WhereElementIsNotElementType():
        if not isinstance(s, Stairs):
            continue
        p = s.get_Parameter(BuiltInParameter.STAIRS_BASE_LEVEL_PARAM)
        if not p or p.AsElementId() != base_level.Id:
            continue
        bb = s.get_BoundingBox(None)
        if bb and bb.Min.X < fx1 and bb.Max.X > fx0 and bb.Min.Y < fy1 and bb.Max.Y > fy0:
            return True
    return False


TransactionManager.Instance.ForceCloseTransaction()

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

    c, n, a = EDGES[edge]
    lane1 = c + a.Multiply(-(FLIGHT_WIDTH + FLIGHT_GAP) / 2.0)   # up-flight
    lane2 = c + a.Multiply((FLIGHT_WIDTH + FLIGHT_GAP) / 2.0)    # return flight
    height = top.Elevation - base.Elevation

    scope = StairsEditScope(doc, "Atrium Stair {} - {}".format(base_name, top_name))
    t = None
    step = "start stair edit scope"
    try:
        if doc.IsModifiable:
            raise Exception("a transaction is still open in the document before the stair edit started")
        stairs_id = scope.Start(base.Id, top.Id)
        step = "start runs transaction"
        t = Transaction(doc, "Atrium stair runs")
        t.Start()

        step = "set riser count / tread depth"
        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))
        if risers % 2:
            risers += 1                      # equal flights so both end at the slab edge
        stairs.DesiredRisersNumber = risers
        try:
            stairs.ActualTreadDepth = TREAD
        except Exception:
            pass
        tread = stairs.ActualTreadDepth
        r = risers // 2

        # Run 1: from the slab edge into the atrium (path line at the story's elevation)
        step = "create run 1"
        run_len = (r - 1) * tread
        s0 = EDGE_SETBACK
        run1 = None
        last_err = None
        for zb in (base.ProjectElevation, base.Elevation):
            try:
                p_start = lane1 + n.Multiply(s0)
                p_end = lane1 + n.Multiply(s0 + run_len)
                run1 = StairsRun.CreateStraightRun(
                    doc, stairs_id,
                    Line.CreateBound(XYZ(p_start.X, p_start.Y, zb), XYZ(p_end.X, p_end.Y, zb)),
                    StairsRunJustification.Center)
                base_z = zb
                break
            except Exception as ex:
                last_err = ex
        if run1 is None:
            raise last_err
        run1.ActualRunWidth = FLIGHT_WIDTH

        # Run 2: back toward the slab edge
        step = "create run 2"
        end_riser = run1.ActualRisersNumber - (r - 1)
        if abs(run1.BaseElevation) < 0.01 and abs(base_z) > 0.01:
            z1 = base_z + run1.TopElevation
        else:
            z1 = run1.TopElevation
        run2_len = (r - end_riser) * tread
        p_start = lane2 + n.Multiply(s0 + run_len)
        p_end = lane2 + n.Multiply(s0 + run_len - run2_len)
        run2 = StairsRun.CreateStraightRun(
            doc, stairs_id,
            Line.CreateBound(XYZ(p_start.X, p_start.Y, z1), XYZ(p_end.X, p_end.Y, z1)),
            StairsRunJustification.Center)
        run2.ActualRunWidth = FLIGHT_WIDTH

        step = "create landing"
        StairsLanding.CreateAutomaticLanding(doc, run1.Id, run2.Id)
        step = "commit runs transaction"
        t.Commit()
        t = None
        step = "commit stair"
        scope.Commit(_StairFailures())

        stairs = doc.GetElement(stairs_id)
        gap = run_len - run2_len
        report.append("{} -> {} ({} edge): {} risers @ {:.2f}\", {:.0f}\" treads, "
                      "projects {:.1f}' into the atrium".format(
                          base_name, top_name, {"S": "south", "N": "north", "E": "east", "W": "west"}[edge],
                          stairs.ActualRisersNumber, stairs.ActualRiserHeight * 12,
                          stairs.ActualTreadDepth * 12, s0 + run_len + FLIGHT_WIDTH))
        if abs(gap) > 0.05:
            report.append("  Note: return flight ends {:.0f}\" {} the slab edge - check the arrival".format(
                abs(gap) * 12, "short of" if gap > 0 else "past"))
    except Exception as ex:
        if t is not None and t.HasStarted() and not t.HasEnded():
            t.RollBack()
        if scope.IsActive:
            scope.Cancel()
        scope.Dispose()
        raise Exception("Atrium stair {} -> {} failed at step '{}': {}".format(base_name, top_name, step, ex))
    scope.Dispose()

    t = Transaction(doc, "Atrium stair railings")
    t.Start()
    try:
        rtypes = list(FilteredElementCollector(doc).OfClass(RailingType))
        if rtypes:
            Railing.Create(doc, stairs_id, rtypes[0].Id, RailingPlacementPosition.Treads)
        t.Commit()
    except Exception as ex:
        t.RollBack()
        report.append("  Railings not added ({}) - add them with Railing > Place on Stair".format(ex))

report.append("Atrium: x {:.0f}'-{:.0f}', y {:.0f}'-{:.0f}'. This is a communicating stair - "
              "it does not count as a required exit.".format(ax0, ax1, ay0, ay1))
OUT = report
