import math
import flow360 as fl
from flow360.examples import TutorialDynamicDerivatives

TutorialDynamicDerivatives.get_files()

project = fl.Project.from_geometry(
    TutorialDynamicDerivatives.geometry,
    name="Tutorial Calculating Dynamic Derivatives using Sliding Interfaces from Python",
)
geometry = project.geometry
geometry.group_faces_by_tag("faceName")

with fl.SI_unit_system:
    cylinder = fl.Cylinder(
        name="cylinder",
        axis=[0, 1, 0],
        center=[0, 0, 0],
        inner_radius=0,
        outer_radius=1.0,
        height=2.5,
    )
    sliding_interface = fl.RotationVolume(
        spacing_axial=0.04,
        spacing_radial=0.04,
        spacing_circumferential=0.04,
        entities=cylinder,
        enclosed_entities=geometry["wing"],
    )
    farfield = fl.AutomatedFarfield(name="farfield")

with fl.SI_unit_system:
    meshing_params = fl.MeshingParams(
        defaults=fl.MeshingDefaults(
            surface_max_edge_length=0.03 * fl.u.m,
            curvature_resolution_angle=8 * fl.u.deg,
            surface_edge_growth_rate=1.15,
            boundary_layer_first_layer_thickness=1e-6,
            boundary_layer_growth_rate=1.15,
        ),
        refinement_factor=1.0,
        volume_zones=[sliding_interface, farfield],
        refinements=[
            fl.SurfaceEdgeRefinement(
                name="leadingEdge",
                method=fl.AngleBasedRefinement(value=1 * fl.u.degree),
                edges=geometry["leadingEdge"],
            ),
            fl.SurfaceEdgeRefinement(
                name="trailingEdge",
                method=fl.HeightBasedRefinement(value=0.001),
                edges=geometry["trailingEdge"],
            ),
        ],
    )

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing_params,
        reference_geometry=fl.ReferenceGeometry(
            moment_center=[0, 0, 0],
            moment_length=[1, 1, 1],
            area=2,
        ),
        operating_condition=fl.AerospaceCondition(
            velocity_magnitude=50,
        ),
        time_stepping=fl.Steady(
            max_steps=10000,
            CFL=fl.RampCFL(initial=1, final=100, ramp_steps=1000),
        ),
        outputs=[
            fl.VolumeOutput(name="VolumeOutput", output_fields=["Mach"]),
            fl.SurfaceOutput(
                name="SurfaceOutput",
                surfaces=geometry["*"],
                output_fields=["Cp", "CfVec"],
            ),
        ],
        models=[
            fl.Rotation(
                volumes=cylinder,
                spec=fl.AngularVelocity(0 * fl.u.rad / fl.u.s),
            ),
            fl.Freestream(surfaces=farfield.farfield, name="Freestream"),
            fl.Wall(surfaces=geometry["wing"], name="NoSlipWall"),
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-9,
                    linear_solver=fl.LinearSolver(max_iterations=35),
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    absolute_tolerance=1e-8,
                    linear_solver=fl.LinearSolver(max_iterations=25),
                ),
            ),
        ],
    )

project.run_case(
    params=params, name="Tutorial Dynamic Derivatives - Steady Initialization"
)

steady_case = fl.Case.from_cloud(project.case.id)

with fl.SI_unit_system:
    params.time_stepping = fl.Unsteady(
        max_pseudo_steps=80,
        steps=400,
        step_size=0.01 * 2.0 * math.pi / 20.0 * fl.u.s,
        CFL=fl.RampCFL(initial=1, final=1e8, ramp_steps=20),
    )
    params.models[0].spec = fl.AngleExpression("0.0349066 * sin(0.05877271 * t)")

project.run_case(
    params=params,
    name="Tutorial Dynamic Derivatives - Unsteady Oscillation",
    fork_from=steady_case,
)
unsteady_case = project.case

import flow360 as fl
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

cmap = plt.colormaps["tab10"]

project = fl.Project.from_example(by_name="Dynamic Derivatives")
case = project.get_case()

case.results.total_forces.download(to_file="total_forces.csv")
case.results.total_forces.load_from_local("total_forces.csv")
total_forces = case.results.total_forces.as_dataframe()

def extract_unique_physical_steps(input_df):
    """Extract the row at the last pseudo step at the end of each physical step."""
    physical_steps_raw = input_df["physical_step"]
    n_rows = len(physical_steps_raw)
    selected_row_index = []
    for i in range(n_rows - 1):
        if physical_steps_raw.iloc[i + 1] > physical_steps_raw.iloc[i]:
            selected_row_index.append(i)
    selected_row_index.append(n_rows - 1)
    output = pd.DataFrame()
    for k, v in input_df.items():
        output[k] = [v.iloc[j] for j in selected_row_index]
    return output.reset_index(drop=True)

total_forces = extract_unique_physical_steps(total_forces)
print(f"Number of physical steps: {len(total_forces)}")

length_unit = case.params.private_attribute_asset_cache.project_length_unit  # m
U_inf = case.params.operating_condition.velocity_magnitude  # m/s
C_inf = 340.29400580821283 * fl.u.m / fl.u.s  # speed of sound
n_chord_per_period = 2.5  # flowthru/rad
chord = 1 * fl.u.m

omega_physical = U_inf / (n_chord_per_period * chord)  # 1/s
omega_nondim = omega_physical / (C_inf / length_unit)
time_step_size_nondim = 0.01 * 2 * np.pi / omega_nondim  # nondim

amplitude_deg = 2 * fl.u.deg
amplitude_rad = amplitude_deg.to(fl.u.rad)

def compute_time_nondim(df):
    """Compute nondimensional time from physical step."""
    n_steps = len(df["physical_step"])
    df["time_nondim"] = np.zeros(n_steps)
    for i in range(n_steps):
        df.loc[i, "time_nondim"] = time_step_size_nondim * (
            df.loc[i, "physical_step"] + 1
        )

def compute_alpha_dot(df):
    """Compute angle of attack and pitch rate."""
    n_steps = len(df["physical_step"])
    df["alpha"] = np.zeros(n_steps)
    df["alpha_dot"] = np.zeros(n_steps)
    alpha_eq_0_index = []

    for i in range(n_steps):
        t = df.loc[i, "time_nondim"]
        df.loc[i, "alpha"] = amplitude_rad * np.sin(omega_nondim * t)

    for i in range(1, n_steps - 1):
        df.loc[i, "alpha_dot"] = (
            (df.loc[i + 1, "alpha"] - df.loc[i - 1, "alpha"])
            / (2 * time_step_size_nondim)
            * C_inf
            / length_unit
            * chord
            / U_inf
        )
        if abs(df.loc[i, "alpha"]) < 1e-10:
            alpha_eq_0_index.append(i)

    return alpha_eq_0_index

compute_time_nondim(total_forces)
alpha_eq_0_index = compute_alpha_dot(total_forces)
print(f"Found {len(alpha_eq_0_index)} points where alpha = 0")

def plot_time_history(df, alpha_eq_0_idx):
    _, axs = plt.subplots(3, figsize=(10, 8))

    axs[0].plot(df["physical_step"], df["CMy"], color=cmap.colors[0])
    axs[0].set_yticks(np.linspace(-0.02, 0.02, 5))
    axs[0].set_ylim([-0.02, 0.02])
    axs[0].set_ylabel("CMy")

    axs[1].plot(df["physical_step"], df["alpha"], color=cmap.colors[1])
    axs[1].set_ylabel(r"$\alpha$ (rad)")

    axs[2].plot(df["physical_step"][1:-2], df["alpha_dot"][1:-2], color=cmap.colors[2])
    axs[2].set_ylabel(r"$\dot{\alpha}$ (rad/flowthru)")
    axs[2].set_xlabel("physical_step")

    for ax in axs:
        ax.label_outer()
        ax.grid()

    for i in alpha_eq_0_idx:
        color = "b" if df["alpha_dot"][i] < 0 else "r"
        axs[0].plot(df["physical_step"][i], df["CMy"][i], "o", color=color)
        axs[1].plot(df["physical_step"][i], df["alpha"][i], "o", color=color)
        axs[2].plot(df["physical_step"][i], df["alpha_dot"][i], "o", color=color)

    axs[0].set_title(
        r"Time History of CMy, $\alpha$ and $\dot{\alpha}$"
        + "\n"
        + f"flowthru/rad={n_chord_per_period:.1f}, amplitude={amplitude_deg:.1f}"
    )
    plt.tight_layout()
    plt.show()

plot_time_history(total_forces, alpha_eq_0_index)

def plot_CMy_vs_q(df, alpha_eq_0_idx):
    plt.figure(figsize=(8, 6))

    n_steps = len(df["physical_step"])
    begin = n_steps - 201
    end = n_steps - 1
    alpha_dot = df.loc[begin:end, "alpha_dot"]
    CMy = df.loc[begin:end, "CMy"]
    plt.plot(alpha_dot, CMy, ".")

    x = []
    y = []
    for i in alpha_eq_0_idx:
        if i >= begin and i <= end:
            x.append(df.loc[i, "alpha_dot"])
            y.append(df.loc[i, "CMy"])

    plt.plot(x[0], y[0], "o", color="r", markersize=10)
    plt.plot(x[1], y[1], "o", color="b", markersize=10)

    coef = np.polyfit(x, y, 1)
    func = np.poly1d(coef)
    x_line = np.linspace(-0.02, 0.02, 101)
    plt.plot(x_line, func(x_line), "--", label=f"slope = {coef[0]:.3f}", color="k")

    plt.xlabel(r"$\dot{\alpha} = q$ (deg/flowthru)")
    plt.ylabel("CMy")
    plt.ylim([-0.02, 0.02])
    plt.yticks(np.linspace(-0.02, 0.02, 5))
    plt.grid()
    plt.legend(loc="lower left")
    plt.title(
        f"Dynamic Derivative dCMy/dq\nflowthru/rad={n_chord_per_period:.1f}, amplitude={amplitude_deg:.1f}"
    )

    tick_locations = np.linspace(-1, 1, 5) / 180 * np.pi
    ax = plt.gca()
    ax.set_xticks(
        tick_locations, [f"{loc * 180 / np.pi:.1f}" for loc in tick_locations]
    )

    ax.text(
        -0.016,
        0.01,
        r"$\alpha = 0$" + "\n" + r"min $\dot{\alpha}$",
        fontsize=14,
        color="b",
    )
    ax.text(
        +0.012,
        -0.014,
        r"$\alpha = 0$" + "\n" + r"max $\dot{\alpha}$",
        fontsize=14,
        color="r",
    )

    plt.tight_layout()
    plt.show()

    print(f"\nDynamic derivative: dCMy/dq = {coef[0]:.3f}")

plot_CMy_vs_q(total_forces, alpha_eq_0_index)