#!/usr/bin/env python3

from pathlib import Path

import acts

mm = acts.UnitConstants.mm

degree = acts.UnitConstants.degree

root = acts.Blueprint(envelope=acts.ExtentEnvelope(r=[10 * mm, 10 * mm]))

pixel = root.addCylinderContainer(direction=acts.AxisDirection.AxisZ, name="Pixel")

print(repr(pixel))

trf = acts.Transform3.Identity() * acts.Translation3(acts.Vector3(0, 0, 0 * mm))

if True:

barrel = acts.CylinderContainerBlueprintNode(

"PixelBarrel",

acts.AxisDirection.AxisR,

attachmentStrategy=acts.CylinderVolumeStack.AttachmentStrategy.Gap,

resizeStrategy=acts.CylinderVolumeStack.ResizeStrategy.Gap,

)

pixel.addChild(barrel)

print("Barrel")

r = 25 * mm

for i in range(0, 4):

r += 50 * mm

bounds = acts.CylinderVolumeBounds(r, r + 20 * mm, 200 * mm)

print(bounds)

brlLayer = barrel.addStaticVolume(trf, bounds, name=f"PixelBarrelLayer{i}")

assert brlLayer.name == f"PixelBarrelLayer{i}"

if True:

with pixel.CylinderContainer("PixelPosEndcap", acts.AxisDirection.AxisZ) as ec:

print("Positive Endcap")

ec.attachmentStrategy = acts.CylinderVolumeStack.AttachmentStrategy.Gap

ec.resizeStrategy = acts.CylinderVolumeStack.ResizeStrategy.Gap

z = 200

for i in range(0, 4):

z += 200 * mm

bounds = acts.CylinderVolumeBounds(100 * mm, 150 * mm, 50 * mm)

print(bounds)

trf = acts.Transform3.Identity() * acts.Translation3(acts.Vector3(0, 0, z))

with ec.StaticVolume(trf, bounds, name=f"PixelPosEndcapDisk{i}") as disc:

print("Add disk", i)

assert disc.name == f"PixelPosEndcapDisk{i}"

if True:

with pixel.Material() as mat:

with mat.CylinderContainer(

direction=acts.AxisDirection.AxisZ, name="PixelNegEndcap"

) as ec:

ec.attachmentStrategy = acts.CylinderVolumeStack.AttachmentStrategy.Gap

print("Negative Endcap")

z = -200

for i in range(0, 4):

z -= 200 * mm

bounds = acts.CylinderVolumeBounds(200 * mm, 300 * mm, 50 * mm)

print(bounds)

trf = acts.Transform3.Identity() * acts.Translation3(

acts.Vector3(0, 0, z)

)

with ec.StaticVolume(

trf, bounds, name=f"PixelNegEndcapDisk{i}"

) as disk:

print("Add disk", i)

assert disk.name == f"PixelNegEndcapDisk{i}"

with open("blueprint.dot", "w") as fh:

root.graphviz(fh)

gctx = acts.GeometryContext.dangerouslyDefaultConstruct()

trackingGeometry = root.construct(

options=acts.BlueprintNode.Options(), gctx=gctx, level=acts.logging.VERBOSE

)

vis = acts.ObjVisualization3D()

trackingGeometry.visualize(vis, gctx)

with Path("blueprint.obj").open("w") as fh:

vis.write(fh)

# print("DONE")