Cantilever

from generative.core import generative_function, FileAsset, Asset, GenerativeType
import trimesh
from pydantic import Field
import shapely


class CantileverGeometricProperties(GenerativeType):
    length_m: float = Field(default=1, gt=0)
    depth_m: float = Field(default=0.1, gt=0)
    width_m: float = Field(default=0.2, gt=0)
    wall_thickness_m: float = Field(default=0.02, gt=0)


class PointLoadScenario(GenerativeType):
    force_newtons: float = Field(default=10)
    force_location_proportion: float = Field(
        default=1,
        ge=0,
        le=1,
    )


class CantileverInputs(GenerativeType):
    geometry: CantileverGeometricProperties
    material_idx: int = Field(default=0, ge=0, le=1)
    load_case: PointLoadScenario


class CantileverAnalysisOutputs(GenerativeType):
    slope_radians: float
    deflection_m: float


class MaterialProperties(GenerativeType):
    material_name: str
    modulus_pascals: float = Field(gt=0)
    cost_gbp_per_kg: float = Field(gt=0)
    density_kg_per_m3: float = Field(gt=0)


class CantileverCostOutputs(GenerativeType):
    total_cost_gbp: float
    material_cost_gbp: float
    fabrication_cost_gbp: float


class CantileverOutputs(GenerativeType):
    cost: CantileverCostOutputs
    structural: CantileverAnalysisOutputs
    model: Asset


@generative_function
def cantilever(inputs: CantileverInputs) -> CantileverOutputs:
    validate_inputs(inputs)

    material = material_database(inputs.material_idx)
    cost = cantilever_cost_model(inputs.geometry, material)
    analysis = cantilever_point_load_analysis(inputs.geometry, inputs.load_case, material)
    model = cantilever_3d_model(inputs.geometry)
    return CantileverOutputs(cost=cost, structural=analysis, model=model)


def validate_inputs(inputs: CantileverInputs) -> None:
    max_possible_wall_thickness = 0.5 * min(inputs.geometry.width_m, inputs.geometry.depth_m)
    if inputs.geometry.wall_thickness_m >= max_possible_wall_thickness:
        raise ValueError("Total wall thickness larger than width or depth")


def cantilever_point_load_analysis(
    geom: CantileverGeometricProperties,
    load: PointLoadScenario,
    material: MaterialProperties,
) -> CantileverAnalysisOutputs:
    """Calculate cantilever performance using standard formulas: https://w.wiki/D2ue"""
    second_mom_area_m4 = second_moment_of_area(geom.width_m, geom.depth_m, geom.wall_thickness_m)
    load_location = load.force_location_proportion * geom.length_m

    slope_numerator = load.force_newtons * load_location**2
    slope = slope_numerator / (2 * material.modulus_pascals * second_mom_area_m4)
    def_numerator = load.force_newtons * load_location**2 * (3 * geom.length_m - load_location)
    deflection = def_numerator / (6 * material.modulus_pascals * second_mom_area_m4)
    return CantileverAnalysisOutputs(slope_radians=slope, deflection_m=deflection)


def second_moment_of_area(outer_width: float, outer_depth: float, wall_thickness: float) -> float:
    inner_width = outer_width - 2 * wall_thickness
    inner_depth = outer_depth - 2 * wall_thickness
    return (outer_width * outer_depth**3 - inner_width * inner_depth**3) / 12


MATERIAL_DATA = {
    0: MaterialProperties(
        material_name="steel", modulus_pascals=210e9, cost_gbp_per_kg=10, density_kg_per_m3=78500
    ),
    1: MaterialProperties(
        material_name="aluminium",
        modulus_pascals=70e9,
        cost_gbp_per_kg=100,
        density_kg_per_m3=2700,
    ),
}


def material_database(material_idx: int) -> MaterialProperties:
    try:
        return MATERIAL_DATA[material_idx]
    except KeyError as e:
        raise ValueError(
            f"Unknown material index '{material_idx}', try {list(MATERIAL_DATA.keys())}"
        ) from e


def cantilever_cost_model(
    geometry: CantileverGeometricProperties,
    material: MaterialProperties,
) -> CantileverCostOutputs:
    outer_cross_section = geometry.width_m * geometry.depth_m
    inner_cross_section = (geometry.width_m - 2 * geometry.wall_thickness_m) * (
        geometry.depth_m - 2 * geometry.wall_thickness_m
    )
    volume = geometry.length_m * (outer_cross_section - inner_cross_section)
    material_cost = material.cost_gbp_per_kg * volume * material.density_kg_per_m3
    fabrication_cost = 10_000 + 100 * volume

    return CantileverCostOutputs(
        total_cost_gbp=fabrication_cost + material_cost,
        fabrication_cost_gbp=fabrication_cost,
        material_cost_gbp=material_cost,
    )


def cantilever_3d_model(geometry: CantileverGeometricProperties) -> Asset:
    half_width, half_depth = geometry.width_m / 2, geometry.depth_m / 2
    wall_offset = geometry.wall_thickness_m
    polygon = shapely.Polygon(
        shell=[
            (-half_width, -half_depth),
            (half_width, -half_depth),
            (half_width, half_depth),
            (-half_width, half_depth),
        ],
        holes=[
            [
                (-half_width + wall_offset, -half_depth + wall_offset),
                (half_width - wall_offset, -half_depth + wall_offset),
                (half_width - wall_offset, half_depth - wall_offset),
                (-half_width + wall_offset, half_depth - wall_offset),
            ]
        ],
    )
    cantilever_mesh = trimesh.creation.extrude_polygon(
        polygon, height=geometry.length_m, engine="triangle"
    )

    asset = FileAsset()
    cantilever_mesh.export(str(asset.path), file_type="glb")
    return asset