from typing import Optional

import ifcopenshell
import ifcopenshell.util.element as el
import ifcopenshell.util.unit as unit
import numpy as np
import pyvista as pv
import triangle


def run(f: ifcopenshell.file, target_density: float, round_digits: Optional[int]) -> ifcopenshell.file:
    scale = unit.calculate_unit_scale(f)
    facesets = f.by_type("IfcConnectedFaceSet")
    if f.schema.upper().startswith("IFC4"):
        facesets += f.by_type("IfcTessellatedFaceSet") 
    el.batch_remove_deep2(f)
    for fs in facesets:
        simplify(scale, f, fs, round_digits, target_density)
    f = el.unbatch_remove_deep2(f)
    return f


def read(fs, nd):
    def triangulate(rs, holes=()):
        if any(len(r) < 3 or len(np.unique(np.asarray(r), axis=0)) < len(r) for r in rs):
            return None

        vs, ss = [], []
        for r in rs:
            n = len(vs)
            vs.extend(r)
            ss.extend((n + i, n + (i + 1) % len(r)) for i in range(len(r)))

        mesh = {
            "vertices": np.array(vs),
            "segments": np.array(ss),
        }
        if holes:
            mesh["holes"] = np.array(holes)

        try:
            mesh = triangle.triangulate(mesh, "p")
        except Exception:
            return None

        vs, ts = mesh.get("vertices"), mesh.get("triangles")
        return None if vs is None or ts is None else (vs, ts)

    def points(fs):
        vs = np.array(fs.Coordinates.CoordList, dtype=float)
        return vs[np.array(fs.PnIndex, dtype=np.int64) - 1] if fs.PnIndex else vs

    def poly(os, hs):
        def frame(ps):
            def normalize(v):
                n = np.linalg.norm(v)
                return None if n < 1e-12 else v / n

            o = ps[0]
            es = np.roll(ps, -1, axis=0) - ps

            for i, a in enumerate(es):
                u = normalize(a)
                if u is None:
                    continue
                for b in es[i + 1 :]:
                    n = normalize(np.cross(a, b))
                    if n is not None:
                        return o, np.array((u, np.cross(n, u), n))

            return None

        fr = frame(os)
        if fr is None:
            return None
        o, m = fr
        rs = [[((p - o) @ m.T)[:2] for p in r] for r in [os] + hs]

        def hole_tri_centroid(r):
            if mesh := triangulate([r]):
                return np.mean(mesh[0][mesh[1][0]], axis=0)

        hs = [h for r in rs[1:] if (h := hole_tri_centroid(r)) is not None]
        if len(hs) != len(rs) - 1:
            return None

        mesh = triangulate(rs, hs)
        if mesh is None:
            return None
        vs, gs = mesh

        ts = []
        for t in gs:
            cs = [vs[i] for i in t]
            ts.append([o + c @ m[:2] for c in cs])
        return ts or None

    def indexed_face(fa, vs):
        os = vs[np.array(fa.CoordIndex, dtype=np.int64) - 1]
        hs = [vs[np.array(h, dtype=np.int64) - 1] for h in getattr(fa, "InnerCoordIndices", ()) or ()]
        return [os[[0, i, i + 1]] for i in range(1, len(os) - 1)] if not hs else poly(os, hs)

    def face(fa):
        def ring(b):
            if not b.Bound or not b.Bound.is_a("IfcPolyLoop"):
                return None
            ps = np.array([p.Coordinates for p in b.Bound.Polygon])
            if b.Orientation is False:
                ps = ps[::-1]
            return ps

        bs = list(fa.Bounds or ())
        if not bs:
            return None
        if len(bs) == 1:
            b = bs[0]
            if b.Bound and b.Bound.is_a("IfcPolyLoop") and len(b.Bound.Polygon) == 3:
                ps = list(b.Bound.Polygon)
                if b.Orientation is False:
                    ps.reverse()
                return [np.array([p.Coordinates for p in ps], dtype=float)]
        ob = next((b for b in bs if b.is_a("IfcFaceOuterBound")), bs[0])
        os = ring(ob)
        hs = [ring(b) for b in bs if b.id() != ob.id()]
        if os is None or len(os) == 0 or any(h is None for h in hs):
            return None
        if not hs and len(os) == 3:
            return [os]
        return poly(os, hs)

    if fs.is_a("IfcConnectedFaceSet"):
        rs = [face(fa) for fa in fs.CfsFaces]
        rs = None if any(r is None for r in rs) else [t for r in rs for t in r]
    elif fs.is_a("IfcTriangulatedFaceSet"):
        vs = points(fs)
        rs = None if len(vs) == 0 or not fs.CoordIndex else vs[np.array(fs.CoordIndex, dtype=np.int64) - 1]
    elif fs.is_a("IfcPolygonalFaceSet"):
        vs = points(fs)
        if len(vs) == 0 or not fs.Faces:
            return None
        rs = [indexed_face(fa, vs) for fa in fs.Faces]
        rs = None if any(r is None for r in rs) else [t for r in rs for t in r]
    else:
        return None

    if rs is None or len(rs) == 0:
        return None
    vs = np.asarray(rs, dtype=float).reshape((-1, 3))
    if nd is not None:
        vs = np.round(vs, nd)
    vs, ix = np.unique(vs, axis=0, return_inverse=True)
    return vs, ix.reshape((-1, 3)).astype(np.int64)


def decimate(vs, ts, r):
    fs = np.empty((len(ts), 4), dtype=np.int64)
    fs[:, 0] = 3
    fs[:, 1:] = ts
    m = pv.PolyData(vs, fs.ravel())
    if m.n_cells < 4:
        return None
    try:
        m = m.decimate_pro(r)
    except Exception:
        return None
    m = m.triangulate()
    fs = m.faces.reshape((-1, 4))
    if len(fs) == 0 or np.any(fs[:, 0] != 3):
        return None
    return np.array(m.points), np.array(fs[:, 1:], dtype=np.int64)


def write(f, src, vs, ts):
    def make_faces(f, vs, ts):
        ps = [f.create_entity("IfcCartesianPoint", Coordinates=tuple(map(float, p))) for p in vs]
        fs = []
        for t in ts:
            lp = f.create_entity("IfcPolyLoop", Polygon=[ps[int(i)] for i in t])
            bd = f.create_entity("IfcFaceOuterBound", Bound=lp, Orientation=True)
            fs.append(f.create_entity("IfcFace", Bounds=[bd]))
        return tuple(fs)

    def make_tfs(f, src, vs, ts):
        ps = f.create_entity("IfcCartesianPointList3D", CoordList=[tuple(map(float, p)) for p in vs])
        kw = {
            "Coordinates": ps,
            "CoordIndex": [tuple(int(i) + 1 for i in t) for t in ts],
        }
        if src.is_a("IfcTessellatedFaceSet") and src.Closed is not None:
            kw["Closed"] = src.Closed
        return f.create_entity("IfcTriangulatedFaceSet", **kw)

    def deps(f, xs):
        ds = {}
        for x in xs:
            for y in f.traverse(x):
                if y.id():
                    ds[y.id()] = y
        return list(ds.values())

    def clean(f, xs):
        ids = {x.id() for x in xs}
        rm = []
        for x in xs:
            try:
                if any(y.id() not in ids for y in f.get_inverse(x)):
                    continue
            except RuntimeError:
                continue
            if x.id():
                rm.append(x)
        if f.to_delete is not None:
            f.to_delete.update(rm)
            return
        f.batch()
        try:
            for x in sorted(rm, key=lambda y: y.id(), reverse=True):
                f.remove(x)
        finally:
            f.unbatch()

    if f.schema.upper().startswith("IFC4"):
        dst = make_tfs(f, src, vs, ts)
        old = deps(f, (src,))
        el.replace_element(src, dst)
        clean(f, old)
    else:
        old = deps(f, tuple(src.CfsFaces))
        src.CfsFaces = make_faces(f, vs, ts)
        clean(f, old)


def simplify(scale, f, fs, round_digits, target_density):
    rs = read(fs, round_digits)
    if rs is None:
        return None
    vs, ts = rs

    def triangle_area(vs, ts):
        tris = vs[ts]
        ab = tris[:, 1] - tris[:, 0]
        ac = tris[:, 2] - tris[:, 0]
        return 0.5 * np.linalg.norm(np.cross(ab, ac), axis=1).sum()

    area = triangle_area(vs, ts) * scale * scale
    num_verts = vs.shape[0]
    density = num_verts / area

    factor = 1.0 - (target_density / density)

    if len(ts) < 32:
        return None
    rs = decimate(vs, ts, factor)
    if rs is None:
        return None
    vs, ts2 = rs
    if len(ts2) >= len(ts):
        return None

    write(f, fs, vs, ts2)
    return len(ts), len(ts2)