# ============================================================
# Blender Script v3: pMHC-II Three-Fate Tolerance Animation (Cycles)
# (Tolera — antigen-specific tolerance platform)
# ============================================================
# WHAT CHANGED FROM v2:
#   v2 showed ONE T cell being anergized. v3 shows Tolera's full
#   three-fold tolerance strategy as three T cells from the same
#   antigen-specific clone, each engaging the SAME wheat pMHC-II and
#   diverging into a different tolerogenic fate:
#       LEFT   -> ANERGY            (survives, goes unresponsive/dormant)
#       CENTER -> TOLERANCE (iTreg) (converted to a regulatory phenotype)
#       RIGHT  -> CLONAL DELETION   (apoptosis: blebs, fades out)
#   Using one antigen (wheat) as the worked example is intentional —
#   the mechanism scales identically across peptides/allergens, so a
#   single example reads clearly without cluttering the frame.
#
# USAGE:
#   1. Blender -> Scripting tab -> paste -> Run Script
#   2. Render Properties -> Engine = Cycles (script sets this); set
#      Device > GPU Compute if you have a GPU.
#   3. Scrub frames 1-420 (~14 s @ 30 fps) or Spacebar to preview.
#      Use Material Preview shading while iterating; Rendered only for
#      final look checks.
#
# STORY BEATS (420 frames @ 30 fps = 14 s):
#   1-60     Establishing: three identical ACTIVATED T cells (green,
#            receptors glowing), camera wide 3/4 push-in
#   60-120   Three wheat pMHC-II complexes fly in, one per cell
#   120-150  Engagement at each interface
#   150-320  DIVERGENT FATES play out in parallel:
#              anergy   -> receptor dim wave + shrink to dormant gray
#              iTreg    -> recolor to regulatory blue + suppression halo
#              deletion -> blebbing + darken + fade to nothing
#   220-330  Fate labels fade in under each cell
#   330-420  Camera settles to wide 3/4 elevated hold (ChimeraX hand-off)
# ============================================================

import bpy
from bpy_extras import anim_utils
import math
import random
import mathutils

# ------------------------------------------------------------
# CONFIG
# ------------------------------------------------------------
CELL_RADIUS = 1.5
N_RECEPTORS = 20
CELL_SPACING = 4.2            # X distance between the three cells

ACTIVE_COLOR   = (0.08, 0.85, 0.55, 1.0)   # activated T cell (green)
ANERGIC_COLOR  = (0.30, 0.32, 0.38, 1.0)   # anergy (desaturated gray)
TREG_COLOR     = (0.15, 0.55, 0.85, 1.0)   # regulatory / iTreg (blue)
DELETION_COLOR = (0.35, 0.10, 0.12, 1.0)   # apoptotic (dark red)
PEPTIDE_COLOR  = (0.95, 0.50, 0.10, 1.0)   # wheat pMHC cargo (orange)

SHRINK_FACTOR = 0.72          # anergy dormant size
CASCADE_SPREAD_FRAMES = 60    # anergy receptor-dim wave duration

FATE_START = 150              # frame all three fates begin resolving

# ------------------------------------------------------------
# RENDER SETTINGS (Cycles)
# ------------------------------------------------------------
scene = bpy.context.scene
scene.render.engine = 'CYCLES'
scene.cycles.samples = 256
scene.cycles.use_denoising = True
scene.cycles.use_adaptive_sampling = True
scene.cycles.adaptive_threshold = 0.01
scene.cycles.max_bounces = 8
scene.render.resolution_x = 1920
scene.render.resolution_y = 1080
scene.render.fps = 30
scene.frame_start = 1
scene.frame_end = 420
scene.render.film_transparent = False
try:
    scene.cycles.device = 'GPU'
except Exception:
    pass

# ------------------------------------------------------------
# CLEAN SCENE
# ------------------------------------------------------------
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)
for block in list(bpy.data.materials):
    bpy.data.materials.remove(block)
for block in list(bpy.data.textures):
    bpy.data.textures.remove(block)

# ------------------------------------------------------------
# KEYFRAME HELPERS
# ------------------------------------------------------------
def kf_loc(obj, frame, loc):
    obj.location = loc
    obj.keyframe_insert(data_path="location", frame=frame)

def kf_scale(obj, frame, scale):
    if isinstance(scale, (int, float)):
        scale = (scale, scale, scale)
    obj.scale = scale
    obj.keyframe_insert(data_path="scale", frame=frame)

def kf_color(node_input, frame, color):
    node_input.default_value = color
    node_input.keyframe_insert(data_path="default_value", frame=frame)

def kf_value(node_input, frame, value):
    node_input.default_value = value
    node_input.keyframe_insert(data_path="default_value", frame=frame)

def kf_mod(mod, prop, frame, value):
    setattr(mod, prop, value)
    mod.keyframe_insert(data_path=prop, frame=frame)

def _iter_action_fcurves(id_data):
    """Yield every F-Curve animating `id_data`, across Blender versions.

    Blender <=4.3 stores F-Curves directly on the Action (action.fcurves).
    Blender 4.4+ / 5.x use slotted Actions: F-Curves live on a channelbag
    reached via layer -> strip -> slot. On 5.0+ the old action.fcurves proxy
    was removed entirely, so we must go through the channelbag."""
    ad = id_data.animation_data
    action = ad.action

    # Legacy path (Blender <= 4.3)
    if hasattr(action, "fcurves"):
        yield from action.fcurves
        return

    # Slotted-Action path (Blender 4.4+ / 5.x)
    slot = ad.action_slot
    if slot is None:
        return
    channelbag = anim_utils.action_get_channelbag_for_slot(action, slot)
    if channelbag is not None:
        yield from channelbag.fcurves

def ease_all_fcurves(id_data):
    if id_data and id_data.animation_data and id_data.animation_data.action:
        for fc in _iter_action_fcurves(id_data):
            for kp in fc.keyframe_points:
                kp.interpolation = 'BEZIER'
                kp.easing = 'EASE_IN_OUT'

# ------------------------------------------------------------
# MOLECULAR NODES: apply a DIFFERENT style per chain
# ------------------------------------------------------------
# The MN import popup only lets you pick ONE style for the whole
# structure. To get cartoon on chains A/B and sticks/spacefill on the
# peptide (chain P), you build the two styles in PARALLEL off the same
# atom geometry -- each gated by a chain selection -- and merge them with
# a Join Geometry node before Group Output. (You cannot stack one style
# after another: a style outputs mesh geometry, not atom-like geometry.)
#
# This function does that wiring for you. Import wheat_pmhc.pdb ONCE with
# any starting style, then call:
#
#   style_by_chain(bpy.data.objects["wheat_pmhc"],
#                  {"Style Cartoon":        ["A", "B"],
#                   "Style Ball and Stick": ["P"]})
#
# For spacefill on the peptide use "Style Spheres" instead of
# "Style Ball and Stick". Every style you name must already exist as a
# node group in the file -- MN loads a style's node group the first time
# it is used, so if a style is missing, add it once via the Geometry
# Nodes editor (Shift+A -> Molecular Nodes -> Style -> ...) and re-run.
# The function tells you exactly which one to add if it can't find it.

def _find_node_group(name):
    """Return an MN node group by tolerant name match, or None.
    Matches 'Style Ball and Stick' against 'Style Ball and Stick',
    'MN_style_ball_and_stick', etc."""
    def norm(s):
        return "".join(c for c in s.lower() if c.isalnum())
    target = norm(name)
    # exact-ish first
    for ng in bpy.data.node_groups:
        if norm(ng.name) == target:
            return ng
    # then containment (handles 'MN_style_' prefixes / suffixes)
    for ng in bpy.data.node_groups:
        if target in norm(ng.name) or norm(ng.name).endswith(target):
            return ng
    return None

def style_by_chain(mol_obj, chain_styles):
    """Wire a per-chain multi-style Molecular Nodes graph.

    mol_obj      : the object MN created on import.
    chain_styles : {style_node_group_name: [chain_id, ...], ...}
                   e.g. {"Style Cartoon": ["A","B"],
                         "Style Ball and Stick": ["P"]}
    Returns the modified GeometryNodes node tree.
    Raises RuntimeError with an actionable message if a required MN node
    group is missing (so nothing breaks silently)."""
    # 1. locate the GeometryNodes modifier + its node tree
    mod = next((m for m in mol_obj.modifiers if m.type == 'NODES'), None)
    if mod is None or mod.node_group is None:
        raise RuntimeError(
            f"'{mol_obj.name}' has no Geometry Nodes modifier -- is this "
            "really the object Molecular Nodes created on import?")
    tree = mod.node_group
    nodes, links = tree.nodes, tree.links

    def geo_out(node):
        return next((s for s in node.outputs if s.type == 'GEOMETRY'), None)
    def geo_in(node):
        return next((s for s in node.inputs if s.type == 'GEOMETRY'), None)

    grp_in  = next(n for n in nodes if n.type == 'GROUP_INPUT')
    grp_out = next(n for n in nodes if n.type == 'GROUP_OUTPUT')

    # 2. find the current Style node and the socket feeding its geometry
    #    input -- that socket carries the colored atom geometry we want to
    #    branch from (so per-chain colors are preserved).
    existing_style = next(
        (n for n in nodes
         if n.type == 'GROUP' and n.node_group
         and n.node_group.name.lower().startswith(("style", "mn_style"))),
        None)
    if existing_style is not None and geo_in(existing_style) and geo_in(existing_style).links:
        source_socket = geo_in(existing_style).links[0].from_socket
    else:
        # fall back to whatever currently feeds Group Output, else raw atoms
        go_in = geo_in(grp_out)
        source_socket = go_in.links[0].from_socket if (go_in and go_in.links) else geo_out(grp_in)

    # 3. locate the Select Chain node group (added once via the GUI, or
    #    auto-generated per molecule as 'Select Chain' / 'Selection ... Chains')
    chain_ng = None
    for ng in bpy.data.node_groups:
        nm = ng.name.lower()
        if "chain" in nm and ("select" in nm or "selection" in nm):
            chain_ng = ng
            break
    if chain_ng is None:
        raise RuntimeError(
            "No 'Select Chain' node group found. In the Geometry Nodes "
            "editor add one once: Shift+A -> Molecular Nodes -> Selection "
            "-> Select Chain, then re-run style_by_chain().")

    # 4. remove the old single-style node so we don't double-draw chains
    if existing_style is not None:
        nodes.remove(existing_style)

    # 5. one Join Geometry to merge all branches
    join = nodes.new("GeometryNodeJoinGeometry")
    join.location = (600, 0)
    links.new(join.outputs["Geometry"], geo_in(grp_out))

    def set_chain_selection(sel_node, chain_ids):
        """Tick the boolean inputs on a Select Chain node that match the
        requested chain letters; untick the rest."""
        wanted = {c.upper() for c in chain_ids}
        matched = []
        for sock in sel_node.inputs:
            if sock.type != 'BOOLEAN':
                continue
            on = sock.name.strip().upper() in wanted
            try:
                sock.default_value = on
            except Exception:
                pass
            if on:
                matched.append(sock.name)
        return matched

    # 6. build one branch per style
    y = 300
    for style_name, chain_ids in chain_styles.items():
        style_ng = _find_node_group(style_name)
        if style_ng is None:
            nodes.remove(join)
            raise RuntimeError(
                f"Style node group '{style_name}' not found in this file. "
                f"Add it once via Shift+A -> Molecular Nodes -> Style, "
                f"then re-run. (Available node groups: "
                f"{[g.name for g in bpy.data.node_groups]})")

        # style node
        style_node = nodes.new("GeometryNodeGroup")
        style_node.node_group = style_ng
        style_node.location = (350, y)
        links.new(source_socket, geo_in(style_node))

        # per-branch chain selection (its own copy so toggles don't clash)
        sel_node = nodes.new("GeometryNodeGroup")
        sel_node.node_group = chain_ng.copy()
        sel_node.location = (120, y)
        matched = set_chain_selection(sel_node, chain_ids)
        if not matched:
            print(f"[style_by_chain] WARNING: none of {chain_ids} matched a "
                  f"chain toggle on the Select Chain node. Its toggles are: "
                  f"{[s.name for s in sel_node.inputs if s.type=='BOOLEAN']}")

        # feed selection into the style node's Selection input if present
        sel_in = next((s for s in style_node.inputs
                       if s.name.lower() == "selection"), None)
        sel_out = next((s for s in sel_node.outputs
                        if s.type == 'BOOLEAN'), None)
        if sel_in and sel_out:
            links.new(sel_out, sel_in)

        links.new(geo_out(style_node), join.inputs[0])
        y -= 300

    print(f"[style_by_chain] Applied {list(chain_styles)} to "
          f"'{mol_obj.name}'. Check the Geometry Nodes editor if a chain "
          f"looks wrong -- confirm the chain letters match the Select "
          f"Chain toggles.")
    return tree

# ------------------------------------------------------------
# BUILD ONE T CELL (body + SSS membrane + emissive receptors)
# Returns a dict of handles so each cell can be animated independently.
# ------------------------------------------------------------
def build_tcell(name, center, seed):
    cx, cy, cz = center

    # --- body with SSS membrane ---
    bpy.ops.mesh.primitive_uv_sphere_add(radius=CELL_RADIUS, location=center,
                                          segments=96, ring_count=48)
    cell = bpy.context.active_object
    cell.name = f"{name}_Body"
    sub = cell.modifiers.new(name="Subdiv", type='SUBSURF')
    sub.levels = 2
    sub.render_levels = 3

    mat_cell = bpy.data.materials.new(name=f"{name}_Membrane")
    mat_cell.use_nodes = True
    nodes = mat_cell.node_tree.nodes
    links = mat_cell.node_tree.links
    nodes.clear()
    output = nodes.new("ShaderNodeOutputMaterial")
    bsdf = nodes.new("ShaderNodeBsdfPrincipled")
    bsdf.inputs["Base Color"].default_value = ACTIVE_COLOR
    bsdf.inputs["Roughness"].default_value = 0.25
    sss_key = "Subsurface Weight" if "Subsurface Weight" in bsdf.inputs else "Subsurface"
    bsdf.inputs[sss_key].default_value = 0.6
    if "Subsurface Radius" in bsdf.inputs:
        bsdf.inputs["Subsurface Radius"].default_value = (0.4, 0.15, 0.1)
    links.new(bsdf.outputs["BSDF"], output.inputs["Surface"])
    cell.data.materials.append(mat_cell)

    # --- receptors (each its own material for individual dimming) ---
    receptors, receptor_mats, emit_nodes = [], [], []
    random.seed(seed)
    for i in range(N_RECEPTORS):
        theta = math.acos(1 - 2 * (i + 0.5) / N_RECEPTORS)
        phi = math.pi * (1 + 5 ** 0.5) * i
        lx = CELL_RADIUS * math.sin(theta) * math.cos(phi)
        ly = CELL_RADIUS * math.sin(theta) * math.sin(phi)
        lz = CELL_RADIUS * math.cos(theta)

        bpy.ops.mesh.primitive_cone_add(radius1=0.09, depth=0.3,
                                        location=(cx + lx, cy + ly, cz + lz))
        rec = bpy.context.active_object
        rec.name = f"{name}_Receptor_{i:02d}"
        rec.rotation_mode = 'QUATERNION'
        up = mathutils.Vector((0, 0, 1))
        target = mathutils.Vector((lx, ly, lz)).normalized()
        rec.rotation_quaternion = up.rotation_difference(target)

        mat = bpy.data.materials.new(name=f"{name}_Rec_Mat_{i:02d}")
        mat.use_nodes = True
        rn = mat.node_tree.nodes
        rl = mat.node_tree.links
        rn.clear()
        r_out = rn.new("ShaderNodeOutputMaterial")
        r_bsdf = rn.new("ShaderNodeBsdfPrincipled")
        r_emit = rn.new("ShaderNodeEmission")
        r_mix = rn.new("ShaderNodeMixShader")
        r_bsdf.inputs["Base Color"].default_value = (0.8, 0.15, 0.25, 1.0)
        r_emit.inputs["Color"].default_value = ACTIVE_COLOR
        r_emit.inputs["Strength"].default_value = 2.5
        r_mix.inputs["Fac"].default_value = 0.5
        rl.new(r_bsdf.outputs["BSDF"], r_mix.inputs[1])
        rl.new(r_emit.outputs["Emission"], r_mix.inputs[2])
        rl.new(r_mix.outputs["Shader"], r_out.inputs["Surface"])
        rec.data.materials.append(mat)

        # Parent each receptor to the cell body so it follows the body's
        # scale/location animation (receptors sit ON the membrane, so when
        # the cell grows/shrinks they must move with it). matrix_parent_inverse
        # is set to the body's current (rest, scale=1) world matrix so the
        # receptor keeps its present world position -- equivalent to Ctrl+P
        # "Object (Keep Transform)".
        rec.parent = cell
        rec.matrix_parent_inverse = cell.matrix_world.inverted()

        receptors.append(rec)
        receptor_mats.append(mat)
        emit_nodes.append(r_emit)

    # docking receptor = the one facing the camera (most -Y), so the
    # engaging pMHC lands on a visible face and the anergy dim-wave
    # radiates from a point the viewer can see.
    dock_idx = min(range(N_RECEPTORS), key=lambda i: receptors[i].location.y)

    return {
        "name": name, "center": mathutils.Vector(center),
        "cell": cell, "mat_cell": mat_cell, "mem_bsdf": bsdf,
        "receptors": receptors, "receptor_mats": receptor_mats,
        "emit_nodes": emit_nodes, "dock_idx": dock_idx,
    }

# ------------------------------------------------------------
# BUILD: wheat pMHC-II "cargo" that flies to a cell
#   PLACEHOLDER cylinder — swap for wheat_pmhc.pdb via Molecular Nodes
#   (see NOTES). All three flyers are the SAME wheat complex.
# ------------------------------------------------------------
def build_pmhc(name, cell):
    dock_rec = cell["receptors"][cell["dock_idx"]]
    dock_loc = dock_rec.location.copy()
    # start: out in front of the cell (toward camera, -Y) and up
    start_loc = (cell["center"].x, cell["center"].y - 10, cell["center"].z + 5)

    bpy.ops.mesh.primitive_cylinder_add(radius=0.18, depth=0.5, location=start_loc)
    pmhc = bpy.context.active_object
    pmhc.name = f"{name}_wheat_pMHC_PLACEHOLDER"  # -> wheat_pmhc.pdb
    mat = bpy.data.materials.new(name=f"{name}_pMHC_Mat")
    mat.use_nodes = True
    mat.node_tree.nodes["Principled BSDF"].inputs["Base Color"].default_value = PEPTIDE_COLOR
    pmhc.data.materials.append(mat)
    return pmhc, start_loc, dock_loc

# ------------------------------------------------------------
# FATE ANIMATIONS
# ------------------------------------------------------------
def animate_engage_and_hold(cell):
    """Subtle idle breathing before the fate resolves (all cells)."""
    c = cell["cell"]
    kf_scale(c, 1, 1.0)
    kf_scale(c, 25, 1.02)
    kf_scale(c, 50, 1.0)
    kf_scale(c, FATE_START, 1.0)

def animate_anergy(cell):
    """Receptor dim-wave radiating from the docking site, then shrink to
    a dormant gray cell that survives (functional unresponsiveness)."""
    dock = cell["receptors"][cell["dock_idx"]]

    def ang(r):
        v1 = (mathutils.Vector(r.location) - cell["center"]).normalized()
        v2 = (mathutils.Vector(dock.location) - cell["center"]).normalized()
        return v1.angle(v2)

    ordered = sorted(range(N_RECEPTORS), key=lambda i: ang(cell["receptors"][i]))
    for rank, idx in enumerate(ordered):
        emit = cell["emit_nodes"][idx]
        delay = (rank / max(1, N_RECEPTORS - 1)) * CASCADE_SPREAD_FRAMES
        f0 = int(FATE_START + delay)
        f1 = f0 + 18
        kf_color(emit.inputs["Color"], f0, ACTIVE_COLOR)
        kf_value(emit.inputs["Strength"], f0, 2.5)
        kf_color(emit.inputs["Color"], f1, ANERGIC_COLOR)
        kf_value(emit.inputs["Strength"], f1, 0.15)

    base = cell["mem_bsdf"].inputs["Base Color"]
    kf_color(base, FATE_START, ACTIVE_COLOR)
    kf_color(base, FATE_START + 90, ANERGIC_COLOR)

    c = cell["cell"]
    kf_scale(c, FATE_START + 90, 1.0)
    kf_scale(c, FATE_START + 150, SHRINK_FACTOR)   # ~frame 300
    kf_scale(c, 420, SHRINK_FACTOR)

def animate_treg(cell):
    """Convert to a regulatory phenotype: recolor blue, keep the cell
    alive/healthy (gentle breathing continues) and emit a soft
    'suppression field' halo — active/dominant tolerance, not shutdown."""
    base = cell["mem_bsdf"].inputs["Base Color"]
    kf_color(base, FATE_START, ACTIVE_COLOR)
    kf_color(base, FATE_START + 80, TREG_COLOR)

    for emit in cell["emit_nodes"]:
        kf_color(emit.inputs["Color"], FATE_START, ACTIVE_COLOR)
        kf_color(emit.inputs["Color"], FATE_START + 80, TREG_COLOR)
        kf_value(emit.inputs["Strength"], FATE_START, 2.5)
        kf_value(emit.inputs["Strength"], FATE_START + 80, 2.0)

    # cell stays alive: keep a slow breathing pulse to the end
    c = cell["cell"]
    kf_scale(c, FATE_START + 60, 1.03)
    kf_scale(c, FATE_START + 130, 1.0)
    kf_scale(c, FATE_START + 200, 1.03)
    kf_scale(c, 420, 1.0)

    # suppression field: ghostly teal sphere that expands outward
    bpy.ops.mesh.primitive_uv_sphere_add(radius=CELL_RADIUS,
                                          location=tuple(cell["center"]),
                                          segments=48, ring_count=24)
    field = bpy.context.active_object
    field.name = f"{cell['name']}_SuppressionField"
    field.modifiers.new(name="Subdiv", type='SUBSURF').levels = 1
    fm = bpy.data.materials.new(name=f"{cell['name']}_Field_Mat")
    fm.use_nodes = True
    fb = fm.node_tree.nodes["Principled BSDF"]
    fb.inputs["Base Color"].default_value = TREG_COLOR
    fb.inputs["Alpha"].default_value = 0.05
    if "Emission Color" in fb.inputs:
        fb.inputs["Emission Color"].default_value = TREG_COLOR
        fb.inputs["Emission Strength"].default_value = 0.8
    elif "Emission" in fb.inputs:
        fb.inputs["Emission"].default_value = TREG_COLOR
    field.data.materials.append(fm)
    kf_scale(field, FATE_START, 0.9)
    kf_scale(field, FATE_START, 0.9)
    kf_scale(field, FATE_START + 10, 1.0)
    kf_scale(field, FATE_START + 120, 1.9)
    kf_scale(field, 420, 2.1)
    # fade the field in then let it hang
    a = fb.inputs["Alpha"]
    kf_value(a, FATE_START, 0.0)
    kf_value(a, FATE_START + 60, 0.06)
    kf_value(a, 420, 0.04)
    return field

def animate_deletion(cell):
    """Apoptosis: membrane blebbing (Displace), darken to apoptotic red,
    shrink and fade to transparent (clonal deletion / removal)."""
    c = cell["cell"]

    # blebbing via a cloud-texture Displace whose strength ramps up
    tex = bpy.data.textures.new(f"{cell['name']}_Bleb", type='CLOUDS')
    tex.noise_scale = 0.35
    disp = c.modifiers.new(name="Bleb", type='DISPLACE')
    disp.texture = tex
    disp.strength = 0.0
    kf_mod(disp, "strength", FATE_START, 0.0)
    kf_mod(disp, "strength", FATE_START + 90, 0.55)   # peak blebbing
    kf_mod(disp, "strength", FATE_START + 160, 0.35)

    base = cell["mem_bsdf"].inputs["Base Color"]
    kf_color(base, FATE_START, ACTIVE_COLOR)
    kf_color(base, FATE_START + 70, DELETION_COLOR)

    for emit in cell["emit_nodes"]:
        kf_value(emit.inputs["Strength"], FATE_START, 2.5)
        kf_value(emit.inputs["Strength"], FATE_START + 70, 0.0)

    # shrink and fade out (Cycles honours the Principled Alpha input)
    kf_scale(c, FATE_START + 40, 1.0)
    kf_scale(c, FATE_START + 170, 0.25)
    alpha = cell["mem_bsdf"].inputs["Alpha"]
    kf_value(alpha, FATE_START + 60, 1.0)
    kf_value(alpha, FATE_START + 170, 0.0)
    kf_value(alpha, 420, 0.0)

# ------------------------------------------------------------
# TEXT LABEL under each cell (emissive; fades in as fate resolves)
# ------------------------------------------------------------
def make_label(text, center, color, fade_start):
    bpy.ops.object.text_add(location=(center.x, center.y - 0.2,
                                       center.z - CELL_RADIUS - 1.1))
    t = bpy.context.active_object
    t.name = f"Label_{text.replace(' ', '_')}"
    t.data.body = text
    t.data.align_x = 'CENTER'
    t.data.size = 0.55
    t.data.extrude = 0.02
    # stand the text upright so its face points at the camera (-Y)
    t.rotation_euler = (math.radians(90), 0, 0)

    m = bpy.data.materials.new(name=f"LblMat_{text}")
    m.use_nodes = True
    nt = m.node_tree
    nt.nodes.clear()
    o = nt.nodes.new("ShaderNodeOutputMaterial")
    e = nt.nodes.new("ShaderNodeEmission")
    e.inputs["Color"].default_value = color
    e.inputs["Strength"].default_value = 0.0
    nt.links.new(e.outputs["Emission"], o.inputs["Surface"])
    t.data.materials.append(m)

    s = e.inputs["Strength"]
    kf_value(s, fade_start, 0.0)
    kf_value(s, fade_start + 30, 4.0)
    kf_value(s, 420, 4.0)
    return t

# ------------------------------------------------------------
# ASSEMBLE THE SCENE
# ------------------------------------------------------------
centers = [(-CELL_SPACING, 0, 0), (0, 0, 0), (CELL_SPACING, 0, 0)]
cell_anergy   = build_tcell("Anergy",   centers[0], seed=7)
cell_treg     = build_tcell("Treg",     centers[1], seed=11)
cell_deletion = build_tcell("Deletion", centers[2], seed=17)
all_cells = [cell_anergy, cell_treg, cell_deletion]

# wheat pMHC flyers (one per cell — same complex)
flyers = []
for cell in all_cells:
    pmhc, start_loc, dock_loc = build_pmhc(cell["name"], cell)
    kf_loc(pmhc, 1, start_loc)
    kf_loc(pmhc, 60, start_loc)
    kf_loc(pmhc, 120, tuple(dock_loc))
    kf_loc(pmhc, 420, tuple(dock_loc))
    flyers.append(pmhc)

# idle breathing then fate resolution
for cell in all_cells:
    animate_engage_and_hold(cell)
animate_anergy(cell_anergy)
field = animate_treg(cell_treg)
animate_deletion(cell_deletion)

# labels (staggered fade-in near each fate's resolution)
lbl1 = make_label("ANERGY",           cell_anergy["center"],   (0.85, 0.85, 0.9, 1),  fade_start=230)
lbl2 = make_label("TOLERANCE (iTreg)", cell_treg["center"],    TREG_COLOR,            fade_start=250)
lbl3 = make_label("CLONAL DELETION",  cell_deletion["center"], (0.9, 0.35, 0.35, 1),  fade_start=300)

# ------------------------------------------------------------
# CAMERA (tracks a central empty; wide 3/4 elevated for hand-off)
# ------------------------------------------------------------
bpy.ops.object.empty_add(location=(0, 0, 0))
cam_target = bpy.context.active_object
cam_target.name = "Camera_Target"

bpy.ops.object.camera_add(location=(2, -15, 6.5))
camera = bpy.context.active_object
camera.name = "Main_Camera"
scene.camera = camera
camera.data.lens = 40
track = camera.constraints.new(type='TRACK_TO')
track.target = cam_target
track.track_axis = 'TRACK_NEGATIVE_Z'
track.up_axis = 'UP_Y'

kf_loc(camera, 1,   (2, -15, 6.5))     # wide establish
kf_loc(camera, 60,  (1, -13, 6.0))     # push in as pMHC arrive
kf_loc(camera, 150, (0, -11.5, 5.5))   # closest during engagement
kf_loc(camera, 300, (0, -12.5, 6.0))   # ease back as fates resolve
kf_loc(camera, 360, (3, -14, 7.0))     # settle to 3/4 elevated
kf_loc(camera, 420, (3, -14, 7.0))     # hold for ChimeraX hand-off

# ------------------------------------------------------------
# EASING
# ------------------------------------------------------------
ease_all_fcurves(camera)
for pmhc in flyers:
    ease_all_fcurves(pmhc)
for cell in all_cells:
    ease_all_fcurves(cell["cell"])
    ease_all_fcurves(cell["mat_cell"].node_tree)
    for mat in cell["receptor_mats"]:
        ease_all_fcurves(mat.node_tree)
ease_all_fcurves(field)
ease_all_fcurves(field.active_material.node_tree)
for lbl in (lbl1, lbl2, lbl3):
    ease_all_fcurves(lbl.active_material.node_tree)

# ------------------------------------------------------------
# LIGHTING (three-point-ish, tuned for Cycles + SSS, widened for 3 cells)
# ------------------------------------------------------------
bpy.ops.object.light_add(type='SUN', location=(6, -6, 9))
key = bpy.context.active_object
key.data.energy = 4.0
key.data.angle = math.radians(2)

bpy.ops.object.light_add(type='AREA', location=(-7, 4, 4))
fill = bpy.context.active_object
fill.data.energy = 500
fill.data.size = 8

bpy.ops.object.light_add(type='AREA', location=(0, 7, -2))
rim = bpy.context.active_object
rim.data.energy = 250
rim.data.size = 5

# Ensure a world exists (a fresh .blend / factory scene may have none)
if scene.world is None:
    scene.world = bpy.data.worlds.new("World")
scene.world.use_nodes = True

# Find the Background shader by TYPE, not by the literal name "Background"
# (a newly created / newer-default world may name it differently or lack it).
_wnt = scene.world.node_tree
bg = next((n for n in _wnt.nodes if n.type == 'BACKGROUND'), None)
if bg is None:
    bg = _wnt.nodes.new("ShaderNodeBackground")
    out = next((n for n in _wnt.nodes if n.type == 'OUTPUT_WORLD'), None)
    if out is None:
        out = _wnt.nodes.new("ShaderNodeOutputWorld")
    _wnt.links.new(bg.outputs["Background"], out.inputs["Surface"])

bg.inputs["Color"].default_value = (0.015, 0.015, 0.02, 1.0)
bg.inputs["Strength"].default_value = 1.0

print("Three-fate scene built (Cycles). Frames 1-420 @ 30 fps (~14 s). "
      "Left=Anergy, Center=iTreg tolerance, Right=Clonal deletion. "
      "Camera tracks a central empty; all camera keyframes are position-only.")

# ============================================================
# NOTES
#
# SWAPPING IN THE REAL WHEAT pMHC (wheat_pmhc.pdb)
#   Install "Molecular Nodes" (Edit > Preferences > Add-ons > search
#   "Molecular Nodes", or https://github.com/BradyAJohnston/MolecularNodes).
#   wheat_pmhc.pdb is the full pMHC-II complex: chain A + B (MHC-II
#   alpha/beta) with chain P (wheat 15-mer IHNVVHAIILHQQQQ) in the groove.
#   Steps:
#     1. Import wheat_pmhc.pdb via Molecular Nodes ONCE.
#     2. Style: cartoon for A/B, sticks/spacefill for chain P. Do this in
#        ONE call with the helper defined above:
#          style_by_chain(bpy.data.objects["wheat_pmhc"],
#                         {"Style Cartoon":        ["A", "B"],
#                          "Style Ball and Stick": ["P"]})
#        (use "Style Spheres" for spacefill on the peptide). Then tint the
#        peptide branch's color with PEPTIDE_COLOR so it reads as the
#        "antigen cargo". Confirm the peptide's chain ID in the MN
#        spreadsheet first -- if MN relabelled it, use that letter.
#     3. Delete the three "*_wheat_pMHC_PLACEHOLDER" cylinders.
#     4. Duplicate your imported MN object three times (one per cell),
#        rename to Anergy_wheat_pMHC / Treg_wheat_pMHC /
#        Deletion_wheat_pMHC, and copy the flyer location keyframes onto
#        each (kf_loc does start_loc @ 1/60, dock_loc @ 120/420).
#   REUSING FOR OTHER ANTIGENS: soy_pmhc.pdb and dairy_pmhc.pdb share the
#   identical A/B/P layout — just import a different file. One antigen
#   (wheat) as the example is sufficient; the mechanism scales across
#   peptides/allergens without changing anything else here.
#
# BIOLOGY / CONTENT NOTES (so captions/VO stay accurate)
#   * The object flying in is a peptide-MHC-II complex being PRESENTED,
#     not a free peptide docking a receptor.
#   * Left  = ANERGY: TCR engagement without costimulation -> functional
#     unresponsiveness. Cell SURVIVES (gray, dormant, smaller). Correct
#     to say "anergic / unresponsive", NOT "destroyed".
#   * Center= TOLERANCE (iTreg): antigen-specific conversion to a
#     regulatory phenotype. Cell stays alive and gains suppressive
#     function (the expanding halo). Say "converted to a regulatory
#     T cell / dominant tolerance".
#   * Right = CLONAL DELETION: apoptosis of the autoreactive/allergen-
#     specific clone (blebbing, fade-out). Say "deleted / removed".
#   These are the three arms of the platform shown as parallel fates of
#   one antigen-specific clone.
#
# TUNING
#   FATE_START (150) sets when all three fates begin. CASCADE_SPREAD_FRAMES
#   sets the anergy dim-wave speed. Labels fade in at 230/250/300 — nudge
#   these to sit just after each fate visually resolves.
#
# RENDER TIME
#   256 samples + denoising + SSS + a transparent field + alpha fade is
#   heavier than v2 (three cells). Do a 64-sample test pass for framing/
#   timing first, then bump to 256+ for the final.
#
# CHIMERAX HAND-OFF
#   Frame-1 camera (2,-15,6.5) is a wide 3/4 elevated establishing shot
#   matching the CLOSING wide of peptide_binding_movie.cxc. The ChimeraX
#   clip plays first (atomic scale), this Blender clip second (cellular),
#   so cross-dissolve the last ~15 frames of the ChimeraX clip into the
#   first ~15 frames here for a seamless atomic -> cellular transition.
# ============================================================
