import flow360 as fl
from flow360.examples import OM6wing

OM6wing.get_files()

project = fl.Project.from_volume_mesh(
    OM6wing.mesh_filename, name="Tutorial UDD alpha controller from Python"
)

volume_mesh = project.volume_mesh

print("Volume mesh boundaries:")
for boundary in volume_mesh.boundary_names:
    print(f"  - {boundary}")

with fl.SI_unit_system:
    operating_condition = fl.AerospaceCondition.from_mach_reynolds(
        reynolds_mesh_unit=11.72e6,
        mach=0.84,
        project_length_unit=0.80167 * fl.u.m,
        temperature=297.78,
        alpha=3.06 * fl.u.deg,
        beta=0 * fl.u.deg,
    )

with fl.SI_unit_system:
    reference_geometry = fl.ReferenceGeometry(
        area=1.15315,
        moment_center=[0, 0, 0],
        moment_length=[1.47601, 0.80167, 1.47601],
    )
    time_stepping = fl.Steady(
        max_steps=2000,
    )

with fl.SI_unit_system:
    models = [
        fl.Wall(surfaces=[volume_mesh["1"]]),
        fl.SlipWall(surfaces=[volume_mesh["2"]]),
        fl.Freestream(surfaces=[volume_mesh["3"]]),
    ]

with fl.SI_unit_system:
    outputs = [
        fl.SurfaceOutput(
            surfaces=volume_mesh["1"],
            output_fields=[
                "Cp",
                "Cf",
                "CfVec",
                "yPlus",
            ],
        )
    ]

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

alpha_controller = fl.UserDefinedDynamic(
    name="alphaController",
    input_vars=["CL"],
    constants={"CLTarget": 0.4, "Kp": 0.2, "Ki": 0.002},
    output_vars={"alphaAngle": "if (pseudoStep > 500) state[0]; else alphaAngle;"},
    state_vars_initial_value=["alphaAngle", "0.0"],
    update_law=[
        "if (pseudoStep > 500) state[0] + Kp * (CLTarget - CL) + Ki * state[1]; else state[0];",
        "if (pseudoStep > 500) state[1] + (CLTarget - CL); else state[1];",
    ],
    input_boundary_patches=[volume_mesh["1"]],
)
params.user_defined_dynamics = [alpha_controller]

case = project.run_case(
    params=params, name="Case of tutorial UDD alpha controller from Python"
)

case.wait()

udd_data = case.results.user_defined_dynamics["alphaController"]

ax0 = udd_data.as_dataframe().plot(
    x="pseudo_step",
    y="alphaAngle",
    title="AoA",
    ylabel="AoA [deg] - blue",
    legend=False,
)
ax1 = ax0.twinx()
udd_data.as_dataframe().plot(
    x="pseudo_step", y="CL", ylabel="CL - orange", ax=ax1, color="orange", legend=False
)

ax0 = udd_data.as_dataframe().plot(
    x="pseudo_step", y="state[0]", title="state", ylabel="state[0] - blue", legend=False
)
ax1 = ax0.twinx()
udd_data.as_dataframe().plot(
    x="pseudo_step",
    y="state[1]",
    ylabel="state[1] - orange",
    ax=ax1,
    color="orange",
    legend=False,
)