import flow360 as fl
from flow360.examples import Windsor
Windsor.get_files()

prj = fl.Project.from_geometry(
    [Windsor.geometry, Windsor.extra["wheel"]], 
    length_unit='mm', 
    name="WindsorWithWheels",
)

geometry = prj.geometry

draft = fl.create_draft(
    new_run_from=geometry, 
    face_grouping='groupByBodyId',
    edge_grouping='edgeId'
)

with draft:
    windsor_body = draft.body_groups["windsorBody.stp"]
    front_wheel = draft.body_groups["wheel.stl"]

    draft.coordinate_systems.assign(
        entities=front_wheel,
        coordinate_system=fl.CoordinateSystem(
            name="shift",
            reference_point=[0, 0, 0] * fl.u.mm,
            axis_of_rotation=(0, 0, 1),
            angle_of_rotation=0 * fl.u.deg,
            scale=(1.0, 1.0, 1.0),
            translation=[-320.025, -167, 70] * fl.u.mm
        )
    )

    mirror_plane = fl.MirrorPlane(
        name="mirror_wheel", 
        normal=(1, 0, 0), 
        center=(0, 0, 0) * fl.u.mm
    )
    mirrored_body_groups, mirrored_surfaces = draft.mirror.create_mirror_of(
        entities=front_wheel, 
        mirror_plane=mirror_plane
    )
    rear_wheel = mirrored_body_groups[0]
    rear_wheel_surface = mirrored_surfaces[0]

    windsor_body.mesh_exterior = True
    front_wheel.mesh_exterior = True

wind_tunnel = fl.WindTunnelFarfield(
    width=1920 * fl.u.mm,
    height=1320 * fl.u.mm,
    inlet_x_position=-1800 * fl.u.mm,
    outlet_x_position=1800 * fl.u.mm,
    floor_z_position=0 * fl.u.mm,
    floor_type=fl.FullyMovingFloor(),
    domain_type="half_body_negative_y"  # Symmetry plane at y=0
)

meshing_params = fl.MeshingParams(
    defaults=fl.MeshingDefaults(
        surface_max_edge_length=10 * fl.u.mm,
        geometry_accuracy=1 * fl.u.mm,
        boundary_layer_first_layer_thickness=1e-3 * fl.u.mm,
        boundary_layer_growth_rate=1.2,
        planar_face_tolerance=1e-3
    ),
    refinements=[
        fl.PassiveSpacing(
            faces=[
                wind_tunnel.left,
                wind_tunnel.symmetry_plane,
                wind_tunnel.inlet,
                wind_tunnel.outlet,
            ], 
            type="projected"
        ),
        fl.PassiveSpacing(
            faces=[
                wind_tunnel.ceiling,
            ],
            type="unchanged"
        ),
        fl.SurfaceRefinement(
            faces=[
                wind_tunnel.left,
                wind_tunnel.symmetry_plane,
                wind_tunnel.inlet,
                wind_tunnel.outlet,
                wind_tunnel.floor,
                wind_tunnel.ceiling
            ],
            max_edge_length=100 * fl.u.mm
        )
    ],
    volume_zones=[wind_tunnel],
    )

with draft:
    stationary_wall_surfaces = [draft.surfaces["body00001"], draft.surfaces["wheel.stl"], rear_wheel_surface]

    BCs = [
        fl.Wall(
            entities=stationary_wall_surfaces,
            use_wall_function=False
        ), 
        fl.SlipWall(
            entities=[
                wind_tunnel.left,
                wind_tunnel.symmetry_plane,
                wind_tunnel.ceiling,
            ]
        ),
        fl.Wall(
            entities=[wind_tunnel.floor],
            velocity=[30, 0, 0] * fl.u.m / fl.u.s,
            use_wall_function=False
        ),
        fl.Freestream(
            entities=[
                wind_tunnel.inlet,
                wind_tunnel.outlet,
            ]
        )
    ]

cd_output = fl.ForceOutput(
    name="CD-windsor",
    models=[BCs[0]],
    moving_statistic=fl.MovingStatistic(moving_window_size=100, method="range", start_step=500),
    output_fields=['CD']
)

run_control = fl.RunControl(stopping_criteria=[
    fl.StoppingCriterion(monitor_output=cd_output, tolerance=0.02, monitor_field="CD")
    ]
)

with fl.SI_unit_system:
    params = fl.SimulationParams(
        meshing=meshing_params,
        reference_geometry=fl.ReferenceGeometry(
            area=0.056,
            moment_length=(0.6375, 0.6375, 0.6375),
            moment_center=(0, 0, 0),
        ),
        operating_condition=fl.AerospaceCondition(
            velocity_magnitude = 30 * fl.u.m / fl.u.s
        ),
        models=[
            fl.Fluid(
                navier_stokes_solver=fl.NavierStokesSolver(
                    linear_solver=fl.LinearSolver(max_iterations=50),
                    absolute_tolerance=1e-6,
                    low_mach_preconditioner=True
                ),
                turbulence_model_solver=fl.SpalartAllmaras(
                    absolute_tolerance=1e-4
                )
            ),
            *BCs
        ],
        time_stepping=fl.Steady(
            CFL=fl.AdaptiveCFL(max=1000),
            max_steps=2000
        ),
        outputs=[
            fl.VolumeOutput(
                output_format="paraview",
                output_fields=[
                    "primitiveVars",
                    "Mach",
                    "mutRatio"
                ]
            ),
            fl.SurfaceOutput(
                entities=[stationary_wall_surfaces],
                output_format="paraview",
                output_fields=[
                    "primitiveVars",
                    "Cp",
                    "Cf",
                    "yPlus"
                ]
            ),
            fl.SliceOutput(
                slices=[
                    fl.Slice(
                        name="Underbody slice",
                        normal=(0, 0, 1),
                        origin=(0, 0, 0.03)*fl.u.m
                    )
                ],
                output_fields=[fl.UserVariable(name="velocity_dim", value=fl.solution.velocity).in_units(fl.u.m/fl.u.s), "velocity"]
            ),
            cd_output # force output included
        ],
        run_control=run_control
    )

with draft:
    case = prj.run_case(
        params=params, 
        name="WindsorWithWheels_0deg", 
        use_geometry_AI=True,  # Enable Geometry AI
        use_beta_mesher=True,  # Required when using Geometry AI
    )

draft.coordinate_systems.get_by_name("shift").angle_of_rotation = -5*fl.u.deg # toe in on the front and toe out on the rear wheel

with draft:
    volume_mesh = prj.generate_volume_mesh(
        params=params, 
        name="WindsorWithWheels_-5deg", 
        use_geometry_AI=True,  # Enable Geometry AI
        use_beta_mesher=True,  # Required when using Geometry AI
    )

with draft:
    case_turned = prj.run_case(
        params=params, 
        name="WindsorWithWheels_-5deg", 
        use_geometry_AI=True,  # Enable Geometry AI
        use_beta_mesher=True,  # Required when using Geometry AI
        fork_from=case,
        interpolate_to_mesh=volume_mesh
    )

from flow360.plugins.report import (
    ReportTemplate, 
    Chart3D, 
    Table, 
    DataItem, 
    Average, 
    TopCamera,
    NonlinearResiduals
)

case.wait()
case_turned.wait()

CD = DataItem(data="surface_forces/totalCD", exclude=["farfield/"+wind_tunnel.floor.name], operations=[Average(fraction=0.1)])

comp_table = Table(data=[CD], section_title="Table comparison of coefficients")

near_ground_flow = Chart3D(
    section_title=f"Slice velocity z=0.03m",
    items_in_row=2,
    force_new_page=True,
    show="slices",
    include=["Underbody slice"],
    field="velocity",
    mode="lic",
    limits=(25*fl.u.m/fl.u.s , 40*fl.u.m/fl.u.s),
    camera=TopCamera(dimensionDirection="width", dimension=1500),
    fig_name="slice_z",
)

report_template = ReportTemplate(
    title="Toe angle comparison",
    items=[
        NonlinearResiduals(),
        comp_table,
        near_ground_flow
    ]
)

report = report_template.create_in_cloud(
    name="report-toe-angle",
    cases=[case, case_turned],
)

report.wait()

report.download("report.pdf")