import flow360 as fl
from flow360.examples import Airplane

import pandas as pd
import matplotlib.pyplot as plt

Airplane.get_files()

project = fl.Project.from_geometry(
    Airplane.geometry,
    name="Warm-Started Alpha Sweep",
)

geometry = project.geometry

with fl.SI_unit_system:
    far_field_zone = fl.AutomatedFarfield()

    meshing_params = fl.MeshingParams(
        defaults=fl.MeshingDefaults(
            boundary_layer_first_layer_thickness=0.01 * fl.u.mm,
            surface_max_edge_length=0.2,
        ),
        volume_zones=[far_field_zone],
    )

with fl.SI_unit_system:
    walls =fl.Wall(surfaces=[geometry["*"]])
    cl_output = fl.ForceOutput(
        output_fields=["CL"],
        models=[walls]
    )
    params = fl.SimulationParams(
        meshing=meshing_params,
        reference_geometry=fl.ReferenceGeometry(),
        operating_condition=fl.AerospaceCondition(
            velocity_magnitude=100,
            alpha=0 * fl.u.deg,
        ),
        time_stepping=fl.Steady(max_steps=4000),
        models=[
            walls,
            fl.Freestream(surfaces=[far_field_zone.farfield]),
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-5
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    absolute_tolerance=1e-5
                )
            )
        ],
        outputs=[cl_output],
        run_control=fl.RunControl(
            stopping_criteria=[
                fl.StoppingCriterion(
                    monitor_field="CL",
                    monitor_output=cl_output,
                    tolerance=0.05,
                    tolerance_window_size=200
                )
            ]
        )
    )

alphas = [-4, 0, 4, 8, 12] * fl.u.deg

case_list = []
previous_case = None

for alpha_angle in alphas:
    params.operating_condition.alpha = alpha_angle

    if alpha_angle == alphas[0]:
        case = project.run_case(
            params=params,
            name=f"alpha_{alpha_angle.value}",
            use_beta_mesher=True,
        )
    else:
        case = project.run_case(
            params=params,
            name=f"alpha_{alpha_angle.value}",
            fork_from=previous_case,
            use_beta_mesher=True,
        )

    previous_case = case

    print(f"Submitted {case.name} (id: {case.id})")
    case_list.append((alpha_angle, case))

print("Waiting for all cases to finish...")
for _, case in case_list:
    case.wait()
print("All cases completed.")

records = []
for alpha_angle, case in case_list:
    averages = case.results.total_forces.get_averages(0.1)
    records.append(
        {
            "alpha_deg": float(alpha_angle.to("deg").value),
            "CL": averages["CL"],
            "CD": averages["CD"],
        }
    )

sweep = pd.DataFrame(records).sort_values("alpha_deg").reset_index(drop=True)
print(sweep)

ax = sweep.plot(
    x="alpha_deg",
    y="CL",
    marker="o",
    legend=False,
    xlabel="AoA [deg]",
    ylabel="CL",
    title="Lift coefficient vs angle of attack",
    figsize=(8, 4),
)
ax.grid(True)
plt.show()

ax = sweep.plot(
    x="alpha_deg",
    y="CD",
    marker="o",
    legend=False,
    xlabel="AoA [deg]",
    ylabel="CD",
    title="Drag coefficient vs angle of attack",
    figsize=(8, 4),
)
ax.grid(True)
plt.show()

ax = sweep.plot(
    x="CD",
    y="CL",
    marker="o",
    legend=False,
    xlabel="CD",
    ylabel="CL",
    title="Drag polar (CL vs CD)",
    figsize=(6, 6),
)
ax.grid(True)
plt.show()