import numpy as np
import matplotlib.pyplot as plt
import flow360 as fl
from flow360.examples import DTU_WindTurbine

DTU_WindTurbine.get_files()

project = fl.Project.from_geometry(
    DTU_WindTurbine.geometry,
    length_unit="m",
    name="Time-Accurate Rotor with Thrust-Convergence Stopping Criterion",
)
geometry = project.geometry

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

    R = 89.166  # blade radius [m] — DTU 10MW reference turbine

    rotor_cylinder = fl.Cylinder(
        center=(0, 0, 0),
        outer_radius=110 * fl.u.m,
        height=40 * fl.u.m,
        axis=(0, 0, 1),
        name="RotorCylinder",
    )

    wake_cylinder = fl.Cylinder(
        center=(0, 0, 55) * fl.u.m,
        outer_radius=200 * fl.u.m,
        height=200 * fl.u.m,
        axis=(0, 0, 1),
        name="WakeCylinder",
    )

rotating_zone = fl.RotationVolume(
    name="RotatingZone",
    spacing_axial=2 * fl.u.m,
    spacing_circumferential=2 * fl.u.m,
    spacing_radial=2 * fl.u.m,
    entities=rotor_cylinder,
    enclosed_entities=geometry["*"],
)

meshing = fl.MeshingParams(
    defaults=fl.MeshingDefaults(
        surface_max_edge_length=2 * fl.u.m,
        boundary_layer_first_layer_thickness=1e-4 * fl.u.m,
    ),
    volume_zones=[farfield, rotating_zone],
    refinements=[
        fl.UniformRefinement(entities=wake_cylinder, spacing=4 * fl.u.m),
        fl.UniformRefinement(entities=rotor_cylinder, spacing=2 * fl.u.m),
    ],
)

omega = 8.836 * fl.u.rpm  # rated angular velocity — DTU 10MW
dt_deg = 6.0  # physical time step [degrees]; change as needed
revs = 50  # maximum number of revolutions to simulate
steps_per_rev = int(360 / dt_deg)
step_size = (dt_deg * fl.u.deg / omega.to("deg/s")).to("s")

time_stepping = fl.Unsteady(
    steps=revs * steps_per_rev,
    step_size=step_size,
    max_pseudo_steps=50,
)

rotation = fl.Rotation(
    name="RotatingZone",
    volumes=rotor_cylinder,
    spec=fl.AngularVelocity(omega),
)

revs_for_mean = 2
window_steps = revs_for_mean * steps_per_rev

walls = fl.Wall(surfaces=geometry["*"], use_wall_function=True)
thrust_output = fl.ForceOutput(
    name="thrust_monitor",
    output_fields=["CFz", "CMz"],
    moving_statistic=fl.MovingStatistic(
        method="mean",
        moving_window_size=window_steps,
    ),
    models=[walls],
)

revs_for_stop_criterion = 3
tolerance_window_size = revs_for_stop_criterion * steps_per_rev
run_control = fl.RunControl(
    stopping_criteria=[
        fl.StoppingCriterion(
            name="thrust_convergence",
            monitor_output=thrust_output,
            monitor_field="CFz",
            tolerance=1e-4,
            tolerance_window_size=tolerance_window_size,
        )
    ]
)

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing,
        operating_condition=fl.AerospaceCondition(
            velocity_magnitude=11.0,  # rated wind speed [m/s]
            alpha=90 * fl.u.deg,  # wind aligned with rotor axis (+Z)
            reference_velocity_magnitude=omega.to("rad/s").to_value()
            * R,  # ≈ tip speed [m/s] for non-dimensionalisation
        ),
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(0, 0, 0),
            moment_length=(R, R, R),
            area=np.pi * R**2,
        ),
        time_stepping=time_stepping,
        run_control=run_control,
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-10,
                    relative_tolerance=1e-1,
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    absolute_tolerance=1e-10,
                    relative_tolerance=1e-1,
                ),
            ),
            walls,
            fl.Freestream(surfaces=farfield.farfield),
            rotation,
        ],
        outputs=[
            fl.SurfaceOutput(
                surfaces=geometry["*"],
                output_fields=["Cp", "Cf", "CfVec"],
            ),
            thrust_output,
        ],
    )

case = project.run_case(
    params=params,
    name="Time-accurate (thrust convergence)",
    use_beta_mesher=True,
)

case.wait()

density = case.params.operating_condition.thermal_state.density
U_ref = case.params.operating_condition.velocity_magnitude
A_ref = case.params.reference_geometry.area
q = 0.5 * float(density) * float(U_ref) ** 2 * float(A_ref)

thrust_monitor_mean = case.results.custom_forces["thrust_monitor_moving_statistic"].as_dataframe()
thrust_monitor_mean["Thrust_mean [N]"] = thrust_monitor_mean["totalCFz_mean"]
thrust_monitor_mean_step = thrust_monitor_mean["end_index"]
thrust_monitor = case.results.custom_forces["thrust_monitor"].as_dataframe()
thrust_monitor["Thrust [N]"] = thrust_monitor["totalCFz"]
thrust_monitor_step = thrust_monitor["physical_step"]

fig, axes = plt.subplots(2, 1, figsize=(10, 8), sharex=False)

axes[0].plot(
    thrust_monitor["physical_step"],
    thrust_monitor["Thrust [N]"],
    color="steelblue",
    alpha=0.5,
    linewidth=0.8,
    label="Instantaneous",
)
axes[0].plot(
    thrust_monitor_mean["end_index"],
    thrust_monitor_mean["Thrust_mean [N]"],
    color="navy",
    linewidth=2,
    label=f"{revs_for_mean}-rev rolling mean",
)
axes[0].set_ylabel("Thrust [N]")
axes[0].legend()
axes[0].grid(True, alpha=0.3)
axes[0].set_title("Thrust convergence")

final_step = thrust_monitor["physical_step"].iloc[-1]
non_linear = case.results.nonlinear_residuals.as_dataframe()
nl = non_linear.copy()
max_ps = nl["pseudo_step"].max() + 1
nl["step_with_pseudo"] = nl["physical_step"] + nl["pseudo_step"] / max_ps
for col in ["0_cont", "1_momx", "2_momy", "3_momz", "4_energ", "5_nuHat"]:
    if col in nl.columns:
        axes[1].semilogy(nl["step_with_pseudo"], nl[col], label=col, alpha=0.7)
axes[1].set_xlim(final_step - steps_per_rev + 1, final_step + 1)
axes[1].set_xlabel("Physical step (with pseudo steps)")
axes[1].set_ylabel("Nonlinear residual")
axes[1].legend(fontsize=8)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print(
    f"Simulation stopped at physical step {final_step} "
    f"({final_step / steps_per_rev:.1f} revolutions)"
)