import flow360 as fl
from flow360.examples import TutorialRANSXv15
import matplotlib.pyplot as plt

TutorialRANSXv15.get_files()
project = fl.Project.from_volume_mesh(
    TutorialRANSXv15.mesh_filename,
    name="Tutorial Time-accurate RANS CFD on XV-15 from Python",
)
volume_mesh = project.volume_mesh

with fl.SI_unit_system:
    rotation_zone = volume_mesh["innerRotating"]
    rotation_zone.center = (0, 0, 0) * fl.u.m
    rotation_zone.axis = (0, 0, -1)

    rotation_model = fl.Rotation(
        volumes=rotation_zone,
        spec=fl.AngularVelocity(600 * fl.u.rpm),
    )

time_model = fl.Unsteady(
    max_pseudo_steps=35,
    steps=600,
    step_size=0.5 / 600 * fl.u.s,
    CFL=fl.AdaptiveCFL(),
)

ddes_model = fl.SpalartAllmaras(
    absolute_tolerance=1e-8,
    linear_solver=fl.LinearSolver(max_iterations=25),
    hybrid_model=fl.DetachedEddySimulation(shielding_function="DDES"),
    rotation_correction=True,
    equation_evaluation_frequency=1,
)

with fl.SI_unit_system:
    params = fl.SimulationParams(
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(0, 0, 0),
            moment_length=(3.81, 3.81, 3.81),
            area=45.604,
        ),
        operating_condition=fl.AerospaceCondition(
            velocity_magnitude=5,
            alpha=-90 * fl.u.deg,
            reference_velocity_magnitude=238.14,
        ),
        time_stepping=time_model,
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-9,
                    linear_solver=fl.LinearSolver(max_iterations=35),
                    limit_velocity=True,
                    limit_pressure_density=True,
                ),
                turbulence_model_solver=ddes_model,
            ),
            fl.Freestream(surfaces=volume_mesh["farField/farField"]),
            fl.Wall(surfaces=volume_mesh["innerRotating/blade"]),
            rotation_model,
        ],
        outputs=[
            fl.VolumeOutput(
                output_fields=[
                    "primitiveVars",
                    "T",
                    "Cp",
                    "Mach",
                    "qcriterion",
                    "VelocityRelative",
                ],
            ),
            fl.SurfaceOutput(
                surfaces=volume_mesh["*"],
                output_fields=[
                    "primitiveVars",
                    "Cp",
                    "Cf",
                    "CfVec",
                    "yPlus",
                    "nodeForcesPerUnitArea",
                ],
            ),
        ],
    )

case = project.run_case(
    params=params,
    name="Tutorial Time-accurate DDES CFD on XV-15 from Python",
)

case.wait()
results = case.results

total_forces = case.results.total_forces.as_dataframe()
total_forces.plot("physical_step", ["CL", "CD"])

non_linear = case.results.nonlinear_residuals.as_dataframe()
non_linear.plot(
    "physical_step",
    ["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"],
    logy=True,
)
plt.xlim(550, 600)
plt.show()