import flow360 as fl
from flow360.examples import TutorialPeriodicBC

TutorialPeriodicBC.get_files()

project = fl.Project.from_volume_mesh(
    TutorialPeriodicBC.mesh_filename,
    length_unit="mm",
    name="Tutorial Periodic Boundary Condition from Python",
)

volume_mesh = project.volume_mesh

with fl.SI_unit_system:

    operating_condition = fl.AerospaceCondition.from_mach_reynolds(
        mach=0.13989,
        reynolds_mesh_unit=3200,
        project_length_unit=1 * fl.u.mm,
        temperature=298.25 * fl.u.K,
    )

with fl.SI_unit_system:
    navier_stokes_solver = fl.NavierStokesSolver(
        absolute_tolerance=1e-11,
        linear_solver=fl.LinearSolver(max_iterations=20),
        order_of_accuracy=2,
        kappa_MUSCL=0.33,
        update_jacobian_frequency=1,
        equation_evaluation_frequency=1,
        numerical_dissipation_factor=1,
    )

    turbulence_model_solver = fl.SpalartAllmaras(
        absolute_tolerance=1e-10,
        linear_solver=fl.LinearSolver(max_iterations=20),
        update_jacobian_frequency=1,
        equation_evaluation_frequency=1,
        order_of_accuracy=2,
    )

    fluid_model = fl.Fluid(
        navier_stokes_solver=navier_stokes_solver,
        turbulence_model_solver=turbulence_model_solver,
    )

with fl.SI_unit_system:
    models_list = [
        fluid_model,
        fl.Wall(
            surfaces=[
                volume_mesh["fluid/vane_*"],
                volume_mesh["fluid/bladeFillet_*"],
                volume_mesh["fluid/shroud"],
                volume_mesh["fluid/hub"],
            ]
        ),
        fl.Freestream(
            surfaces=[
                volume_mesh["fluid/inlet"],
            ]
        ),
        fl.Outflow(
            surfaces=[
                volume_mesh["fluid/outlet"],
            ],
            spec=fl.Pressure(
                1.0032 * operating_condition.thermal_state.pressure
            ),  # Slightly above ambient
        ),
        fl.SlipWall(
            surfaces=[
                volume_mesh["fluid/bottomFront"],
                volume_mesh["fluid/topFront"],
            ]
        ),
    ]

models_list.append(
    fl.Periodic(
        surface_pairs=[
            (volume_mesh["fluid/periodic-1"], volume_mesh["fluid/periodic-2"])
        ],  # Connects the two radially oppose mesh surfaces
        spec=fl.Rotational(
            axis_of_rotation=(1, 0, 0)
        ),  # Specify rotation around X-axis
    ),
)

with fl.SI_unit_system:
    slice_inlet = fl.Slice(
        name="Inlet",
        normal=[1, 0, 0],
        origin=[-179, 0, 0] * fl.u.m,
    )
    slice_outlet = fl.Slice(
        name="Outlet",
        normal=[1, 0, 0],
        origin=[539, 0, 0] * fl.u.m,
    )
    slice_trailing_edge = fl.Slice(
        name="TrailingEdge",
        normal=[1, 0, 0],
        origin=[183, 0, 0] * fl.u.m,
    )
    slice_wake = fl.Slice(
        name="Wake",
        normal=[1, 0, 0],
        origin=[294.65, 0, 0] * fl.u.m,
    )
    outputs_list = [
        fl.VolumeOutput(
            output_format="tecplot",
            output_fields=[
                "primitiveVars",
                "vorticity",
                "Cp",
                "Mach",
                "qcriterion",
                "mut",
                "nuHat",
                "mutRatio",
                "gradW",
                "T",
                "residualNavierStokes",
            ],
        ),
        fl.SurfaceOutput(
            surfaces=volume_mesh["*"],
            output_format="both",
            output_fields=[
                "primitiveVars",
                "Cp",
                "Cf",
                "CfVec",
                "yPlus",
                "nodeForcesPerUnitArea",
            ],
        ),
        fl.SliceOutput(
            slices=[slice_inlet, slice_outlet, slice_trailing_edge, slice_wake],
            output_format="both",
            output_fields=["Cp", "primitiveVars", "T", "Mach", "gradW"],
        ),
    ]

with fl.SI_unit_system:
    params = fl.SimulationParams(
        reference_geometry=fl.ReferenceGeometry(
            moment_center=[0, 0, 0], moment_length=[1, 1, 1], area=209701.3096271187
        ),
        operating_condition=operating_condition,
        time_stepping=fl.Steady(max_steps=5000, CFL=fl.AdaptiveCFL()),
        models=models_list,
        outputs=outputs_list,
    )

case = project.run_case(
    params=params, name="Tutorial Periodic Boundary Condition from Python"
)

case.wait()

results = case.results

nonlinear_residuals = results.nonlinear_residuals.as_dataframe()
ax = nonlinear_residuals.plot(
    x="pseudo_step",
    y=["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"],
    xlim=(0, None),
    xlabel="Pseudo Step",
    secondary_y=["5_nuHat"],
    figsize=(10, 7),
    logy=True,
    title="Nonlinear residuals",
)
ax.grid(True)