import flow360 as fl
from flow360.examples import TutorialUDDStructural
import numpy as np

TutorialUDDStructural.get_files()
project = fl.Project.from_volume_mesh(
    TutorialUDDStructural.mesh_filename, name="UDD structural tutorial from Python"
)
vm = project.volume_mesh

u_inf = 15 * fl.u.m / fl.u.s
ref_length = (4 * fl.u.inch).to("m")
ref_density = 1.226 * fl.u.kg / fl.u.m**3
ref_T = 288.15 * fl.u.K
kinematic_viscosity = 14.61e-6 * fl.u.m**2 / fl.u.s

R = 287.058 * fl.u.J / (fl.u.kg * fl.u.K)
speed_of_sound = np.sqrt(1.4 * R * ref_T).to("m/s")
Mach = u_inf / speed_of_sound
Reynolds = u_inf * ref_length / kinematic_viscosity

mesh_unit = project.length_unit
Re_mesh_unit = Reynolds * mesh_unit.to_value() / ref_length
operating_condition = fl.AerospaceCondition.from_mach_reynolds(
    mach=Mach,
    reynolds_mesh_unit=Re_mesh_unit,
    project_length_unit=project.length_unit,
)

freestream = fl.Freestream(entities=vm["farFieldBlock/farField"])
slip_wall = fl.SlipWall(entities=vm["*/slipWall"])
no_slip_wall = fl.Wall(entities=vm["plateBlock/noSlipWall"])
BCs = [freestream, slip_wall, no_slip_wall]

volume_rotation = vm["plateBlock"]
volume_rotation.center = (0, 0, 0) * fl.u.m
volume_rotation.axis = [0, 1, 0]
rotation_model = fl.Rotation(
    entities=[volume_rotation], spec=fl.FromUserDefinedDynamics()
)

deg2rad = np.pi / 180

I = 0.443768309310345  # Structural moment of inertia
zeta = 0.005  # Structural damping ratio
theta0 = 5 * deg2rad  # Initial equilibrium angle
iniPert = 45 * deg2rad  # Initial perturbation
K = 0.005  # Stiffness of the spring
omegaN = np.sqrt(K / I)  # Structural natural angular frequency
initTheta = theta0 + iniPert
initOmegaDot = (0 - K * (initTheta - theta0)) / I

deltaNonDimensional = 0.5
delt_SI = deltaNonDimensional * (ref_length / speed_of_sound)
n_steps = 3000
time_stepping = fl.Unsteady(
    max_pseudo_steps=100,
    step_size=delt_SI,
    steps=n_steps,
)

dynamic = fl.UserDefinedDynamic(
    name="dynamicTheta",
    input_vars=["momentY"],
    constants={
        "I": I,
        "zeta": zeta,
        "K": K,
        "omegaN": omegaN,
        "theta0": theta0,
    },
    output_vars={
        "omegaDot": "state[0];",
        "omega": "state[1];",
        "theta": "state[2];",
    },
    state_vars_initial_value=[str(initOmegaDot), "0.0", str(initTheta)],
    update_law=[
        "if (pseudoStep == 0) (momentY - K * ( state[2] - theta0 ) - 2 * zeta * omegaN * I *state[1] ) / I; else state[0];",
        "if (pseudoStep == 0) state[1] + state[0] * timeStepSize; else state[1];",
        "if (pseudoStep == 0) state[2] + state[1] * timeStepSize; else state[2];",
    ],
    input_boundary_patches=vm["plateBlock/noSlipWall"],
    output_target=vm["plateBlock"],
)

with fl.SI_unit_system:
    outputs = [
        fl.SliceOutput(
            slices=[fl.Slice(name="y-slice", origin=(0, 0, 0), normal=(0, 1, 0))],
            output_format="tecplot",
            output_fields=["primitiveVars", "Mach"],
            frequency=10,
            frequency_offset=200,
        ),
        fl.VolumeOutput(
            output_format="tecplot",
            output_fields=["primitiveVars", "qcriterion", "Mach"],
        ),
    ]

with fl.SI_unit_system:
    param = fl.SimulationParams(
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(0, 0, 0),
            moment_length=(1, 1, 1) * mesh_unit,
            area=0.5325 * mesh_unit * mesh_unit,
        ),
        operating_condition=operating_condition,
        time_stepping=time_stepping,
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-9,
                    relative_tolerance=1e-3,
                    linear_solver=fl.LinearSolver(max_iterations=35),
                    kappa_MUSCL=0.33,
                    order_of_accuracy=2,
                    update_jacobian_frequency=1,
                    equation_evaluation_frequency=1,
                    low_mach_preconditioner=True,
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    absolute_tolerance=1e-8,
                    relative_tolerance=1e-2,
                    linear_solver=fl.LinearSolver(max_iterations=35),
                    order_of_accuracy=2,
                    rotation_correction=True,
                    update_jacobian_frequency=1,
                    reconstruction_gradient_limiter=0.5,
                    equation_evaluation_frequency=1,
                ),
            ),
            *BCs,
            rotation_model,
        ],
        outputs=outputs,
        user_defined_dynamics=[dynamic],
    )

case = project.run_case(params=param, name="UDD structural tutorial from Python")

import matplotlib.pyplot as plt

case.wait()

results = case.results

total_forces = results.total_forces.as_dataframe()
total_forces_2000 = total_forces[total_forces["physical_step"] >= 2000]
total_forces_2000.plot("physical_step", ["CL", "CD"])

non_linear = results.nonlinear_residuals
non_linear.filter_physical_steps_only()
non_linear = non_linear.as_dataframe()
non_linear_2000 = non_linear[non_linear["physical_step"] >= 2000]
non_linear_2000.plot(
    "physical_step",
    ["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"],
    logy=True,
)

plt.show()