# -----------------------------------------------------------------------------
# Perimeter Offices, Levels 2 & 3  -  Dynamo 3.x (CPython3 / PythonNet3)
#
# Builds a ring of offices between the curtain wall and the atrium on each
# level, following the office sketch:
#   * ~15' wide offices (sketch), depth reduced to ~16'-4" so the corridor
#     around the atrium is 4'-6" clear to the atrium edge (4'-0" clear past
#     the W10 columns at the atrium corners)
#   * Door at the corridor-side corner, bookshelves on the corridor wall,
#     desk mid-room with desk chair on the window side and two guest chairs
#     on the corridor side
#   * A larger corner office in each corner, reached from the corridor
#     (the office ring is a "pinwheel" so every corner office has a door)
#   * Rooms named "Office", numbered 201.. / 301..
#
#   Plan (one level; # = atrium, . = corridor):
#
#     +-----+------+------+------+-------------+
#     | W4  |  N3  |  N2  |  N1  |  N corner   |
#     +-----+......................+-----------+
#     | W3  :                      :    E3     |
#     +-----+   ##################  +----------+
#     | W2  :   ##   ATRIUM     ##  :    E2    |
#     +-----+   ##  40' x 40'   ##  +----------+
#     | W1  :   ##################  :    E1    |
#     +-----------+......................+-----+
#     | S corner  |  S1  |  S2  |  S3  |  E4   |
#     +-----------+------+------+------+-----+-+
#       (each side: 1 corner office + 3 standard offices = 16 per level)
#
# Families:
#   Walls  : first Basic wall type with "Interior" in its name (e.g.
#            Interior - 4 7/8" Partition); change PARTITION_TYPE to pick one.
#   Doors  : Single-Flush 36" x 84" (loaded from the Revit library if needed)
#   Desk / chairs / shelving : loaded families are used first; otherwise the
#            script loads the first match from the Revit library's Furniture
#            folder. If none can be found it says which one is missing.
#
# Safe to re-run: skips walls that already exist and any office that
# already has a room.
#
# HOW TO USE
#   Dynamo 3.x > Python Script node (CPython3 or PythonNet3) > paste > Run.
#   Ctrl+Z in Revit undoes the whole run.
# -----------------------------------------------------------------------------

import math
import os
import glob

# =================================================================== settings
LEVELS = [("Level 2", "Level 3"), ("Level 3", "Level 4")]   # (office level, level above)
CORRIDOR_CLEAR = 4.5          # ft, atrium edge -> corridor wall face (4' past columns)
TARGET_OFFICE_WIDTH = 15.0    # ft, from the sketch
WALL_TOP_OFFSET = -11.0 / 12  # partitions stop under the 11" slab above
PARTITION_TYPE = None         # exact wall type name, or None = first "Interior" basic wall
DOOR_FAMILY, DOOR_TYPE = "Single-Flush", '36" x 84"'
DOOR_EDGE = 0.5               # ft, side wall face -> door frame (standard offices)
FRONT_IS_NEG_Y = True         # furniture families face their -Y side (Revit convention);
                              # set False if desks/chairs come out backwards
# Extra folders to search for furniture .rfa files (searched before the Revit
# library), e.g. [r"C:\Users\me\Documents\Revit Families"]
FAMILY_SEARCH_FOLDERS = []
# If a furniture family can't be found, skip that piece (True) or stop (False)
SKIP_MISSING_FURNITURE = True

# Furniture layout as fractions of the office (u = across, from the door side;
# v = depth, from the window). Taken from the sketch.
DESK = (0.36, 0.40)
DESK_CHAIR = (0.36, 0.24)
GUEST_CHAIRS = [(0.22, 0.58), (0.50, 0.58)]
SHELVES_U = 0.66              # bookshelves centered here, against the corridor wall
SHELF_COUNT = 2


# ============================================================ PURE LAYOUT START
def build_layout(perim, atrium, wall_t, clear=CORRIDOR_CLEAR, target_w=TARGET_OFFICE_WIDTH):
    """Pure geometry, no Revit. perim/atrium = (x0, y0, x1, y1) in feet.

    Returns (bands, walls, offices) in plan coordinates:
      walls   : [(band_index, kind, (x, y), (x, y))]  kind = 'corridor' | 'demising'
      offices : [dict(band, corner, ul, ur, D, door_u, origin, U, V)]
    """
    x0, y0, x1, y1 = perim
    ax0, ay0, ax1, ay1 = atrium
    # Each band: origin corner, U (along band), V (toward atrium),
    # band length, atrium start along U, atrium distance along V.
    bands = [
        dict(name="S", O=(x0, y0), U=(1, 0), V=(0, 1), L=x1 - x0, au=ax0 - x0, av=ay0 - y0),
        dict(name="E", O=(x1, y0), U=(0, 1), V=(-1, 0), L=y1 - y0, au=ay0 - y0, av=x1 - ax1),
        dict(name="N", O=(x1, y1), U=(-1, 0), V=(0, -1), L=x1 - x0, au=x1 - ax1, av=y1 - ay1),
        dict(name="W", O=(x0, y1), U=(0, -1), V=(1, 0), L=y1 - y0, au=y1 - ay1, av=ax0 - x0),
    ]
    for b in bands:
        b["D"] = b["av"] - clear - wall_t / 2.0      # corridor wall centerline depth

    def pt(b, u, v):
        return (b["O"][0] + u * b["U"][0] + v * b["V"][0],
                b["O"][1] + u * b["U"][1] + v * b["V"][1])

    walls, offices = [], []
    n = len(bands)
    for i, b in enumerate(bands):
        nxt, prv = bands[(i + 1) % n], bands[(i - 1) % n]
        D = b["D"]
        u_end = b["L"] - nxt["D"]          # where the next band's corridor wall line is
        b["u_end"] = u_end
        # Corridor wall: from the curtain wall (u=0) to the next band's corridor line.
        # Its first stretch (u < prv D) is the end wall of the previous band.
        walls.append((i, "corridor", pt(b, 0, D), pt(b, u_end, D)))

        # Corner office: u 0 .. atrium start; corridor frontage = prv["D"] .. au
        u1 = b["au"]
        frontage = (u1 - wall_t / 2.0) - (prv["D"] + wall_t / 2.0)
        offices.append(dict(band=i, corner=True, ul=0.0, ur=u1, D=D,
                            door_u=(prv["D"] + u1) / 2.0, frontage=frontage))
        # Standard offices: u1 .. u_end
        count = max(1, int(round((u_end - u1) / target_w)))
        w = (u_end - u1) / count
        edges = [u1 + k * w for k in range(count + 1)]
        for k in range(count):
            offices.append(dict(band=i, corner=False, ul=edges[k], ur=edges[k + 1], D=D,
                                door_u=None, frontage=None))
        # Demising walls (curtain wall -> corridor wall), excluding u_end
        # (that line is the next band's corridor wall).
        for u in edges[:-1]:
            walls.append((i, "demising", pt(b, u, 0), pt(b, u, D)))

    for o in offices:
        b = bands[o["band"]]
        o.update(origin=b["O"], U=b["U"], V=b["V"])
    return bands, walls, offices


def office_point(o, fu, fv, mirror):
    """Plan point at fraction fu across / fv deep in an office."""
    w = o["ur"] - o["ul"]
    u = o["ul"] + ((1.0 - fu) if mirror else fu) * w
    v = fv * o["D"]
    return (o["origin"][0] + u * o["U"][0] + v * o["V"][0],
            o["origin"][1] + u * o["U"][1] + v * o["V"][1])
# ============================================================== PURE LAYOUT END


if __name__ != "__layout_test__":
    import clr
    clr.AddReference("RevitAPI")
    clr.AddReference("RevitServices")

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

    doc = DocumentManager.Instance.CurrentDBDocument
    TOL = 0.01
    report = []

    def type_name(e):
        p = e.get_Parameter(BuiltInParameter.SYMBOL_NAME_PARAM)
        return p.AsString() if p else e.Name

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

    # ------------------------------------------- perimeter + atrium (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)
            (xs if abs(p1.Y - p0.Y) > abs(p1.X - p0.X) else ys).append(p0.X if abs(p1.Y - p0.Y) > abs(p1.X - p0.X) else 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

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

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

    ax0, ax1 = bay(xs, (xs[0] + xs[-1]) / 2.0)
    ay0, ay1 = bay(ys, (ys[0] + ys[-1]) / 2.0)

    # Perimeter = curtain wall location lines (fallback: outer grids + 1')
    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:
        perim = (min(p.X for p in cw_pts), min(p.Y for p in cw_pts),
                 max(p.X for p in cw_pts), max(p.Y for p in cw_pts))
    else:
        perim = (xs[0] - 1.0, ys[0] - 1.0, xs[-1] + 1.0, ys[-1] + 1.0)
        report.append("No curtain walls found - using outer grids + 1' as the perimeter")

    # ------------------------------------------------------- wall type ------
    basic = [wt for wt in FilteredElementCollector(doc).OfClass(WallType) if wt.Kind == WallKind.Basic]
    if PARTITION_TYPE:
        wall_type = next((wt for wt in basic if type_name(wt) == PARTITION_TYPE), None)
    else:
        wall_type = next((wt for wt in basic if "interior" in type_name(wt).lower()), None) \
            or min(basic, key=lambda wt: wt.Width)
    if wall_type is None:
        raise Exception("Partition wall type not found.")
    t = wall_type.Width

    bands, walls, offices = build_layout(perim, (ax0, ay0, ax1, ay1), t)

    # --------------------------------------------------- family loading ----
    def library_roots():
        roots = glob.glob(os.path.join(r"C:\ProgramData\Autodesk", "RVT *", "Libraries", "*"))
        return sorted(roots, reverse=True)

    def find_library_file(subfolder, predicate):
        bases = list(FAMILY_SEARCH_FOLDERS)
        for root in library_roots():
            bases += [os.path.join(root, subfolder), root]   # preferred folder, then whole library
        for base in bases:
            if not os.path.isdir(base):
                continue
            hits = []
            for dirpath, _, files in os.walk(base):
                for f in files:
                    if f.lower().endswith(".rfa") and predicate(f.lower()):
                        hits.append(os.path.join(dirpath, f))
            if hits:
                return sorted(hits, key=len)[0]
        return None

    def loaded_symbol(categories, predicate):
        for cat in categories:
            for s in FilteredElementCollector(doc).OfCategory(cat).OfClass(FamilySymbol):
                if predicate(s.Family.Name.lower()):
                    return s
        return None

    def get_furniture(label, predicate, preferred=None):
        try:
            return _get_furniture(label, predicate, preferred)
        except Exception as ex:
            if not SKIP_MISSING_FURNITURE:
                raise
            warnings.append("Skipped {}s: {}".format(label, ex))
            return None

    def _get_furniture(label, predicate, preferred=None):
        cats = [BuiltInCategory.OST_Furniture, BuiltInCategory.OST_Casework,
                BuiltInCategory.OST_FurnitureSystems, BuiltInCategory.OST_SpecialityEquipment,
                BuiltInCategory.OST_GenericModel]
        if preferred:
            s = loaded_symbol(cats, preferred)
            if s:
                return s
        s = loaded_symbol(cats, predicate)
        if s:
            return s
        for pred in ([preferred] if preferred else []) + [predicate]:
            path = find_library_file("Furniture", pred)
            if path and doc.LoadFamily(path):
                report.append("Loaded {} family: {}".format(label, os.path.basename(path)))
                s = loaded_symbol(cats, lambda n, stem=os.path.splitext(os.path.basename(path))[0].lower(): n == stem)
                if s:
                    return s
        raise Exception("no {} family loaded or found in the Revit library - load one via "
                        "Insert > Load Family (or add its folder to FAMILY_SEARCH_FOLDERS) and re-run"
                        .format(label))

    def get_door():
        def match(s):
            return s.Family.Name.lower() == DOOR_FAMILY.lower() and type_name(s).replace(" ", "") == DOOR_TYPE.replace(" ", "")
        doors = list(FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_Doors).OfClass(FamilySymbol))
        s = next((d for d in doors if match(d)), None)
        if s:
            return s
        path = find_library_file("Doors", lambda f: f == DOOR_FAMILY.lower() + ".rfa")
        if path:
            try:
                doc.LoadFamilySymbol(path, DOOR_TYPE)
            except Exception:
                pass
            doors = list(FilteredElementCollector(doc).OfCategory(BuiltInCategory.OST_Doors).OfClass(FamilySymbol))
            s = next((d for d in doors if match(d)), None)
            if s:
                report.append("Loaded door {} : {}".format(DOOR_FAMILY, DOOR_TYPE))
                return s
        s = next((d for d in doors if "single" in d.Family.Name.lower()), None) or (doors[0] if doors else None)
        if s is None:
            raise Exception("No door families loaded. Load Single-Flush via Insert > Load Family and re-run.")
        report.append("Using door {} : {}".format(s.Family.Name, type_name(s)))
        return s

    # ------------------------------------------------------- placement ------
    def xyz(p, z):
        return XYZ(p[0], p[1], z)

    def vec(v):
        return XYZ(v[0], v[1], 0)

    def plan_extent(inst, direction):
        bb = inst.get_BoundingBox(None)
        vals = [XYZ(x, y, 0).DotProduct(direction) for x in (bb.Min.X, bb.Max.X) for y in (bb.Min.Y, bb.Max.Y)]
        return min(vals), max(vals)

    def bbox_center(inst):
        bb = inst.get_BoundingBox(None)
        return XYZ((bb.Min.X + bb.Max.X) / 2.0, (bb.Min.Y + bb.Max.Y) / 2.0, 0)

    def place(symbol, level, p, front):
        """Place a level-based family centered (by bounding box) at p, front facing `front`."""
        z = level.ProjectElevation
        inst = doc.Create.NewFamilyInstance(xyz(p, z), symbol, level, StructuralType.NonStructural)
        doc.Regenerate()
        by = inst.GetTransform().BasisY
        target = front.Negate() if FRONT_IS_NEG_Y else front
        ang = math.atan2(by.X * target.Y - by.Y * target.X, by.X * target.X + by.Y * target.Y)
        if abs(ang) > 1e-6:
            o = inst.Location.Point
            ElementTransformUtils.RotateElement(doc, inst.Id, Line.CreateBound(o, o + XYZ(0, 0, 1)), ang)
            doc.Regenerate()
        c = bbox_center(inst)
        ElementTransformUtils.MoveElement(doc, inst.Id, XYZ(p[0] - c.X, p[1] - c.Y, 0))
        return inst

    existing_walls = []
    for w in FilteredElementCollector(doc).OfClass(Wall):
        if isinstance(w.Location, LocationCurve):
            c = w.Location.Curve
            existing_walls.append((w, c.GetEndPoint(0), c.GetEndPoint(1)))

    def find_wall(level, a, b):
        for w, p, q in existing_walls:
            if w.LevelId != level.Id:
                continue
            if (abs(p.X - a[0]) < TOL and abs(p.Y - a[1]) < TOL and abs(q.X - b[0]) < TOL and abs(q.Y - b[1]) < TOL) or \
               (abs(p.X - b[0]) < TOL and abs(p.Y - b[1]) < TOL and abs(q.X - a[0]) < TOL and abs(q.Y - a[1]) < TOL):
                return w
        return None

    # ---------------------------------------------------------- build -------
    counts = dict(walls=0, doors=0, rooms=0, furniture=0, skipped=0)
    warnings = []
    is_chair = lambda n: any(k in n for k in ("chair", "seat", "stool"))

    TransactionManager.Instance.EnsureInTransaction(doc)
    try:
        door_sym = get_door()
        desk_sym = get_furniture("desk", lambda n: "desk" in n and "chair" not in n,
                                 preferred=lambda n: n == "desk")
        task_sym = get_furniture("desk chair", is_chair,
                                 preferred=lambda n: is_chair(n) and ("executive" in n or "task" in n))
        guest_sym = get_furniture("guest chair", is_chair,
                                  preferred=lambda n: is_chair(n) and ("breuer" in n or "guest" in n or "side" in n)) \
            or task_sym
        task_sym = task_sym or guest_sym
        shelf_sym = get_furniture("bookshelf", lambda n: "shelv" in n or "book" in n,
                                  preferred=lambda n: "book" in n)
        for s in (door_sym, desk_sym, task_sym, guest_sym, shelf_sym):
            if s is not None and not s.IsActive:
                s.Activate()
        doc.Regenerate()

        for o in offices:
            if o["corner"] and o["frontage"] < 3.5:
                warnings.append("Corner office on {} side has only {:.1f}' of corridor frontage for its door"
                                .format(bands[o["band"]]["name"], o["frontage"]))

        for li, (level_name, above_name) in enumerate(LEVELS):
            level, above = levels[level_name], levels[above_name]
            z = level.ProjectElevation
            height = above.Elevation - level.Elevation + WALL_TOP_OFFSET

            # Walls
            corridor_walls = {}
            for band_i, kind, a, b in walls:
                w = find_wall(level, a, b)
                if w is None:
                    w = Wall.Create(doc, Line.CreateBound(xyz(a, z), xyz(b, z)), wall_type.Id,
                                    level.Id, height, 0.0, False, False)
                    w.get_Parameter(BuiltInParameter.WALL_HEIGHT_TYPE).Set(above.Id)
                    w.get_Parameter(BuiltInParameter.WALL_TOP_OFFSET).Set(WALL_TOP_OFFSET)
                    existing_walls.append((w, xyz(a, z), xyz(b, z)))
                    counts["walls"] += 1
                if kind == "corridor":
                    corridor_walls[band_i] = w
            doc.Regenerate()

            # Offices
            number = int(''.join(ch for ch in level_name if ch.isdigit()) or (li + 2)) * 100
            for o in offices:
                number += 1
                V, U = vec(o["V"]), vec(o["U"])
                mirror = o["corner"]               # corner office: door on the far side
                center = office_point(o, 0.5, 0.5, False)
                if doc.GetRoomAtPoint(XYZ(center[0], center[1], z + 1.0)) is not None:
                    counts["skipped"] += 1
                    continue

                # Door in the corridor wall, swinging into the office
                if o["corner"]:
                    du = o["door_u"]
                else:
                    dw = door_sym.get_Parameter(BuiltInParameter.DOOR_WIDTH)
                    dwidth = dw.AsDouble() if dw else 3.0
                    du = o["ul"] + t / 2.0 + DOOR_EDGE + dwidth / 2.0
                dp = (o["origin"][0] + du * o["U"][0] + o["D"] * o["V"][0],
                      o["origin"][1] + du * o["U"][1] + o["D"] * o["V"][1])
                door = doc.Create.NewFamilyInstance(xyz(dp, z), door_sym, corridor_walls[o["band"]],
                                                    level, StructuralType.NonStructural)
                doc.Regenerate()
                # Facing points to the side the door swings away from -> corridor
                if door.FacingOrientation.DotProduct(V) < 0:
                    door.flipFacing()
                counts["doors"] += 1

                # Desk, desk chair (window side), guest chairs (corridor side)
                if desk_sym:
                    place(desk_sym, level, office_point(o, DESK[0], DESK[1], mirror), V.Negate())
                    counts["furniture"] += 1
                if task_sym:
                    place(task_sym, level, office_point(o, DESK_CHAIR[0], DESK_CHAIR[1], mirror), V)
                    counts["furniture"] += 1
                if guest_sym:
                    for fu, fv in GUEST_CHAIRS:
                        place(guest_sym, level, office_point(o, fu, fv, mirror), V.Negate())
                        counts["furniture"] += 1

                # Bookshelves against the corridor wall, side by side
                anchor = office_point(o, SHELVES_U, 0.5, mirror)
                if shelf_sym is None:
                    shelf_count = 0
                else:
                    shelf_count = SHELF_COUNT
                if shelf_count:
                    first = place(shelf_sym, level, anchor, V.Negate())
                    lo, hi = plan_extent(first, U)
                    width = hi - lo
                    shelves = [first]
                    for k in range(1, shelf_count):
                        shelves.append(place(shelf_sym, level, anchor, V.Negate()))
                    doc.Regenerate()
                    wall_face = o["D"] - t / 2.0 - 0.02
                    origin = XYZ(o["origin"][0], o["origin"][1], 0)
                    for k, s in enumerate(shelves):
                        du_shift = (k - (shelf_count - 1) / 2.0) * width
                        vmin, vmax = plan_extent(s, V)
                        dv = wall_face - (vmax - origin.DotProduct(V))
                        ElementTransformUtils.MoveElement(doc, s.Id, U.Multiply(du_shift) + V.Multiply(dv))
                    counts["furniture"] += shelf_count

                # Room
                room = doc.Create.NewRoom(level, UV(center[0], center[1]))
                room.Name = "Office"
                room.Number = str(number)
                counts["rooms"] += 1
    finally:
        TransactionManager.Instance.TransactionTaskDone()

    std = [o for o in offices if not o["corner"]]
    cor = [o for o in offices if o["corner"]]
    report.insert(0, "Per level: {} offices ({} corner ~{:.0f}' x {:.1f}', {} standard ~{:.1f}' x {:.1f}')".format(
        len(offices), len(cor), cor[0]["ur"] - cor[0]["ul"], cor[0]["D"],
        len(std), std[0]["ur"] - std[0]["ul"], std[0]["D"]))
    report.append("Created {} walls, {} doors, {} furniture items, {} rooms on {}".format(
        counts["walls"], counts["doors"], counts["furniture"], counts["rooms"],
        " & ".join(l for l, _ in LEVELS)))
    report.append("Corridor: {:.1f}' clear from atrium edge to wall face".format(CORRIDOR_CLEAR))
    if counts["skipped"]:
        report.append("Skipped {} offices that already have rooms".format(counts["skipped"]))
    report.extend(warnings)
    OUT = report
