# -*- coding: utf-8 -*-
# TorsoCAE Validation Journal — flexible_multibody_dynamics case 2: double pendulum
# (substitute for slider-crank; see README.md for why and for the validation method)

import sys
import numpy as np

sys.path.insert(0, "./server/")
from mesh_generators.structured import build_structured_mesh
from pytorsocae import TorsoCAESession

L1, L2 = 1.0, 1.0
WY, WZ = 0.08, 0.04
E, NU, RHO = 210e9, 0.3, 7800.0
G = 9.81


def _combine_disconnected_meshes(*meshes):
    coordinates, hexes, volume_tags, quads, surface_tags = [], [], [], [], []
    node_offset = 0
    for mesh in meshes:
        points = np.asarray(mesh["x"], dtype=np.float64)
        hex_block = mesh["topologies"][5]
        quad_block = mesh["topologies"][3]
        coordinates.append(points)
        hexes.append(np.asarray(hex_block["topology"], dtype=np.int32) + node_offset)
        volume_tags.append(np.asarray(hex_block["cell_data"], dtype=np.int32))
        quads.append(np.asarray(quad_block["topology"], dtype=np.int32) + node_offset)
        surface_tags.append(np.asarray(quad_block["cell_data"], dtype=np.int32))
        node_offset += len(points)
    return {
        "x": np.vstack(coordinates),
        "topologies": {
            5: {"topology": np.vstack(hexes), "cell_data": np.concatenate(volume_tags)},
            3: {"topology": np.vstack(quads), "cell_data": np.concatenate(surface_tags)},
        },
        "type_props": {5: {"dim": 3, "num_nodes": 8}, 3: {"dim": 2, "num_nodes": 4}},
        "gdim": 3,
    }


link_a = build_structured_mesh(
    bounds=((0.0, L1), (-WY / 2, WY / 2), (-WZ / 2, WZ / 2)), divisions=(10, 1, 1),
    surface_tags={"xmin": 101, "xmax": 102, "ymin": 103, "ymax": 104, "zmin": 105, "zmax": 106}, volume_tag=1,
)
link_b = build_structured_mesh(
    bounds=((L1, L1 + L2), (-WY / 2, WY / 2), (-WZ / 2, WZ / 2)), divisions=(10, 1, 1),
    surface_tags={"xmin": 201, "xmax": 202, "ymin": 203, "ymax": 204, "zmin": 205, "zmax": 206}, volume_tag=2,
)
combined = _combine_disconnected_meshes(link_a, link_b)

session = TorsoCAESession()
session.inline_mesh(combined, mesh_id="double_pend_mesh", name="Double Pendulum",
                     volume_names={1: "Link A", 2: "Link B"})
session.solid("Link A").material(E=E, nu=NU, rho=RHO)
session.solid("Link B").material(E=E, nu=NU, rho=RHO)
session.set_flexible_multibody(
    frames=[
        {"id": "ground", "origin": [0.0, 0.0, 0.0], "grounded": True},
        {"id": "link_a_root", "origin": [0.0, 0.0, 0.0]},
        {"id": "link_a_tip", "origin": [L1, 0.0, 0.0]},
        {"id": "link_b_root", "origin": [L1, 0.0, 0.0]},
    ],
    bodies=[
        {"id": "body_a", "kind": "flexible", "reference_frame_id": "link_a_root"},
        {"id": "body_b", "kind": "flexible", "reference_frame_id": "link_b_root"},
    ],
    attachments=[
        {"frame_id": "link_a_root", "surface_tag": 101, "body_id": "body_a", "dependent_axes": ["x", "y", "z"]},
        {"frame_id": "link_a_tip", "surface_tag": 102, "body_id": "body_a", "dependent_axes": ["x", "y", "z"]},
        {"frame_id": "link_b_root", "surface_tag": 201, "body_id": "body_b", "dependent_axes": ["x", "y", "z"]},
    ],
    joints=[
        {"id": "ground_hinge", "frame_a": "ground", "frame_b": "link_a_root", "kind": "revolute", "axis": "z"},
        {"id": "inter_link_hinge", "frame_a": "link_a_tip", "frame_b": "link_b_root", "kind": "revolute", "axis": "z"},
    ],
    kkt_augmentation=0.0,
    projection={"tolerance": 1.0e-9, "max_iterations": 15},
)
session.set_model_options(gravity=[0.0, -G, 0.0])
session.set_physics("structural", submodel="flexible_multibody_dynamics", backend="dolfinx")
session.set_solver_options(
    algo="auto", tol=1.0e-7, max_iter=800, num_steps=1200, dt=0.0025,
    time_scheme="generalized-alpha", spectral_radius_infinity=0.98,
    max_inner_iter=500, inner_tol=1.0e-5, n_cores=1, device="cpu",
)
session.compute(mesh_ids=["double_pend_mesh"])