
import flow360 as fl
from flow360.examples.DARPA import DARPA_SUBOFF

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import trapezoid

DARPA_SUBOFF.get_files()
project = fl.Project.from_geometry(
    DARPA_SUBOFF.geometry,
    name="Tutorial DARPA Actuator Disk from Python",
)
geo = project.geometry
geo.group_faces_by_tag("faceId")  # group faces by faceId

R_out = 0.3  # Outer radius of the actuator disk
R_hub = 0.0597 / 2  # Blade root radius
R_in = (
    R_hub * 0.98
)  # Inner radius of the actuator disk, smaller than the blade root radius to prevent hub vortex affecting the surface flow
NBlades = 5  # Number of blades

thrust = 2000  # Thrust force
torque = 100  # Torque force
omega = 100 * 2 * np.pi / 60  # Rotation rate
rho = 1000  # Density
Uinf = 5  # Freestream velocity

r_from_data = [R_in, R_out * 0.7, R_out]
thrust_max = (2 * thrust) / (R_out - R_in)
torque_max = (2 * torque) / (R_out - R_in)
thrust_radial = [0, thrust_max, 0]  # Triangular distribution
torque_radial = [0, torque_max, 0]  # Triangular distribution

r_from_data = np.array(r_from_data)
thrust_radial_from_data = np.array(thrust_radial)
torque_radial_from_data = np.array(torque_radial)
f_ax_from_data = thrust_radial_from_data / (2 * np.pi * r_from_data)
f_cir_from_data = torque_radial_from_data / (2 * np.pi * r_from_data**2)

n_points = 200
r_theory = np.linspace(R_in, R_out, n_points)
xi = r_theory / R_out

v_ax = 0.5 * (np.sqrt((2 * thrust) / (rho * np.pi * R_out**2) + Uinf**2) + Uinf)

lambda_w = v_ax / (omega * R_out)

f_tip = (NBlades / 2) * (1 - xi) * np.sqrt(1 + 1 / lambda_w**2)
f_hub = (NBlades / 2) * (xi - R_in / R_out) * np.sqrt(1 + 1 / lambda_w**2)
F_tip = (2 / np.pi) * np.arccos(np.exp(-f_tip))
F_hub = (2 / np.pi) * np.arccos(np.exp(-f_hub))
F = F_tip * F_hub

gamma_tmp = (F * xi**2) / (lambda_w**2 + xi**2)

f_ax = NBlades * (rho * omega * gamma_tmp) / (2 * np.pi)
f_cir = NBlades * (rho * v_ax * gamma_tmp) / (2 * np.pi * r_theory)

integrand_thrust = f_ax * 2 * np.pi * r_theory
integrand_torque = f_cir * r_theory * 2 * np.pi * r_theory
thrust_tmp = trapezoid(integrand_thrust, r_theory)
torque_tmp = trapezoid(integrand_torque, r_theory)
K_thrust = thrust / thrust_tmp
K_torque = torque / torque_tmp
f_ax_from_theory = K_thrust * f_ax
f_cir_from_theory = K_torque * f_cir

data_is_given = False
if data_is_given:
    r = r_from_data
    f_ax = f_ax_from_data
    f_cir = f_cir_from_data
else:
    r = r_theory
    f_ax = f_ax_from_theory
    f_cir = f_cir_from_theory

thrust_check = trapezoid(
    f_ax * 2 * np.pi * r, r
)  # Pressure force integrated over the radius
torque_check = trapezoid(
    f_cir * r * 2 * np.pi * r, r
)  # Pressure force integrated over the radius

print(f"Desired thrust: {thrust}, computed thrust: {thrust_check}")
print(f"Desired torque: {torque}, computed torque: {torque_check}")

with fl.SI_unit_system:
    x_AD = 4.2609  # x-location of the actuator disk center
    AD_zone = fl.Cylinder(
        name="AD",
        axis=(1, 0, 0),
        center=(x_AD, 0, 0),
        inner_radius=R_in,
        outer_radius=R_out,
        height=R_out * 0.1,
    )

    AD_model = fl.ActuatorDisk(
        name="AD_model",
        volumes=AD_zone,
        force_per_area=fl.ForcePerArea(
            radius=r,
            thrust=f_ax,
            circumferential=f_cir,
        ),
    )

d_axial = R_out * 0.1 / 30  # 30 nodes in the axial direction
d_radial = d_axial * 2  # x2 spatial distance in the radial direction
d_circumferential = d_axial * 2  # x2 spatial distance in the circumferential direction

first_layer_thickness = 1e-5
growth_rate = 1.2
surface_elements_length = 0.001
n_layers = 20

total_height = first_layer_thickness * growth_rate**n_layers

inner_radius = R_hub + total_height

with fl.SI_unit_system:
    refinements = [
        fl.AxisymmetricRefinement(
            name="AD_axisym",
            spacing_axial=d_axial,
            spacing_radial=d_radial,
            spacing_circumferential=d_circumferential,
            entities=[
                fl.Cylinder(
                    name="AD_axisym",
                    axis=(1, 0, 0),
                    center=(x_AD, 0, 0),
                    inner_radius=inner_radius,
                    outer_radius=R_out,
                    height=R_out * 0.1,
                ),
            ],
        )
    ]

with fl.SI_unit_system:
    faces_close_to_AD = [
        geo[f"body00001_face000{n:02d}"] for n in [2, 4, 19, 21]
    ]  # Surface faces close to axisymmetric refinement
    surface_refinement = fl.SurfaceRefinement(
        name="surface_refinement",
        max_edge_length=1 * fl.u.mm,
        faces=faces_close_to_AD,
    )

    height = R_out * 3  # Height of the volume refinement
    front_space = (
        R_out * 0.5
    )  # Space in front of the actuator disk to the volume refinement
    center_volume = x_AD + (height / 2 - front_space)
    volume_refinement = fl.UniformRefinement(
        name="volume_refinement",
        spacing=10 * fl.u.mm,
        entities=[
            fl.Cylinder(
                name="volume_refinement",
                axis=(1, 0, 0),
                center=(center_volume, 0, 0),
                outer_radius=R_out * 1.5,
                height=height,
            )
        ],
    )

    L = 4.3562
    height_general = L * 2
    front_space_general = 0.5
    center_general = height_general / 2 - front_space_general
    general_refinement = fl.UniformRefinement(
        name="general_refinement",
        spacing=20 * fl.u.mm,
        entities=[
            fl.Cylinder(
                name="general_refinement",
                axis=(1, 0, 0),
                center=(center_general, 0, 0),
                outer_radius=R_out * 2,
                height=height_general,
            )
        ],
    )
    refinements += [
        surface_refinement,
        volume_refinement,
        general_refinement,
    ]

with fl.SI_unit_system:
    operating_condition = fl.LiquidOperatingCondition(
        velocity_magnitude=5, alpha=0 * fl.u.deg
    )

with fl.SI_unit_system:
    farfield = fl.AutomatedFarfield()
    params = fl.SimulationParams(
        meshing=fl.MeshingParams(
            defaults=fl.MeshingDefaults(
                boundary_layer_first_layer_thickness=first_layer_thickness,
                surface_max_edge_length=20 * fl.u.mm,
            ),
            refinements=refinements,  # All refinements previously defined
            volume_zones=[farfield],  # Farfield volume zone
        ),
        operating_condition=operating_condition,  # Liquid operating condition
        time_stepping=fl.Steady(max_steps=5000),
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(),
                turbulence_model_solver=fl.SpalartAllmaras(),
            ),
            fl.Freestream(surfaces=farfield.farfield),
            fl.Wall(surfaces=geo["*"]),
            AD_model,  # Actuator disk model
        ],
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(0, 0, 0),
            moment_length=(3.81, 3.81, 3.81),
            area=45.604,
        ),
        outputs=[
            fl.VolumeOutput(
                output_fields=[
                    "primitiveVars",
                    "Cp",
                    "qcriterion",
                ],
            ),
            fl.SurfaceOutput(
                surfaces=geo["*"],
                output_fields=[
                    "primitiveVars",
                    "Cp",
                    "Cf",
                    "CfVec",
                    "yPlus",
                ],
            ),
        ],
    )

case = project.run_case(
    params=params,
    name="Tutorial DARPA Actuator Disk from Python",
)

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

length_unit = project.length_unit
density = case.params.operating_condition.material.density
U_ref = (
    case.params.operating_condition.reference_velocity_magnitude
    or case.params.operating_condition.velocity_magnitude
)
A_ref = case.params.reference_geometry.area

body_forces = case.results.total_forces.as_dataframe().iloc[-1]
CD = body_forces["CD"]

actuator_disk_output = case.results.actuator_disks.as_dataframe().iloc[-1]
C_p = actuator_disk_output["Disk0_Power"]
C_t = actuator_disk_output["Disk0_Force"]
C_q = actuator_disk_output["Disk0_Moment"]

drag_force = (CD * 0.5 * density * U_ref**2 * A_ref).to("N")
power = (C_p * density * U_ref**3 * length_unit**2).to("W")
thrust = (C_t * density * U_ref**2 * length_unit**2).to("N")
torque = (C_q * density * U_ref**2 * length_unit**3).to("N*m")

print(f"Drag force: {drag_force}")
print(f"Power: {power}")
print(f"Thrust: {thrust}")
print(f"Torque: {torque}")
