"""Finer CAD workflow shown in the gallery animation.

The numerical gallery comparison uses the separate 14,027-element reference
mesh, assessed by run_validation.py. This journal generates its own fine mesh.
Upload cadex/Bearing_Bracket_Assy_STEP214.STEP before replay.
"""
import numpy as np
from pytorsocae import TorsoCAESession

session = TorsoCAESession()
session.upload("cadex/Bearing_Bracket_Assy_STEP214.STEP", body_id="bracket_source")

# Remove the separate bearing insert, retaining the bracket CAD solid.
session.delete_solid(
    solid_id="solid:2",
    occurrence_path="free_shape[1]/0:1:1:1/component[2]/0:1:1:1:2/solid[1]",
)
# The reference preprocessing replaces the spherical bearing seat with a
# 10.3 mm radius cylindrical bore centred at (101.6, -2.07) mm.
session.csg_make_shape(
    "cylinder",
    {"r": 10.3, "h": 40.0, "cx": 101.6, "cy": -2.07, "cz": -20.0},
    name="Bearing bore tool",
    body_id="bearing_tool",
)
session.csg_boolean(
    "cut",
    "bracket_source",
    "bearing_tool",
    name="Airplane bearing bracket",
    body_id="bracket",
)
session.csg_delete("bearing_tool")
session.import_delete(body_id="bracket_source")
# Partition the CAD into the original protected bearing and bolt regions.
# Partition retains both sides of each interface, making the mesh conform to it.
session.csg_make_shape(
    "cylinder",
    {"r": 13.97, "h": 40.0, "cx": 101.6, "cy": -2.20713268418, "cz": -20.0},
    name="Protected bearing region",
    body_id="partition_tool",
)
session.csg_boolean(
    "partition",
    "bracket",
    "partition_tool",
    name="Airplane bearing bracket",
    body_id="bracket",
)
session.csg_delete("partition_tool")
for x, z in ((36.83, 8.255), (36.83, -8.255), (0.0, 9.525), (0.0, -9.525)):
    session.csg_make_shape(
        "cylinder",
        {"r": 4.572, "h": 3.175, "cx": x, "cy": -3.175, "cz": z},
        name="Protected bolt region",
        body_id="partition_tool",
    )
    session.csg_transform(
        "partition_tool",
        "rotate",
        {
            "ax": 1.0,
            "ay": 0.0,
            "az": 0.0,
            "angle_deg": -90.0,
            "cx": x,
            "cy": -3.175,
            "cz": z,
        },
    )
    session.csg_boolean(
        "partition",
        "bracket",
        "partition_tool",
        name="Airplane bearing bracket",
        body_id="bracket",
    )
    session.csg_delete("partition_tool")

# STEP dimensions are millimetres; the analysis uses SI metres.
session.csg_transform(
    "bracket",
    "scale",
    {"sx": 0.001, "sy": 0.001, "sz": 0.001},
    linear_deflection=0.00005,
)
session.select_body("bracket")
# Fine CAD mesh with the published fine-case 4 mm sensitivity filter below.
# This mesh is generated here; it is not the website's 166,000-element mesh.
bolt_centres = np.array(
    [[0.03683, 0.008255], [0.03683, -0.008255], [0.0, 0.009525], [0.0, -0.009525]]
)
session.mesh(
    mesh_id="bracket_mesh",
    algo_id="hxt_frontal",
    size_min=0.0015,
    size_max=0.0015,
    order=1,
    dim=3,
)
session.solid("Airplane bearing bracket").material(E=200e9, nu=0.27, rho=7833.0)
session.set_physics("structural", "linear_elastic", backend="dolfinx")
session.set_solver_options(algo="direct", n_cores=4)

# Identify the new mesh's bore faces geometrically, without copied mesh tags.
mesh = session.mesh_data("bracket_mesh")
points = np.asarray(mesh["x"])
triangles = np.asarray(mesh["topologies"][2]["topology"])
tags = np.asarray(mesh["topologies"][2]["cell_data"])
bearing_tags = []
bolt_tags = [[] for _ in bolt_centres]
for tag in np.unique(tags):
    vertices = points[np.unique(triangles[tags == tag])]
    if np.allclose(
        np.linalg.norm(vertices[:, :2] - [0.1016, -0.00207], axis=1),
        0.0103,
        atol=1e-7,
        rtol=0,
    ):
        bearing_tags.append(int(tag))
    for index, centre in enumerate(bolt_centres):
        radius = np.linalg.norm(vertices[:, [0, 2]] - centre, axis=1)
        if (
            np.allclose(radius, 0.002413, atol=1e-7, rtol=0)
            and vertices[:, 1].min() >= -0.003175 - 1e-7
            and vertices[:, 1].max() <= 1e-7
        ):
            bolt_tags[index].append(int(tag))
if len(bearing_tags) != 1 or not all(bolt_tags):
    raise ValueError(
        f"Expected one bearing bore and four bolt bores; found {bearing_tags}, {bolt_tags}"
    )
for group in bolt_tags:
    for tag in group:
        session.surface(tag, scope_id="bracket").bc("fixed")
bearing_tag = bearing_tags[0]
session.surface(bearing_tag, scope_id="bracket").bc(
    "rbe2", reference=[0.1016, -0.00207, 0.0]
)
for case_id, force in (
    ("F1", [281.0, 0.0, 0.0]),
    ("F2", [298.0, -298.0, 0.0]),
    ("F3", [0.0, -562.0, 0.0]),
):
    session.analysis_case(case_id).surface(bearing_tag, scope_id="bracket").bc(
        "rbe2", reference_force=force
    )

# Identify the CAD-partitioned protected cells on this new mesh:
# the bearing annulus and four bolt sleeves. No reference element IDs are reused.
cells = np.asarray(mesh["topologies"][4]["topology"])
vertices = points[cells]
centres = vertices.mean(axis=1)
volumes = np.abs(np.linalg.det(vertices[:, 1:] - vertices[:, :1])) / 6.0
edge_squared = np.stack(
    [
        np.sum((vertices[:, i] - vertices[:, j]) ** 2, axis=1)
        for i, j in ((0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3))
    ],
    axis=1,
)
mean_ratio = 12.0 * (3.0 * volumes) ** (2.0 / 3.0) / edge_squared.sum(axis=1)
if mean_ratio.min() < 0.1 or np.sqrt(edge_squared.max()) > 0.015:
    raise ValueError(
        "Mesh-quality check failed: "
        f"minimum mean ratio={mean_ratio.min():.6g} (required >= 0.1); "
        f"{np.count_nonzero(mean_ratio < 0.1)} cells below the quality limit; "
        f"longest edge={np.sqrt(edge_squared.max()):.6g} m (required <= 0.015 m)"
    )
if abs(volumes.sum() / 0.0001091733129844137 - 1.0) > 0.01:
    raise ValueError("Bracket mesh volume differs from the reference by more than 1%")
print(
    "Bracket mesh:", len(cells), "tetrahedra; minimum quality", float(mean_ratio.min())
)
protected = (
    np.linalg.norm(centres[:, :2] - [0.1016, -0.00220713268418], axis=1) <= 0.01397
)
for centre in bolt_centres:
    protected |= (
        (np.linalg.norm(centres[:, [0, 2]] - centre, axis=1) <= 0.004572)
        & (centres[:, 1] >= -0.003175 - 1e-9)
        & (centres[:, 1] <= 1e-9)
    )

# Select whole CAD subdomains; do not classify individual interface cells.
volume_tags = np.asarray(mesh["topologies"][4]["cell_data"])
protected_tags = []
for tag in np.unique(volume_tags):
    region = volume_tags == tag
    fraction = volumes[region & protected].sum() / volumes[region].sum()
    if 0.05 < fraction < 0.95:
        raise ValueError(
            "A CAD volume mixes design and protected material; partitioning failed"
        )
    if fraction >= 0.95:
        protected_tags.append(int(tag))
if len(protected_tags) != 5 or len(np.unique(volume_tags)) != 6:
    raise ValueError("Expected one design volume and five protected CAD volumes")
protected = np.isin(volume_tags, protected_tags)

if abs(volumes[protected].sum() / 2.5506967618015697e-6 - 1.0) > 0.05:
    raise ValueError(
        "Protected-region volume differs from the reference by more than 5%"
    )

study = session.optimization_study(
    name="CAD-first BESO airplane bearing bracket",
    mesh_id="bracket_mesh",
    optimizer="beso",
    volume_fraction=0.1,
    filter_radius=0.004,
    minimum_stiffness=1e-6,
    preserve_cell_ids=np.flatnonzero(protected).tolist(),
    cases=["F1", "F2", "F3"],
    beso_options={
        "addition_ratio": 0.04,
        "removal_ratio": 0.08,
        "stress_limit": 1e9,
        "failure_tolerance": 1,
        "decay": -0.2,
        "void_density": 1e-6,
        "protect_boundaries": False,
    },
)
result = session.optimize(study, max_iterations=150, checkpoint_interval=5)
print("Optimization stopping reason:", result["optimization"]["reason"])
session.save_project(
    "topology_beso_bracket.tcae",
    include_results=True,
    post_scene={
        "field_key": {"name": "Retained Material"},
        "field_view_threshold": "0.5",
    },
)