import flow360 as fl
from math import cos, sin, radians
import numpy as np
from flow360.examples import F1_2025

F1_2025.get_files()

project = fl.Project.from_geometry(
    F1_2025.geometry,
    name="F1 racecar",
    length_unit="m"
)

geometry = project.geometry

geometry.group_faces_for_snappy()

rad_lhs = fl.SeedpointVolume(name="rad_lhs",
                                    point_in_mesh=[2040.53452, -405.234, 395.34243]*fl.u.mm,
                                    axes=[[19.98, 6.62, 29.21],
                                        [48.076066, -210.05, 14.72]])

rad_rhs = fl.SeedpointVolume(name="rad_rhs",
                            point_in_mesh=[2040.53452, 405.234, 395.34243]*fl.u.mm,
                            axes=[[19.98, -6.62, 29.21],
                                [48.076066, 210.05, 14.72]])

fluid = fl.SeedpointVolume(point_in_mesh=[0.234, 0.18172, 2.1324]*fl.u.m, name="fluid")

surface_defaults = fl.snappy.SurfaceMeshingDefaults(
    min_spacing=8 * fl.u.mm,
    max_spacing=20 * fl.u.mm,
    gap_resolution=0.001 * fl.u.mm, # gaps of this length should be sealed by the mesher, can be specified per face
)  

spacing = fl.OctreeSpacing(base_spacing=1 * fl.u.mm)

smooth_controls=fl.snappy.SmoothControls(
    lambda_factor=0.7,
    mu_factor=0.71,
    iterations=5,
)

castellated_mesh_controls=fl.snappy.CastellatedMeshControls(
    resolve_feature_angle=18.0 * fl.u.deg,
)

with fl.SI_unit_system:
    plinth_pucks = fl.UniformRefinement(
        entities=[
            fl.Cylinder(
                name="fr-lhs-puck",
                center=[200, -800, 10] * fl.u.mm,
                axis=[0, 0, 1],
                outer_radius=150 * fl.u.mm,
                height=15 * fl.u.mm,
            ),
            fl.Cylinder(
                name="fr-rhs-puck",
                center=[200, 800, 10] * fl.u.mm,
                axis=[0, 0, 1],
                outer_radius=150 * fl.u.mm,
                height=15 * fl.u.mm,
            ),
            fl.Cylinder(
                name="rr-lhs-puck",
                center=[3650, -765, -10] * fl.u.mm,
                axis=[0, 0, 1],
                outer_radius=200 * fl.u.mm,
                height=15 * fl.u.mm,
            ),
            fl.Cylinder(
                name="rr-rhs-puck",
                center=[3650, 765, -10] * fl.u.mm,
                axis=[0, 0, 1],
                outer_radius=200 * fl.u.mm,
                height=15 * fl.u.mm,
            ),
        ],
        spacing=0.5 * fl.u.mm,
    )

with fl.SI_unit_system:
    surface_refinements = [
        fl.snappy.BodyRefinement(
            bodies=geometry.snappy_bodies["tunnel"],
            min_spacing=spacing[-11],
            max_spacing=spacing[-11],
        ),
        fl.snappy.RegionRefinement(
            regions=geometry.snappy_bodies["*"]["*plinth*"],
            min_spacing=spacing[1],
            max_spacing=spacing[-2],
        ),
        fl.snappy.SurfaceEdgeRefinement(
            entities=geometry.snappy_bodies["*"]["*plinth*"],
            spacing=[2 * fl.u.mm],
            distances=[5 * fl.u.mm],
        ),
        fl.snappy.SurfaceEdgeRefinement(
            entities=geometry.snappy_bodies["tunnel"],
            spacing=[spacing[-11]],
            distances=[2 * fl.u.mm],
            included_angle=140 * fl.u.deg,
        ),
        fl.snappy.SurfaceEdgeRefinement(
            entities=[
                geometry.snappy_bodies["toint-rad-main-*"],
                geometry.snappy_bodies["body-main-rad-duct"],
                geometry.snappy_bodies["*-int-brake-duct-*"],
            ],
            spacing=[1 * fl.u.mm],
            distances=[2 * fl.u.mm],
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["*wh-rim*"],
                geometry.snappy_bodies["body-helmet"],
                geometry.snappy_bodies["body-pit"],
            ],
            min_spacing=1 * fl.u.mm,
            max_spacing=5 * fl.u.m,
        ),
        fl.snappy.BodyRefinement(
            bodies=geometry.snappy_bodies["body-misc"],
            min_spacing=0.5 * fl.u.mm,
            max_spacing=3 * fl.u.mm,
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["body-camera"],
                geometry.snappy_bodies["body-aero-camera"],
                geometry.snappy_bodies["body-steering-wheel"],
                geometry.snappy_bodies["uf-strakes"],
                geometry.snappy_bodies["body-mirror"],
                geometry.snappy_bodies["rw-pillar"],
                geometry.snappy_bodies["rr-int-vanes-*"],
                geometry.snappy_bodies["body-front-af"],
                geometry.snappy_bodies["body-main-rad-duct"],
            ],
            min_spacing=2 * fl.u.mm,
            max_spacing=10 * fl.u.m,
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["*susp*"],
                geometry.snappy_bodies["fr-int-winglet-*"],
                geometry.snappy_bodies["fr-int-caketin-*"],
            ],
            min_spacing=2 * fl.u.mm,
            max_spacing=10 * fl.u.m,
            proximity_spacing=0.5 * fl.u.mm,
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["body"],
                geometry.snappy_bodies["body-nose"],
            ],
            min_spacing=5 * fl.u.mm,
            max_spacing=10 * fl.u.m,
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["*wh-cover-OB*"],
                geometry.snappy_bodies["fw*"],
                geometry.snappy_bodies["rw-ep"],
                geometry.snappy_bodies["rw-mp"],
                geometry.snappy_bodies["rw-flap"],
            ],
            min_spacing=1 * fl.u.mm,
            max_spacing=10 * fl.u.m,
        ),
        fl.snappy.BodyRefinement(
            bodies=[
                geometry.snappy_bodies["uf"],
                geometry.snappy_bodies["diffuser"],
            ],
            min_spacing=1 * fl.u.mm,
            max_spacing=10 * fl.u.m,
        ),
        fl.snappy.SurfaceEdgeRefinement(
            entities=[
                geometry.snappy_bodies["fw-flap-*"],
                geometry.snappy_bodies["fw-fine"],
                geometry.snappy_bodies["fw-ep"],
                geometry.snappy_bodies["rw-mp"],
                geometry.snappy_bodies["rw-flap"],
                geometry.snappy_bodies["diffuser"],
                geometry.snappy_bodies["uf-strakes"],
                geometry.snappy_bodies["body-pit"],
            ],
            spacing=1 * fl.u.mm,
            min_elem=3,
            retain_on_smoothing=True,
        ),
        fl.snappy.SurfaceEdgeRefinement(
            entities=[
                geometry.snappy_bodies["fw-mp"],
                geometry.snappy_bodies["rw-gf"],
                geometry.snappy_bodies["rw-ep"],
                geometry.snappy_bodies["*velocity*"],
                geometry.snappy_bodies["eb-exhaust"],
            ],
            spacing=[0.5 * fl.u.mm],
            distances=[1.5 * fl.u.mm],
            min_elem=3,
            retain_on_smoothing=True,
        )
    ]

with fl.SI_unit_system:
    surface_meshing = fl.snappy.SurfaceMeshingParams(
        defaults=surface_defaults,
        refinements=surface_refinements+[plinth_pucks],
        base_spacing=spacing,
        smooth_controls=smooth_controls,
        castellated_mesh_controls=castellated_mesh_controls
    )

with fl.SI_unit_system:
    volume_refinements = [
        fl.UniformRefinement(
            entities=[
                fl.Box(
                    name="fr_wing",
                    center=[-650, 0, 250] * fl.u.mm,
                    size=[900, 2100, 600] * fl.u.mm,
                ),
                fl.Box(
                    name="rr_wing",
                    center=[4050, 0, 700] * fl.u.mm,
                    size=[700, 1400, 800] * fl.u.mm,
                ),
            ],
            spacing=6 * fl.u.mm,
        ),
        fl.UniformRefinement(
            entities=[
                fl.Box(
                    name="whole",
                    center=[2000, 0, 750] * fl.u.mm,
                    size=[6500, 2300, 1600] * fl.u.mm,
                )
            ],
            spacing=10 * fl.u.mm,
        ),
    ]

    volume_meshing = fl.VolumeMeshingParams(
        defaults=fl.VolumeMeshingDefaults(
            boundary_layer_first_layer_thickness=0.5 * fl.u.mm,
            boundary_layer_growth_rate=1.2,
        ),
        refinements=[
            fl.PassiveSpacing(
                faces=[
                    geometry.snappy_bodies["tunnel"]["top"],
                    geometry.snappy_bodies["tunnel"]["side*"],
                ],
                type="unchanged",
            ),
            fl.PassiveSpacing(
                faces=[
                    geometry.snappy_bodies["tunnel"]["inlet"],
                    geometry.snappy_bodies["tunnel"]["outlet"],
                ],
                type="projected",
            ),
            fl.BoundaryLayer(
                faces=[geometry.snappy_bodies["tunnel"]["ground"]],
                first_layer_thickness=1 * fl.u.mm,
                growth_rate=1.4,
            ),
            fl.BoundaryLayer(
                faces=[
                    geometry.snappy_bodies["fw*"]["*"],
                    geometry.snappy_bodies["rw*"]["*"],
                    geometry.snappy_bodies["uf*"]["*"],
                ],
                first_layer_thickness=0.2 * fl.u.mm,
            ),
            fl.BoundaryLayer(
                faces=[geometry.snappy_bodies["*susp*"]["*"]],
                first_layer_thickness=0.35 * fl.u.mm,
            ),
            *volume_refinements,
        ],
        planar_face_tolerance=1e-3,
        gap_treatment_strength=1,
    )

meshing = fl.ModularMeshingWorkflow(
    surface_meshing=surface_meshing,
    volume_meshing=volume_meshing,
    zones=[fl.CustomZones(entities=[fluid, rad_lhs, rad_rhs])],
)

operating_condition = fl.AerospaceCondition(
    velocity_magnitude=50 * fl.u.m / fl.u.s,
    alpha=-0.273031 * fl.u.deg,
    beta=0.5 * fl.u.deg,
)

time_stepping = fl.Steady(
    CFL=fl.AdaptiveCFL(
        min=0.1,
        max=15000,
        max_relative_change=1,
        convergence_limiting_factor=0.25,
    ),
    max_steps=6000,  # more steps than the mesh-only setup
)

wall_model_BC = [
    fl.Wall(
        surfaces=[
            geometry.snappy_bodies["fw*"]["*"],
            geometry.snappy_bodies["rw*"]["*"],
            geometry.snappy_bodies["body*"]["*"],
            geometry.snappy_bodies["uf*"]["*"],
            geometry.snappy_bodies["eb*"]["*"],
            geometry.snappy_bodies["fr-int*"]["*"],
            geometry.snappy_bodies["fr-susp*"]["*"],
            geometry.snappy_bodies["rr-int*"]["*"],
            geometry.snappy_bodies["rr-susp*"]["*"],
            geometry.snappy_bodies["diffuser*"]["*"],
        ],
        use_wall_function=True,
    )
]

moving_wall = [
    fl.Wall(
        surfaces=[geometry.snappy_bodies["tunnel"]["ground"]],
        use_wall_function=True,
        velocity=[
            50 * cos(radians(0.5)) + 50 * cos(radians(-0.273031)),
            50 * sin(radians(0.5)),
            50 * sin(radians(-0.273031)),
        ]
        * fl.u.m
        / fl.u.s,
    )
]

rotating_walls = [
    fl.Wall(
        surfaces=[geometry.snappy_bodies["fr-wh*lhs"]["*"]],
        use_wall_function=True,
        velocity=fl.WallRotation(
            center=[
                0.20344675860875966,
                -0.7650740158453145,
                0.3592397445075839,
            ]
            * fl.u.m,
            axis=[
                -0.015622695450066178,
                -0.995547991275173,
                0.09295229128344651,
            ],
            angular_velocity=141.50663390275048 * fl.u.rad / fl.u.s,
        ),
        roughness_height=1 * fl.u.mm,
    ),
    fl.Wall(
        surfaces=[geometry.snappy_bodies["fr-wh*rhs"]["*"]],
        use_wall_function=True,
        roughness_height=1 * fl.u.mm,
        velocity=fl.WallRotation(
            center=[
                0.20344675860875977,
                0.7650740158453145,
                0.3592397445075839,
            ]
            * fl.u.m,
            axis=[
                0.015622695450066141,
                -0.995547991275173,
                -0.0929522912834463,
            ],
            angular_velocity=141.50663390275048 * fl.u.rad / fl.u.s,
        ),
    ),
    fl.Wall(
        surfaces=[geometry.snappy_bodies["rr-wh*lhs"]["*"]],
        use_wall_function=True,
        roughness_height=1 * fl.u.mm,
        velocity=fl.WallRotation(
            center=[
                3.6466031503561447,
                -0.742088395684798,
                0.34259442839962195,
            ]
            * fl.u.m,
            axis=[
                0.001488022855271076,
                -0.9993832015999016,
                0.03508564019528322,
            ],
            angular_velocity=142.1293673304936 * fl.u.rad / fl.u.s,
        ),
    ),
    fl.Wall(
        surfaces=[geometry.snappy_bodies["rr-wh*rhs"]["*"]],
        use_wall_function=True,
        roughness_height=1 * fl.u.mm,
        velocity=fl.WallRotation(
            center=[
                3.6466031503561456,
                0.7420883956847978,
                0.342594428389622,
            ]
            * fl.u.m,
            axis=[
                -0.001488022855274553,
                -0.9993832015999016,
                -0.0350856401952817,
            ],
            angular_velocity=142.1293673304936 * fl.u.rad / fl.u.s,
        ),
    ),
]

slip_wall = [fl.SlipWall(surfaces=[geometry.snappy_bodies["tunnel"]["top"]])]
freestream = [
    fl.Freestream(
        surfaces=[
            geometry.snappy_bodies["tunnel"]["inlet"],
            geometry.snappy_bodies["tunnel"]["side*"],
            geometry.snappy_bodies["tunnel"]["outlet"],
        ]
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-inlet-exhaust-out"]["*"],
        velocity=[100, 0, 0] * fl.u.m / fl.u.s,
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-inlet-fr-int-duct-*-out"]["*"],
        velocity=[25.38, 0, 0] * fl.u.m / fl.u.s,
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-inlet-rr-int-duct-*-out"]["*"],
        velocity=[25.38, 0, 0] * fl.u.m / fl.u.s,
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-outlet-engine-intake-in"]["*"],
        velocity=[21.991020894901634, 0, -21.991020894901634] * fl.u.m / fl.u.s,
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-outlet-fr-int-duct-*-in"]["*"],
        velocity=[15.36, 0, 0] * fl.u.m / fl.u.s,
    ),
    fl.Freestream(
        surfaces=geometry.snappy_bodies["velocity-outlet-rr-int-duct-*-in"]["*"],
        velocity=[24.93, 0, 0] * fl.u.m / fl.u.s,
    ),
]

porous_medium = [
    fl.PorousMedium(
        volumes=[rad_lhs, rad_rhs],
        darcy_coefficient=(3.75e7, 3.75e8, 3.75e8) / fl.u.m**2,
        forchheimer_coefficient=(212.8, 2128, 2128) / fl.u.m,
    )
]

user_defined_fields = [
    fl.UserDefinedField(
        name="CpT",
        expression=(
            "double gamma = 1.40; double pow1 = gamma/(gamma-1); double pow2 = (gamma-1) / 2; "
            "double MachRefSq = MachRef * MachRef; "
            "double Mach = sqrt(primitiveVars[1] * primitiveVars[1] + primitiveVars[2] * primitiveVars[2] + primitiveVars[3] * primitiveVars[3]) / "
            "sqrt(gamma * primitiveVars[4] / primitiveVars[0]); "
            "double MachSq = Mach * Mach;"
            "CpT = 2 /(gamma*MachRefSq)* (pow((1+pow2*MachSq),pow1)*gamma *primitiveVars[4]-1 );"
        ),
    ),
    fl.UserDefinedField(
        name="vel_adim",
        expression=(
            "vel_adim[0]=primitiveVars[1]/ MachRef; "
            "vel_adim[1]=primitiveVars[2] / MachRef; "
            "vel_adim[2]=primitiveVars[3] / MachRef;"
        ),
    ),
]

with fl.SI_unit_system:
    x_positions = np.arange(-1.1, 5.475, 0.025)
    y_positions = np.arange(-1.1, 1.1, 0.025)
    z_positions = np.arange(-0.05, 1.25, 0.025)

    outputs = [
        fl.SurfaceOutput(
            surfaces=geometry.snappy_bodies["*"]["*"],
            output_format="paraview",
            output_fields=["primitiveVars", "Cp", "yPlus", "CfVec"],
        ),
        fl.VolumeOutput(
            name="fl.VolumeOutput",
            output_format="paraview",
            output_fields=[
                "primitiveVars",
                "vorticity",
                "residualNavierStokes",
                "residualTurbulence",
                "Cp",
                "CpT",
                "qcriterion",
                "vel_adim",
            ],
        ),
        fl.SliceOutput(
            name="slices",
            entities=[
                *[
                    fl.Slice(
                        name=f"slice_x_{i}",
                        normal=(
                            cos(radians(-0.2594)),
                            0,
                            sin(radians(-0.2594)),
                        ),
                        origin=(x, 0, 0),
                    )
                    for i, x in enumerate(x_positions)
                ],
                *[
                    fl.Slice(
                        name=f"slice_y_{i}",
                        normal=(0, 1, 0),
                        origin=(0, y, 0),
                    )
                    for i, y in enumerate(y_positions)
                ],
                *[
                    fl.Slice(
                        name=f"slice_z_{i}",
                        normal=(0, 0, 1),
                        origin=(0, 0, z),
                    )
                    for i, z in enumerate(z_positions)
                ],
            ],
            output_format="paraview",
            output_fields=["CpT", "Cp", "vel_adim", "vorticity", "qcriterion"],
        ),
    ]

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing,
        operating_condition=operating_condition,
        time_stepping=time_stepping,
        models=[
            *wall_model_BC,
            *moving_wall,
            *rotating_walls,
            *slip_wall,
            *freestream,
            *porous_medium,
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    absolute_tolerance=1e-10,
                    linear_solver=fl.LinearSolver(max_iterations=50),
                    kappa_MUSCL=0.33,
                    limit_pressure_density=False,
                    limit_velocity=False,
                    low_mach_preconditioner=True,
                    numerical_dissipation_factor=0.2,
                ),
                turbulence_model_solver=fl.KOmegaSST(
                    absolute_tolerance=1e-8,
                    linear_solver=fl.LinearSolver(max_iterations=25),
                    equation_evaluation_frequency=4,
                    rotation_correction=True,
                    modeling_constants=fl.KOmegaSSTModelConstants(C_alpha1=0.4),
                ),
            ),
        ],
        user_defined_fields=user_defined_fields,
        outputs=outputs,
        reference_geometry=fl.ReferenceGeometry(
            moment_center=(1.4001, 0, -0.3176),  # reference moment center
            moment_length=(1, 2.7862, 1),  # reference lengths
            area=1.7095,  # reference area
        ),
    )

case = project.run_case(
    params=params,
    name="Tutorial F1 Car from Python",
    use_beta_mesher=True,
)