import ifcopenshell
import pandas as pd
from ifcopenshell.util import cost, element
from bonsai import tool

# Step 1: Get the model
model = tool.Ifc.get()

# Step 2: Identify the IfcCostSchedule
myCostSchedule = model.by_type("IfcCostSchedule")[0]

# Step 3: Get all IfcCostItem elements in the cost schedule
CostItemLst = list(cost.get_schedule_cost_items(myCostSchedule))

# Step 4: Create a list of IfcCostItem elements that have NO nested ones
FinalCostItemLst = [cstitm for cstitm in CostItemLst if not list(cost.get_nested_cost_items(cstitm))]

# Step 5: Define the allowed element types and their corresponding quantity properties + units
FILTER_TYPES = {
    "IfcColumn": ("Qto_ColumnBaseQuantities", "NetVolume", "m³"),
    "IfcBeam": ("Qto_BeamBaseQuantities", "NetVolume", "m³"),
    "IfcSlab": ("Qto_SlabBaseQuantities", "NetVolume", "m³"),
    "IfcFooting": ("Qto_FootingBaseQuantities", "NetVolume", "m³"),
    "IfcWall": ("Qto_WallBaseQuantities", "NetSideArea", "m²"),
    "IfcCovering": ("Qto_CoveringBaseQuantities", "NetArea", "m²"),  # Default for IfcCovering
    "IfcDoor": ("Qto_DoorBaseQuantities", "Area", "m²"),
    "IfcWindow": ("Qto_WindowBaseQuantities", "Area", "m²"),
}

# Step 6: Set output directory and filename
output_file = r"C:\Users\steve\OneDrive\Desktop\CostItem_Assignments.csv"

# List to store all data
all_data = []

# Step 7: Loop through cost items
for costitem in FinalCostItemLst:
    assigned_elements = cost.get_cost_item_assignments(costitem)  # Get assigned elements

    # Manually filter elements by type
    filtered_elements = [elem for elem in assigned_elements if elem.is_a() in FILTER_TYPES]

    if not filtered_elements:  # Skip cost items if no relevant elements remain
        continue

    # Get the cost item name (cleaned for CSV)
    cost_item_name = costitem.Name if costitem.Name else f"CostItem_{costitem.id()}"

    # Collect data for the DataFrame
    for elem in filtered_elements:
        global_id = elem.GlobalId  # NEW: Extract GlobalId
        elem_type = elem.is_a()  # The IFC class (e.g., IfcColumn, IfcBeam)
        type_name = elem.IsTypedBy[0].RelatingType.Name if elem.IsTypedBy else "Unknown"

        # Get the correct quantity set, property name, and unit
        qto_set_name, property_name, unit = FILTER_TYPES[elem_type]
        qto_set = element.get_psets(elem, qtos_only=True).get(qto_set_name, {})

        # Special handling for IfcCovering with PredefinedType == SKIRTINGBOARD
        if elem_type == "IfcCovering" and getattr(elem, "PredefinedType", None) == "SKIRTINGBOARD":
            property_name = "Width"  # Use Width directly
            unit = "m"  # Update unit to meters
            quantity_value = qto_set.get("Width", "N/A")
        else:
            # Extract relevant quantity and round to 2 decimal places
            quantity_value = qto_set.get(property_name, "N/A")

        if isinstance(quantity_value, (int, float)):  # Ensure it's a number before rounding
            quantity_value = round(quantity_value, 2)

        all_data.append({
            "GlobalId": global_id,  # NEW: GlobalId as first column
            "Element": cost_item_name,
            "IFCType": elem_type,
            "TypeName": type_name,
            "Selector": property_name,
            "Unit": unit,
            "Quantity": quantity_value,
            
        })

# Step 8: Convert list to DataFrame
df = pd.DataFrame(all_data)

# Step 9: Save to CSV
df.to_csv(output_file, index=False)

print(f"CSV file saved at: {output_file}")