import ifcopenshell
import ifcopenshell.api
import ifcopenshell.api.aggregate
import ifcopenshell.api.context
import ifcopenshell.api.feature
import ifcopenshell.api.geometry
import ifcopenshell.api.root
import ifcopenshell.api.spatial
import ifcopenshell.api.type
import ifcopenshell.api.unit
import ifcopenshell.geom
import ifcopenshell.util.representation
import ifcopenshell.util.unit
import numpy as np
import propertygroups

O = 0.0, 0.0, 0.0
X = 1.0, 0.0, 0.0
Y = 0.0, 1.0, 0.0
Z = 0.0, 0.0, 1.0


class Context:
    model = None
    body = None
    storey = None

    def __init__(self):
        self._create_empty_model()

    def clear(self):
        self._create_empty_model()

    def open(self, file_content):
        self.model = ifcopenshell.file.from_string(file_content)
        body = [
            ctx for ctx in self.model.by_type("IfcGeometricRepresentationContext") if ctx.ContextIdentifier == "Body"
        ]
        if body:
            self.body = body[0]
        else:
            context = ifcopenshell.api.context.add_context(self.model, context_type="Model")
            self.body = ifcopenshell.api.context.add_context(
                self.model,
                context_type="Model",
                context_identifier="Body",
                target_view="MODEL_VIEW",
                parent=context,
            )
        axis = [
            ctx for ctx in self.model.by_type("IfcGeometricRepresentationContext") if ctx.ContextIdentifier == "Axis"
        ]
        if body:
            self.axis = axis[0]
        else:
            context = ifcopenshell.api.context.add_context(self.model, context_type="Model")
            self.axis = ifcopenshell.api.context.add_context(
                self.model,
                context_type="Model",
                context_identifier="Axis",
                target_view="GRAPH_VIEW",
                parent=context,
            )
        self.storey = self.model.by_type("IfcBuildingStorey")[0]

    def _create_empty_model(self):
        # Create a blank model
        self.model = ifcopenshell.file()
        # All projects must have one IFC Project element
        project = ifcopenshell.api.root.create_entity(self.model, ifc_class="IfcProject", name="My Project")
        # Geometry is optional in IFC, but because we want to use geometry in this example, let's define units
        # Assigning without arguments defaults to metric units
        ifcopenshell.api.unit.assign_unit(self.model)
        # Let's create a modeling geometry context, so we can store 3D geometry (note: IFC supports 2D too!)
        context = ifcopenshell.api.context.add_context(self.model, context_type="Model")
        # In particular, in this example we want to store the 3D "body" geometry of objects, i.e. the body shape
        self.body = ifcopenshell.api.context.add_context(
            self.model,
            context_type="Model",
            context_identifier="Body",
            target_view="MODEL_VIEW",
            parent=context,
        )
        self.axis = ifcopenshell.api.context.add_context(
            self.model,
            context_type="Model",
            context_identifier="Axis",
            target_view="GRAPH_VIEW",
            parent=context,
        )
        # Create a site, building, and storey. Many hierarchies are possible.
        site = ifcopenshell.api.root.create_entity(self.model, ifc_class="IfcSite", name="My Site")
        building = ifcopenshell.api.root.create_entity(self.model, ifc_class="IfcBuilding", name="Building A")
        self.storey = ifcopenshell.api.root.create_entity(
            self.model,
            ifc_class="IfcBuildingStorey",
            name="Ground Floor",
        )
        # Since the site is our top level location, assign it to the project
        # Then place our building on the site, and our storey in the building
        ifcopenshell.api.aggregate.assign_object(
            self.model,
            relating_object=project,
            products=[site],
        )
        ifcopenshell.api.aggregate.assign_object(
            self.model,
            relating_object=site,
            products=[building],
        )
        ifcopenshell.api.aggregate.assign_object(
            self.model,
            relating_object=building,
            products=[self.storey],
        )

    def create_2pt_wall(self, p1, p2, elevation, height, thickness, container, wall_type=None):
        p1 = np.array([p1[0], p1[1]])
        p2 = np.array([p2[0], p2[1]])

        wall = ifcopenshell.api.root.create_entity(self.model, ifc_class="IfcWall")
        length = float(np.linalg.norm(p2 - p1))
        representation = ifcopenshell.api.geometry.add_wall_representation(
            self.model,
            context=self.body,
            length=length,
            height=height,
            thickness=thickness,
        )
        ifcopenshell.api.geometry.assign_representation(
            self.model,
            product=wall,
            representation=representation,
        )
        representation = ifcopenshell.api.geometry.add_axis_representation(
            self.model,
            context=self.axis,
            axis=[(0.0, 0.0), (length, 0.0)],
        )
        ifcopenshell.api.geometry.assign_representation(
            self.model,
            product=wall,
            representation=representation,
        )
        v = p2 - p1
        v = np.divide(v, float(np.linalg.norm(v)), casting="unsafe")
        matrix = np.array(
            [
                [v[0], -v[1], 0, p1[0]],
                [v[1], v[0], 0, p1[1]],
                [0, 0, 1, elevation],
                [0, 0, 0, 1],
            ]
        )
        ifcopenshell.api.geometry.edit_object_placement(self.model, product=wall, matrix=matrix)
        ifcopenshell.api.spatial.assign_container(
            self.model,
            relating_structure=container,
            products=[wall],
        )
        if wall_type:
            ifcopenshell.api.type.assign_type(
                self.model,
                related_object=wall,
                relating_type=wall_type,
            )

        return wall

    def get_element(self, guid):
        return self.model[guid]

    def get_model(self):
        return self.model

    def create_fill(self, ty, pt, wall):
        if isinstance(wall, str):
            wall = self.model[wall]
        if not wall.is_a("IfcWall"):
            raise ValueError("Only 'wall' hosts are supported")
        if ty == "door":
            props = propertygroups.BIMDoorProperties()
        elif ty == "window":
            props = propertygroups.BIMWindowProperties()
        else:
            raise ValueError("Only 'door' or 'window' fills are supported")
        si_conversion = ifcopenshell.util.unit.calculate_unit_scale(self.model)
        body = ifcopenshell.util.representation.get_context(self.model, "Model", "Body", "MODEL_VIEW")
        representation_data = props.to_dict(si_conversion=si_conversion)
        representation_data["context"] = body
        if ty == "door":
            door_representation = ifcopenshell.api.geometry.add_door_representation(self.model, **representation_data)
        else:
            door_representation = ifcopenshell.api.geometry.add_window_representation(self.model, **representation_data)
        door = ifcopenshell.api.root.create_entity(self.model, ifc_class=f"ifc{ty}")
        door.OverallWidth = props.overall_width / si_conversion
        door.OverallHeight = props.overall_height / si_conversion
        ifcopenshell.api.geometry.assign_representation(
            self.model,
            product=door,
            representation=door_representation,
        )
        ifcopenshell.api.spatial.assign_container(
            self.model,
            relating_structure=self.storey,
            products=[door],
        )

        r = [r for r in wall.Representation.Representations if r.RepresentationIdentifier == "Axis"]
        if not r:
            raise ValueError("Axis representation is needed")
        r = r[0]
        axis_geometry = ifcopenshell.geom.create_shape(
            ifcopenshell.geom.settings(
                DIMENSIONALITY=ifcopenshell.ifcopenshell_wrapper.CURVES_SURFACES_AND_SOLIDS,
                USE_WORLD_COORDS=True,
            ),
            wall,
            r,
        )
        vs = np.array(axis_geometry.geometry.verts).reshape((-1, 3))
        es = np.array(axis_geometry.geometry.edges).reshape((-1, 2))
        A, B = vs[es[0]]
        v = B - A
        P = np.zeros(3)
        P[0 : len(pt)] = pt
        AP = P - A
        AP_dot_v = np.dot(AP, v)
        v_dot_v = np.dot(v, v)
        t = AP_dot_v / v_dot_v * np.linalg.norm(v) / si_conversion

        opening = ifcopenshell.api.root.create_entity(
            self.model,
            ifc_class="IfcOpeningElement",
            predefined_type="OPENING",
            name="Opening",
        )

        position_3d = None
        if self.model.schema == "IFC2X3":
            position_3d = self.model.createIfcAxis2Placement2D(self.model.createIfcCartesianPoint([0.0, 0.0, 0.0]))
        position_2d = self.model.createIfcAxis2Placement2D(
            self.model.createIfcCartesianPoint([door.OverallWidth / 2.0, 0.0])
        )

        opening.Representation = self.model.createIfcProductDefinitionShape(
            Representations=[
                self.model.createIfcShapeRepresentation(
                    body,
                    "Body",
                    "SweptSolid",
                    Items=[
                        self.model.createIfcExtrudedAreaSolid(
                            self.model.createIfcRectangleProfileDef(
                                "AREA",
                                None,
                                position_2d,
                                door.OverallWidth,
                                1.2 / si_conversion,
                            ),
                            position_3d,
                            self.model.createIfcDirection((0.0, 0.0, 1.0)),
                            door.OverallHeight,
                        )
                    ],
                )
            ]
        )

        ifcopenshell.api.feature.add_feature(self.model, feature=opening, element=wall)
        ifcopenshell.api.feature.add_filling(self.model, opening=opening, element=door)

        z_offsets = {"door": 0, "window": 1}
        opening.ObjectPlacement = self.model.createIfcLocalPlacement(
            wall.ObjectPlacement,
            self.model.createIfcAxis2Placement3D(
                self.model.createIfcCartesianPoint((float(t), 0.0, z_offsets[ty] / si_conversion))
            ),
        )

        door.ObjectPlacement = self.model.createIfcLocalPlacement(
            opening.ObjectPlacement,
            self.model.createIfcAxis2Placement3D(self.model.createIfcCartesianPoint((0.0, 0.0, 0.0))),
        )

        return door

    def to_obj_file(self, fn):
        st = ifcopenshell.geom.settings(USE_WORLD_COORDS=True, WELD_VERTICES=False)
        it = ifcopenshell.geom.iterator(st, self.model, exclude=("IfcOpeningElement",))
        sr = ifcopenshell.geom.serializers.obj(fn, fn + ".mtl", st, ifcopenshell.geom.serializer_settings())
        if it.initialize():
            for el in ifcopenshell.geom.consume_iterator(it):
                sr.write(el)
            sr.finalize()


if __name__ == "__main__":
    m = Context()
    w1 = m.create_2pt_wall((0.0, 0.0), (4.0, 0.0), 0.0, 3.0, 0.1, m.storey)
    w2 = m.create_2pt_wall((1.0, 3.0), (1.0, 0.0), 0.0, 3.0, 0.1, m.storey)
    m.create_fill("door", [2, 0.0], w1)
    m.create_fill("window", [1, 1.5], w2)
    m.model.write("out.ifc")
    m.to_obj_file("out.obj")