# -----------------------------------------------------------------------------
# Floors + Roof with Central Atrium  -  Dynamo 3.x (CPython3 / PythonNet3)
#
#   Level 1 : "Concrete Slab - 6in"          6" cast-in-place concrete
#             (solid by default - the atrium needs a floor; set
#              LEVEL1_ATRIUM_OPENING = True to cut it here too)
#   Level 2 : "Bubble Deck Slab - 11in"      voided biaxial slab, with atrium opening
#   Level 3 : "Bubble Deck Slab - 11in"      voided biaxial slab, with atrium opening
#   Level 4 : roof (project's basic roof type), flat, with atrium opening,
#             extended over the south cantilever to the spandrel line
#
# Bubble deck type layers (top -> bottom):
#   3" Concrete top layer | 5" "Concrete - Voided (Bubble Deck)" core | 3" Concrete bottom
#   The plastic voids aren't modeled as geometry; the voided core material
#   identifies the zone for schedules and takeoffs.
#
# Geometry is derived from the model:
#   * Perimeter = outer grid lines + EDGE_OFFSET (matches the curtain walls)
#   * Atrium   = the grid bay at the center of the building (grids 2-3 x D-C)
#   * Roof     = perimeter, grown to cover any Level 4 framing beyond it
#
#   Plan, Levels 2-4:
#       +--------------------------------------+
#       |                                      |
#       |          +--------------+            |
#       |          |   ATRIUM     |            |
#       |          |  (opening)   |            |
#       |          +--------------+            |
#       |                                      |
#       +--------------------------------------+
#       :.... roof only: overhang to y=-10' ...:
#
# Safe to re-run: skips any level that already has a floor/roof of that type.
#
# 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, Line, CurveLoop, CurveArray, ModelCurveArray, Grid, Level, Floor, FloorType,
    RoofType, FootPrintRoof, Material, CompoundStructure, CompoundStructureLayer,
    MaterialFunctionAssignment, LocationCurve, ElementId,
    FilteredElementCollector, BuiltInCategory, BuiltInParameter,
)
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager

doc = DocumentManager.Instance.CurrentDBDocument

# ---------------------------------------------------------------- settings ---
EDGE_OFFSET = 1.0                 # feet outside outer grids (curtain wall line)
LEVEL1_ATRIUM_OPENING = False     # True = cut the atrium out of the Level 1 slab too
ROOF_COVERS_LEVEL4_FRAMING = True # extend roof over cantilevers at Level 4

SLAB_TYPE = "Concrete Slab - 6in"
SLAB_THICKNESS = 6.0 / 12

BUBBLE_TYPE = "Bubble Deck Slab - 11in"
BUBBLE_LAYERS = [(3.0 / 12, "Concrete"),
                 (5.0 / 12, "Concrete - Voided (Bubble Deck)"),
                 (3.0 / 12, "Concrete")]

PLAN = [
    # (level, kind, type name, atrium opening)
    ("Level 1", "floor", SLAB_TYPE, LEVEL1_ATRIUM_OPENING),
    ("Level 2", "floor", BUBBLE_TYPE, True),
    ("Level 3", "floor", BUBBLE_TYPE, True),
    ("Level 4", "roof", None, True),
]

TOL = 0.01

# ------------------------------------------------------------- levels -------
levels = {l.Name: l for l in FilteredElementCollector(doc).OfClass(Level)}
missing = [p[0] for p in PLAN if p[0] not in levels]
if missing:
    raise Exception("Level(s) not found: {}".format(", ".join(missing)))

# ------------------------------------------------- grid-based geometry ------
xs, ys = [], []
for g in FilteredElementCollector(doc).OfClass(Grid):
    if not isinstance(g.Curve, Line):
        continue
    a, b = g.Curve.GetEndPoint(0), g.Curve.GetEndPoint(1)
    if abs(b.Y - a.Y) > abs(b.X - a.X):
        xs.append(a.X)      # vertical grid
    else:
        ys.append(a.Y)      # horizontal grid


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


xs, ys = uniq(xs), uniq(ys)
if len(xs) < 3 or len(ys) < 3:
    raise Exception("Need at least 3 vertical and 3 horizontal grids to find the atrium bay.")

cx, cy = (xs[0] + xs[-1]) / 2.0, (ys[0] + ys[-1]) / 2.0


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


ax0, ax1 = bay_around(xs, cx)
ay0, ay1 = bay_around(ys, cy)

perim = [xs[0] - EDGE_OFFSET, ys[0] - EDGE_OFFSET, xs[-1] + EDGE_OFFSET, ys[-1] + EDGE_OFFSET]


def level_framing_extents(level):
    """Bounding box (x0,y0,x1,y1) of beams referenced to this level, or None."""
    pts = []
    for b in (FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_StructuralFraming)
              .WhereElementIsNotElementType()):
        p = b.get_Parameter(BuiltInParameter.INSTANCE_REFERENCE_LEVEL_PARAM)
        if p and p.AsElementId() == level.Id and isinstance(b.Location, LocationCurve):
            c = b.Location.Curve
            pts.extend([c.GetEndPoint(0), c.GetEndPoint(1)])
    if not pts:
        return None
    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 rect_lines(x0, y0, x1, y1, z, ccw=True):
    pts = [XYZ(x0, y0, z), XYZ(x1, y0, z), XYZ(x1, y1, z), XYZ(x0, y1, z)]
    if not ccw:
        pts.reverse()
    return [Line.CreateBound(pts[i], pts[(i + 1) % 4]) for i in range(4)]


def curve_loop(lines):
    loop = CurveLoop()
    for ln in lines:
        loop.Append(ln)
    return loop


# ------------------------------------------------------ types & materials ---
def type_name(e):
    p = e.get_Parameter(BuiltInParameter.SYMBOL_NAME_PARAM)
    return p.AsString() if p else e.Name


def get_material(name, report):
    mats = list(FilteredElementCollector(doc).OfClass(Material))
    for m in mats:
        if m.Name.lower() == name.lower():
            return m.Id
    if name.lower() == "concrete":
        for m in mats:  # e.g. "Concrete, Cast-in-Place gray"
            if m.Name.lower().startswith("concrete"):
                return m.Id
    mid = Material.Create(doc, name)
    report.append("Created material '{}'".format(name))
    return mid


def get_floor_type(name, layers, report):
    fts = [ft for ft in FilteredElementCollector(doc).OfClass(FloorType) if not ft.IsFoundationSlab]
    for ft in fts:
        if type_name(ft) == name:
            return ft
    source = next((ft for ft in fts if "generic" in type_name(ft).lower()), None) or fts[0]
    ft = source.Duplicate(name)
    cs_layers = List[CompoundStructureLayer]()
    for width, mat in layers:
        cs_layers.Add(CompoundStructureLayer(width, MaterialFunctionAssignment.Structure,
                                             get_material(mat, report)))
    cs = CompoundStructure.CreateSimpleCompoundStructure(cs_layers)
    cs.StructuralMaterialIndex = 0
    ft.SetCompoundStructure(cs)
    report.append("Created floor type '{}' ({:.0f}\" thick)".format(
        name, sum(w for w, _ in layers) * 12))
    return ft


def exists_on_level(category, level, type_id=None):
    for e in FilteredElementCollector(doc).OfCategory(category).WhereElementIsNotElementType():
        lp = e.get_Parameter(BuiltInParameter.LEVEL_PARAM) or \
             e.get_Parameter(BuiltInParameter.ROOF_BASE_LEVEL_PARAM)
        if lp and lp.AsElementId() == level.Id and (type_id is None or e.GetTypeId() == type_id):
            return True
    return False


# ---------------------------------------------------------------- create ----
report = []

TransactionManager.Instance.EnsureInTransaction(doc)
try:
    for level_name, kind, tname, opening in PLAN:
        level = levels[level_name]
        z = level.ProjectElevation
        x0, y0, x1, y1 = perim

        if kind == "floor":
            layers = [(SLAB_THICKNESS, "Concrete")] if tname == SLAB_TYPE else BUBBLE_LAYERS
            ft = get_floor_type(tname, layers, report)
            if exists_on_level(BuiltInCategory.OST_Floors, level, ft.Id):
                report.append("Skipped {} floor - '{}' already there".format(level_name, tname))
                continue
            loops = List[CurveLoop]()
            loops.Add(curve_loop(rect_lines(x0, y0, x1, y1, z, ccw=True)))
            if opening:
                loops.Add(curve_loop(rect_lines(ax0, ay0, ax1, ay1, z, ccw=False)))
            fl = Floor.Create(doc, loops, ft.Id, level.Id)
            sp = fl.get_Parameter(BuiltInParameter.FLOOR_PARAM_IS_STRUCTURAL)
            if sp and not sp.IsReadOnly:
                sp.Set(1)
            report.append("{}: '{}' floor {:.0f}' x {:.0f}'{}".format(
                level_name, tname, x1 - x0, y1 - y0,
                " with {:.0f}' x {:.0f}' atrium opening".format(ax1 - ax0, ay1 - ay0) if opening else ""))

        else:  # roof
            if exists_on_level(BuiltInCategory.OST_Roofs, level):
                report.append("Skipped {} roof - a roof is already there".format(level_name))
                continue
            if ROOF_COVERS_LEVEL4_FRAMING:
                ext = level_framing_extents(level)
                if ext:
                    x0, y0 = min(x0, ext[0]), min(y0, ext[1])
                    x1, y1 = max(x1, ext[2]), max(y1, ext[3])
            rts = list(FilteredElementCollector(doc).OfClass(RoofType))
            if not rts:
                raise Exception("No roof types in the project.")
            rt = next((r for r in rts if "generic" in type_name(r).lower()), rts[0])

            footprint = CurveArray()
            for ln in rect_lines(x0, y0, x1, y1, z, ccw=True):
                footprint.Append(ln)
            if opening:
                for ln in rect_lines(ax0, ay0, ax1, ay1, z, ccw=False):
                    footprint.Append(ln)

            mapping = ModelCurveArray()
            result = doc.Create.NewFootPrintRoof(footprint, level, rt, mapping)
            roof, mapping = (result[0], result[1]) if isinstance(result, tuple) else (result, mapping)

            # Flat roof: no edge defines slope
            for mc in mapping:
                try:
                    roof.set_DefinesSlope(mc, False)
                except Exception:
                    try:
                        roof.DefinesSlope[mc] = False
                    except Exception:
                        pass
            report.append("{}: roof '{}' {:.0f}' x {:.0f}'{} (y {:.0f}' to {:.0f}')".format(
                level_name, type_name(rt), x1 - x0, y1 - y0,
                " with atrium opening" if opening else "", y0, y1))
finally:
    TransactionManager.Instance.TransactionTaskDone()

report.append("Atrium bay: x {:.0f}'-{:.0f}', y {:.0f}'-{:.0f}'".format(ax0, ax1, ay0, ay1))

OUT = report
