import math

import flow360 as fl
from flow360.examples import IsolatedPropellerNacelle

import matplotlib.pyplot as plt
import pandas as pd

IsolatedPropellerNacelle.get_files()
project = fl.Project.from_geometry(
    IsolatedPropellerNacelle.geometry,
    length_unit="m",
    name="Propeller with nacelle, rotating zone with GeometryAI",
)
geometry = project.geometry

draft = fl.create_draft(new_run_from=geometry, face_grouping="faceName", edge_grouping="edgeId")

with draft:
    all_surfaces = draft.surfaces["*"]
    blades = [draft.surfaces[f"blade{i}"] for i in range(1, 6)]
    spinner = draft.surfaces["spinner"]
    nacelle = draft.surfaces["nacelle"]

print(sorted(surface.name for surface in all_surfaces))

with fl.SI_unit_system:
    rotor_cylinder = fl.Cylinder(
        name="rotor",
        axis=(1, 0, 0),
        center=(-0.025, 0, 0),
        outer_radius=1.6,
        height=0.55,
    )

    rotation_zone = fl.RotationVolume(
        name="rotor_zone",
        entities=rotor_cylinder,
        spacing_axial=0.0625,
        spacing_radial=0.0625,
        spacing_circumferential=0.0625,
        enclosed_entities=blades + [spinner, nacelle],
        stationary_enclosed_entities=[nacelle],
    )

    rotor_refinement = fl.UniformRefinement(
        name="rotor_refinement",
        spacing=0.03125,
        entities=fl.Cylinder(
            name="rotor_refinement",
            axis=(1, 0, 0),
            center=(-0.025, 0, 0),
            outer_radius=1.58,
            height=0.53,
        ),
    )

farfield = fl.AutomatedFarfield()

with fl.SI_unit_system:
    meshing = fl.MeshingParams(
        defaults=fl.MeshingDefaults(
            geometry_accuracy=1 * fl.u.mm,
            curvature_resolution_angle=15 * fl.u.deg,
            surface_max_edge_length=0.0625,
            boundary_layer_first_layer_thickness=1e-5,
            boundary_layer_growth_rate=1.25,
        ),
        refinements=[
            rotor_refinement,
            fl.SurfaceRefinement(name="blades", max_edge_length=0.03, faces=blades),
        ],
        volume_zones=[farfield, rotation_zone],
    )

rpm = 500
velocity = 20
diameter = 3.0

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing,
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(0, 0, 0),
            moment_length=(diameter / 2, diameter / 2, diameter / 2),
            area=math.pi * (diameter / 2) ** 2,
        ),
        operating_condition=fl.AerospaceCondition(velocity_magnitude=velocity),
        time_stepping=fl.Steady(max_steps=1000),
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(),
                turbulence_model_solver=fl.SpalartAllmaras(),
            ),
            fl.Wall(name="propeller_and_nacelle", surfaces=all_surfaces),
            fl.Freestream(surfaces=farfield.farfield),
            fl.Rotation(
                name="propeller_rotation",
                volumes=rotor_cylinder,
                spec=fl.AngularVelocity(rpm * fl.u.rpm),
            ),
        ],
        outputs=[
            fl.SurfaceOutput(
                surfaces=all_surfaces,
                output_fields=["Cp", "Cf", "CfVec", "yPlus", "velocity"],
            ),
            fl.SliceOutput(
                slices=[fl.Slice(name="y0", normal=(0, 1, 0), origin=(0, 0, 0))],
                output_fields=["velocity", "Cp", "Mach"],
            ),
        ],
    )

with draft:
    case = project.run_case(
        params=params,
        name="Propeller with nacelle, J = 0.8",
        use_geometry_AI=True,
        use_beta_mesher=True,
    )

case.wait()

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

surface_forces = case.results.surface_forces.as_dataframe().iloc[-1]

density = case.params.operating_condition.thermal_state.density
A_ref = case.params.reference_geometry.area
L_ref = case.params.reference_geometry.moment_length[0]
q_inf = 0.5 * density * (velocity * fl.u.m / fl.u.s) ** 2

boundaries = sorted(column.removesuffix("_CFx") for column in surface_forces.index if column.endswith("_CFx"))
loads = pd.DataFrame(
    {
        "Fx [N]": [float((surface_forces[f"{b}_CFx"] * q_inf * A_ref).to("N")) for b in boundaries],
        "Mx [N*m]": [float((surface_forces[f"{b}_CMx"] * q_inf * A_ref * L_ref).to("N*m")) for b in boundaries],
    },
    index=boundaries,
)
loads.round(2)

propeller = [b for b in boundaries if "nacelle" not in b]
thrust = -loads.loc[propeller, "Fx [N]"].sum()
torque = -loads.loc[propeller, "Mx [N*m]"].sum()

n = rpm / 60
rho = float(density.to("kg/m**3"))
J = velocity / (n * diameter)
C_T = thrust / (rho * n**2 * diameter**4)
C_Q = torque / (rho * n**2 * diameter**5)
C_P = 2 * math.pi * C_Q

print(f"Advance ratio J = {J:.2f}")
print(f"Thrust        T = {thrust:8.1f} N    C_T = {C_T:.4f}")
print(f"Torque        Q = {torque:8.1f} N*m  C_Q = {C_Q:.5f}")
print(f"Power         P = {2 * math.pi * n * torque / 1000:8.2f} kW   C_P = {C_P:.4f}")
print(f"Efficiency    eta = {J * C_T / C_P:.2f}")