import flow360 as fl
from flow360.examples import Tutorial2DCRM

Tutorial2DCRM.get_files()

project = fl.Project.from_geometry(
    Tutorial2DCRM.geometry, name="Tutorial 2D CRM from Python"
)

geometry = project.geometry

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

with fl.SI_unit_system:

    cylinders = [
        fl.Cylinder(
            name=f"cylinder{i+1}",
            axis=[0, 1, 0],
            center=[0.7, 0.5, 0],
            outer_radius=outer_radius,
            height=1.0,
        )
        for i, outer_radius in enumerate([1.1, 2.2, 3.3, 4.5])
    ]
    cylinder5 = fl.Cylinder(
        name="cylinder5",
        axis=[-1, 0, 0],
        center=[6.5, 0.5, 0],
        outer_radius=6.5,
        height=10,
    )

    farfield = fl.AutomatedFarfield(name="farfield", method="quasi-3d")

with fl.SI_unit_system:
    meshing_params = fl.SimulationParams(
        meshing=fl.MeshingParams(
            defaults=fl.MeshingDefaults(
                surface_edge_growth_rate=1.17,
                surface_max_edge_length=1.1,
                curvature_resolution_angle=12 * fl.u.deg,
                boundary_layer_growth_rate=1.17,
                boundary_layer_first_layer_thickness=1.8487111e-06,
            ),
            refinement_factor=1.35,
            gap_treatment_strength=0.5,
            volume_zones=[farfield],
            refinements=[
                fl.UniformRefinement(
                    name="refinement1", spacing=0.1, entities=[cylinders[0]]
                ),
                fl.UniformRefinement(
                    name="refinement2", spacing=0.15, entities=[cylinders[1]]
                ),
                fl.UniformRefinement(
                    name="refinement3", spacing=0.225, entities=[cylinders[2]]
                ),
                fl.UniformRefinement(
                    name="refinement4", spacing=0.275, entities=[cylinders[3]]
                ),
                fl.UniformRefinement(
                    name="refinement5", spacing=0.325, entities=[cylinder5]
                ),
                fl.SurfaceRefinement(
                    name="wing", max_edge_length=0.74, faces=[geometry["wing"]]
                ),
                fl.SurfaceRefinement(
                    name="flap-slat",
                    max_edge_length=0.55,
                    faces=[geometry["flap"], geometry["slat"]],
                ),
                fl.SurfaceRefinement(
                    name="trailing",
                    max_edge_length=0.36,
                    faces=[
                        geometry["*Trailing"],
                    ],
                ),
                fl.SurfaceEdgeRefinement(
                    name="edges",
                    method=fl.HeightBasedRefinement(value=0.0007),
                    edges=[
                        geometry["*trailingEdge"],
                        geometry["*leadingEdge"],
                    ],
                ),
                fl.SurfaceEdgeRefinement(
                    name="symmetry",
                    method=fl.ProjectAnisoSpacing(),
                    edges=[geometry["symmetry"]],
                ),
            ],
        ),
    )

with fl.SI_unit_system:
    reference_geometry_params = fl.SimulationParams(
        reference_geometry=fl.ReferenceGeometry(
            moment_center=[0.25, 0, 0], moment_length=[1, 1, 1], area=0.01
        )
    )

with fl.SI_unit_system:
    operating_condition_params = fl.SimulationParams(
        operating_condition=fl.AerospaceCondition.from_mach_reynolds(
            mach=0.2,
            reynolds_mesh_unit=5e6,
            temperature=272.1,
            alpha=16 * fl.u.deg,
            beta=0 * fl.u.deg,
            project_length_unit=1 * fl.u.m,
        ),
    )

with fl.SI_unit_system:
    time_stepping_params = fl.SimulationParams(
        time_stepping=fl.Steady(max_steps=3000, CFL=fl.AdaptiveCFL()),
    )

with fl.SI_unit_system:
    models_params = fl.SimulationParams(
        models=[
            fl.Wall(
                surfaces=[
                    geometry["*"],
                ],
                name="wall",
            ),
            fl.Freestream(surfaces=farfield.farfield, name="Freestream"),
            fl.SlipWall(surfaces=farfield.symmetry_planes, name="slipwall"),
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    linear_solver=fl.LinearSolver(max_iterations=35), kappa_MUSCL=0.33
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    linear_solver=fl.LinearSolver(max_iterations=25),
                    equation_evaluation_frequency=1,
                ),
            ),
        ],
    )

volume_outputs = [
    "primitiveVars",
    "vorticity",
    "residualNavierStokes",
    "residualTurbulence",
    "Cp",
    "Mach",
    "qcriterion",
    "mut",
]
surface_outputs = [
    "primitiveVars",
    "Cp",
    "Cf",
    "CfVec",
    "yPlus",
]

with fl.SI_unit_system:
    outputs_params = fl.SimulationParams(
        outputs=[
            fl.VolumeOutput(
                output_format="tecplot",
                name="fl.VolumeOutput",
                output_fields=volume_outputs,
            ),
            fl.SurfaceOutput(
                output_format="tecplot",
                name="fl.SurfaceOutput",
                surfaces=geometry["*"],
                output_fields=surface_outputs,
            ),
            fl.SliceOutput(
                output_format="tecplot",
                output_fields=volume_outputs,
                slices=[
                    fl.Slice(
                        name="Y=0.005m slice", normal=[0, 1, 0], origin=[0, 0.005, 0]
                    )
                ],
            ),
        ],
    )

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing_params.meshing,
        reference_geometry=reference_geometry_params.reference_geometry,
        operating_condition=operating_condition_params.operating_condition,
        time_stepping=time_stepping_params.time_stepping,
        models=models_params.models,
        outputs=outputs_params.outputs,
    )

case = project.run_case(params=params, name="Case of tutorial 2D CRM from Python")

case.wait()

total_forces = case.results.total_forces
total_forces = total_forces.as_dataframe()

nonlinear_residuals = case.results.nonlinear_residuals
nonlinear_residuals = nonlinear_residuals.as_dataframe()

cfl = case.results.cfl
cfl = cfl.as_dataframe()

x_slicing_force_distribution = case.results.x_slicing_force_distribution
x_slicing_force_distribution.wait()
x_slicing_force_distribution = x_slicing_force_distribution.as_dataframe()

print(cfl)

ax = nonlinear_residuals.plot(
    x="pseudo_step",
    y=["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"],
    logy=True,
    xlabel="Pseudo Step",
    ylabel="residuals",
    figsize=(10, 5),
)
ax.grid(True)

ax = cfl.plot(
    x="pseudo_step",
    y=["0_NavierStokes_cfl", "1_SpalartAllmaras_cfl"],
    xlim=(0, 3000),
    figsize=(10, 5),
)
ax.grid(True)

ax = total_forces.plot(
    x="pseudo_step",
    y="CL",
    xlabel="Pseudo Step",
    ylabel="CL",
    figsize=(10, 5),
    style="orange",
)
ax.grid(True)

ax = total_forces.plot(
    x="pseudo_step",
    y="CD",
    xlabel="Pseudo Step",
    ylabel="CD",
    figsize=(10, 5),
    style="lightgreen",
)
ax.grid(True)

ax = x_slicing_force_distribution.plot(
    x="X", y="totalCumulative_CD_Curve", xlabel="X", ylabel="CD", figsize=(10, 5)
)
ax.grid(True)