import bpy
import bonsai.tool as tool
import ifcopenshell.api
import ifcopenshell.util.placement
import random
from mathutils import Vector
from mathutils.geometry import delaunay_2d_cdt
import pandas as pd


def make_ifcBorehole(name, location):
    bpy.ops.object.empty_add(type='PLAIN_AXES', radius=1, align='WORLD', location=location, scale=(1, 1, 1))
    new_empty = bpy.context.active_object
    new_empty.name = name
    bpy.ops.bim.assign_class(ifc_class="IfcBorehole", predefined_type="", userdefined_type="")
    new_empty.select_set(False)
    
    
def make_ifcGeotechnicalStratum_for_borehole(name, location, depthTop, depth, diameter):
    bpy.ops.mesh.primitive_cylinder_add(vertices=16, radius=diameter/2, depth=depth, enter_editmode=False, align='WORLD', 
        location=location, scale=(1, 1, 1))
    new_stratum = bpy.context.active_object
    new_stratum.name = name
    new_stratum.location.z -= depthTop+depth/2
    bpy.ops.bim.assign_class(ifc_class="IfcGeotechnicalStratum", predefined_type="SOLID", userdefined_type="")


def make_material_stratum(name_material):
    material = ifcopenshell.api.material.add_material(model, name=name_material, category="earth")
    
    style = ifcopenshell.api.style.add_style(model, name="style_"+name_material)
    ifcopenshell.api.style.add_surface_style(model,style=style, ifc_class="IfcSurfaceStyleShading", attributes={
        "SurfaceColour": { "Name": None, "Red": random.random(), "Green": random.random(), "Blue": random.random()},
        "Transparency": 0.
        })
                
    context = None
    contexts = model.by_type("IfcGeometricRepresentationContext")
    for c in contexts: 
        if c.ContextIdentifier=='Body':
            context = c
            break
        
    ifcopenshell.api.style.assign_material_style(model, material=material, style=style, context=context)

   
def get_stratum(name_stratum):
    stratum = None
    for element in model.by_type("IfcGeotechnicalStratum"):
        if hasattr(element, "Name") and element.Name == name_stratum:
            stratum = element
            break
    return stratum


def get_material(name_material):
    material = None
    for element in model.by_type("IfcMaterial"):
        if hasattr(element, "Name") and element.Name == name_material:
            material = element
            break
    return material


def get_borehole(name_borehole):
    borehole = None 
            
    for element in model.by_type("IfcBorehole"):
        if hasattr(element, "Name") and element.Name == name_borehole:
            borehole = element
            break
    return borehole


def assing_material_to_stratum(stratum, material):
    ifcopenshell.api.material.assign_material(model, products=[stratum], type="IfcMaterial", material=material)


def agregatte_stratum_to_borehole(borehole, stratum):
    ifcopenshell.api.aggregate.assign_object(model, products=[stratum], relating_object=borehole)
   

def assing_qto_to_stratum(stratum, depth, diameter):
    qto = ifcopenshell.api.pset.add_qto(model, product=stratum, name="Qto_LinearStratumBaseQuantities")
    ifcopenshell.api.pset.edit_qto(model, qto=qto, properties={"Diameter": diameter, "Length": depth})
    

def assing_pset_to_stratum(stratum, name_material):
    name_stratum_composite = name_material.strip().upper().replace(" ","_")
    
    PEnum_SoilCompositeFractions =[
    "BOULDERS",
    "BOULDERS_WITH_COBBLES",
    "BOULDERS_WITH_FINER_SOILS",
    "CLAY",
    "CLAYEY_SILT",
    "COBBLES",
    "COBBLES_WITH_BOULDERS",
    "COBBLES_WITH_FINER_SOILS",
    "FILL",
    "GRAVEL",
    "GRAVELLY_SAND",
    "GRAVEL_WITH_CLAY_OR_SILT",
    "GRAVEL_WITH_COBBLES",
    "ORGANIC_CLAY",
    "ORGANIC_SILT",
    "SAND",
    "SANDY_CLAYEY_SILT",
    "SANDY_GRAVEL",
    "SANDY_GRAVELLY_CLAY",
    "SANDY_GRAVELLY_SILT",
    "SANDY_GRAVEL_WITH_COBBLES",
    "SANDY_PEAT",
    "SANDY_SILT",
    "SAND_WITH_CLAY_AND_SILT",
    "SILT",
    "SILTY_CLAY",
    "OTHER",
    "NOTKNOWN",
    "UNSET"
    ]
    
    
    tmp_composite = []
    
    if name_stratum_composite in PEnum_SoilCompositeFractions:
        tmp_composite.append(name_stratum_composite)
    else:
        tmp_composite.append("OTHER")
    
    pset = ifcopenshell.api.pset.add_pset(model, product=stratum, name="Pset_SolidStratumComposition")
    ifcopenshell.api.pset.edit_pset(model, pset=pset, properties={"CompositeFractions": tmp_composite})



************************************************************************************************************************************
#change file path
file_borehole = "file_path/boreholes_data.xlsx" 
file_strata = "file_path/strata_data.xlsx"


model = tool.Ifc.get()

df_boreholes = pd.read_excel(file_borehole)
df_strata = pd.read_excel(file_strata)


#iterates over the stratum types and makes the materials
strata_material = set(df_strata['Geology'].unique())

for name_material in strata_material:
    make_material_stratum(name_material)


#make boreholes
for _, row in df_boreholes.iterrows():
    name_borehole = row['Location_ID']
    location = (row['X/Easting'], row['Y/Northing'], row['Z/Ground_elevation'])
    make_ifcBorehole(name_borehole, location)


#iterate by boreholes and make stratum
df_strata_grouped = df_strata.groupby('Location_ID')
diameter = 0.1

for name, group in df_strata_grouped:
    st_numb = 1
    borehole = get_borehole(name)
    borehole_location = ifcopenshell.util.placement.get_local_placement(borehole.ObjectPlacement)
    location = (borehole_location[0][3], borehole_location[1][3], borehole_location[2][3])
    for _, row in group.iterrows():
        name_material = row['Geology']
        name_stratum = name+"_strat_"+str(st_numb)+"_"+name_material
        st_numb+=1
        depthTop = row['Depth_Top']
        depthBottom = row['Depth Bottom']
        depth = depthBottom - depthTop
        make_ifcGeotechnicalStratum_for_borehole(name_stratum, location, depthTop, depth, diameter)
        stratum = get_stratum(name_stratum)
        ifcopenshell.api.attribute.edit_attributes(model, product=stratum, attributes={"Description": row['Description']})
        material = get_material(name_material)
        assing_material_to_stratum(stratum, material)
        assing_qto_to_stratum(stratum, depth, diameter)
        assing_pset_to_stratum(stratum, name_material)
        agregatte_stratum_to_borehole(borehole, stratum)
    

bpy.ops.bim.save_project(filepath=tool.Ifc.get_path(), should_save_as=True, use_relative_path=False)