import flow360 as fl
from flow360.examples import TutorialUDDForcesMoments

TutorialUDDForcesMoments.get_files()

project = fl.Project.from_geometry(
    TutorialUDDForcesMoments.geometry, name="Hinge Torques Monitoring"
)

geometry = project.geometry

geometry.group_edges_by_tag("edgeName")
geometry.group_faces_by_tag("groupName")

print("Geometry boundaries:")
for boundary in geometry.entity_info.get_boundaries():
    print(f"- {boundary.name}")

farfield = fl.AutomatedFarfield()

with fl.SI_unit_system:
    meshing = fl.MeshingParams(
        defaults=fl.MeshingDefaults(
            surface_max_edge_length=0.5,
            curvature_resolution_angle=10 * fl.u.deg,
            boundary_layer_first_layer_thickness=2e-6,
            boundary_layer_growth_rate=1.2,
        ),
        volume_zones=[farfield],
    )

with fl.SI_unit_system:
    box1 = fl.Box(name="box1", size=[2, 6, 3], center=[6.5, 9, 0], axis_of_rotation=[0, 1, 0])
    box2 = fl.Box(name="box2", size=[2, 6, 3], center=[6.5, -9, 0], axis_of_rotation=[0, 1, 0])
    box3 = fl.Box(name="box3", size=[4, 8, 3], center=[12, 0, 2], axis_of_rotation=[0, 1, 0])

    meshing.refinements.extend(
        [
            fl.SurfaceEdgeRefinement(
                method=fl.AngleBasedRefinement(value=1 * fl.u.deg),
                edges=[geometry["leadingEdge"]],
            ),
            fl.SurfaceEdgeRefinement(
                method=fl.HeightBasedRefinement(value=5e-3),
                edges=[geometry["trailingEdge"]],
            ),
            fl.SurfaceRefinement(max_edge_length=0.5, faces=[geometry["wing*"]]),
            fl.UniformRefinement(name="box_refinement1", entities=[box1, box2, box3], spacing=0.2),
        ]
    )

with fl.SI_unit_system:
    reference_geometry = fl.ReferenceGeometry(
        area=60, moment_center=[5.7542, 0, 0], moment_length=[1, 1, 1]
    )
    operating_condition = fl.AerospaceCondition(
        velocity_magnitude=50,
        alpha=10 * fl.u.deg,
        atmosphere=fl.ThermalState(temperature=288.15),
    )
    models = [
        fl.Fluid(
            navier_stokes_solver=fl.NavierStokesSolver(
                absolute_tolerance=1e-9,
                linear_solver=fl.LinearSolver(max_iterations=35),
            ),
            turbulence_model_solver=fl.SpalartAllmaras(
                linear_solver=fl.LinearSolver(max_iterations=25)
            ),
        ),
        fl.Wall(
            surfaces=[geometry["*Left"], geometry["*Right"], geometry["fuselage"]],
        ),
        fl.Freestream(surfaces=[farfield.farfield]),
    ]
    time_stepping = fl.Steady(max_steps=5000)

right_aileron_torque_center = [5.7542, 7, 0] * fl.u.m
right_aileron_torque_axis = [0, 1, 0]

right_aileron_hinge_torque_partial_dim = fl.UserVariable(
    name="right_aileron_hinge_torque_partial",
    value=fl.math.dot(
        fl.math.cross(
            fl.math.subtract(fl.solution.coordinate, right_aileron_torque_center),
            fl.solution.node_forces_per_unit_area,
        ),
        right_aileron_torque_axis,
    ),
).in_units(new_unit=fl.u.Pa * fl.u.m)

left_aileron_torque_center = [5.7542, -7, 0] * fl.u.m
left_aileron_torque_axis = [0, -1, 0]

left_aileron_hinge_torque_partial_dim = fl.UserVariable(
    name="left_aileron_hinge_torque_partial",
    value=fl.math.dot(
        fl.math.cross(
            fl.math.subtract(fl.solution.coordinate, left_aileron_torque_center),
            fl.solution.node_forces_per_unit_area,
        ),
        left_aileron_torque_axis,
    ),
).in_units(new_unit=fl.u.Pa * fl.u.m)

right_rudder_torque_center = [12.01, 0.861, 0.861] * fl.u.m
right_rudder_torque_axis = [0, fl.math.sqrt(2) / 2, fl.math.sqrt(2) / 2]

right_rudder_hinge_torque_partial_dim = fl.UserVariable(
    name="right_rudder_hinge_torque_partial",
    value=fl.math.dot(
        fl.math.cross(
            fl.math.subtract(fl.solution.coordinate, right_rudder_torque_center),
            fl.solution.node_forces_per_unit_area,
        ),
        right_rudder_torque_axis,
    ),
).in_units(new_unit=fl.u.Pa * fl.u.m)

left_rudder_torque_center = [12.01, -0.861, 0.861] * fl.u.m
left_rudder_torque_axis = [0, -fl.math.sqrt(2) / 2, fl.math.sqrt(2) / 2]

left_rudder_hinge_torque_partial_dim = fl.UserVariable(
    name="left_rudder_hinge_torque_partial",
    value=fl.math.dot(
        fl.math.cross(
            fl.math.subtract(fl.solution.coordinate, left_rudder_torque_center),
            fl.solution.node_forces_per_unit_area,
        ),
        left_rudder_torque_axis,
    ),
).in_units(new_unit=fl.u.Pa * fl.u.m)

outputs = [
    fl.SurfaceIntegralOutput(
        name="right_aileron_hinge_torque",
        output_fields=[right_aileron_hinge_torque_partial_dim],
        surfaces=geometry["aileronRight"],
    ),
    fl.SurfaceIntegralOutput(
        name="left_aileron_hinge_torque",
        output_fields=[left_aileron_hinge_torque_partial_dim],
        surfaces=geometry["aileronLeft"],
    ),
    fl.SurfaceIntegralOutput(
        name="right_rudder_hinge_torque",
        output_fields=[right_rudder_hinge_torque_partial_dim],
        surfaces=geometry["rudderRight"],
    ),
    fl.SurfaceIntegralOutput(
        name="left_rudder_hinge_torque",
        output_fields=[left_rudder_hinge_torque_partial_dim],
        surfaces=geometry["rudderLeft"],
    ),
]

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing,
        operating_condition=operating_condition,
        reference_geometry=reference_geometry,
        models=models,
        outputs=outputs,
    )

case = project.run_case(params, name="Steady RANS (hinge moments)", use_beta_mesher=True)

case.wait()
results = case.results

nonlinear_residuals_data = results.nonlinear_residuals.as_dataframe()
nonlinear_residuals_data.plot(
    "pseudo_step",
    ["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"],
    logy=True,
    xlim=[10, 2000],
    grid=True,
)

left_rudder_hinge_torque_history = results.monitors["left_rudder_hinge_torque"].as_dataframe()
left_rudder_hinge_torque_history.plot(
    "pseudo_step", "left_rudder_hinge_torque_partial_integral", grid=True
)