import numpy as np
import math
from MeshCleaner import MeshCleaner


IN_PATH = 'test/orig_clean.obj'
OUT_PATH = 'test/seg.obj'

THRESH_ANGLE = 55
PERCENTAGE = .7


class Point:
    def __init__(self, i):
        self.i = i
        self.norm = []
        self.t = []
        self.n = []
        self.angle = None

    def add_neighbor(self, n):
        if n not in self.n:
            self.n.append(n)

    def add_tri(self, t):
        if t not in self.t:
            self.t.append(t)

    def add_normal(self, n):
        self.norm.append(n)


class FileFormater:
    def __init__(self):
        self.input = open(IN_PATH, 'r')
        self.t = []
        self.x = []
        self.y = []
        self.z = []
        self.p = []
        self.center = []
        self.norm = [0., 0., 0.]

    def format(self):
        line = self.input.readline()
        while line[0] != 'v':
            line = self.input.readline()

        xavg = 0.
        yavg = 0.
        zavg = 0.
        v = 0
        while line[0] == 'v':
            words = line.split()
            x = float(words[1])
            y = float(words[2])
            z = float(words[3])
            xavg += x
            yavg += y
            zavg += z
            self.x.append(x)
            self.y.append(y)
            self.z.append(z)
            self.p.append(Point(v))
            line = self.input.readline()
            v += 1

        self.center.append(xavg / v)
        self.center.append(yavg / v)
        self.center.append(zavg / v)

        while line[0] != 'f':
            line = self.input.readline()

        t = 0
        while line[0] == 'f':
            words = line.split()
            x = int(words[1]) - 1
            y = int(words[2]) - 1
            z = int(words[3]) - 1
            self.t.append((x, y, z))
            self.p[x].add_neighbor(y)
            self.p[x].add_neighbor(z)
            self.p[y].add_neighbor(x)
            self.p[y].add_neighbor(z)
            self.p[z].add_neighbor(y)
            self.p[z].add_neighbor(x)

            u = [self.x[y] - self.x[x], self.y[y] - self.y[x],
                 self.z[y] - self.z[x]]
            v = [self.x[z] - self.x[x], self.y[z] - self.y[x],
                 self.z[z] - self.z[x]]

            normal = []
            a = (u[1] * v[2]) - (u[2] * v[1])
            b = (u[2] * v[0]) - (u[0] * v[2])
            c = (u[0] * v[1]) - (u[1] * v[0])
            self.norm[0] += a
            self.norm[1] += b
            self.norm[2] += c
            normal.append(a)
            normal.append(b)
            normal.append(c)
            self.p[x].add_normal(normal)
            self.p[y].add_normal(normal)
            self.p[z].add_normal(normal)
            self.p[x].add_tri(t)
            self.p[y].add_tri(t)
            self.p[z].add_tri(t)
            line = self.input.readline()
            t += 1

        self.norm[0] /= t
        self.norm[1] /= t
        self.norm[2] /= t

        self.input.close()


def main():
    mc = MeshCleaner('test/orig.obj', 'test/orig_clean.obj')
    mc.clean()
    ff = FileFormater()
    ff.format()
    x = np.r_[ff.x]
    y = np.r_[ff.y]
    z = np.r_[ff.z]

    f = open(OUT_PATH, "w")

    for i in range(len(x)):
        f.write("v {0} {1} {2} 1.0 0.0 0.0\n".format(x[i], y[i], z[i]))

    n1 = np.array(ff.norm)
    mag1 = (n1[0]**2 + n1[1]**2 + n1[2]**2)**.5
    # tri = np.zeros(0, dtype=np.int)
    tri = []

    for p in ff.p:
        if len(p.norm) != 0:
            pavg = [0., 0., 0.]
            for i in range(len(p.norm)):
                pavg[0] += p.norm[i][0]
                pavg[1] += p.norm[i][1]
                pavg[2] += p.norm[i][2]

            pavg[0] /= len(p.norm)
            pavg[1] /= len(p.norm)
            pavg[2] /= len(p.norm)
            n2 = np.array(pavg)

            p.norm = pavg
            numer = np.dot(n1, n2)
            mag2 = (n2[0]**2 + n2[1]**2 + n2[2]**2)**.5
            denom = mag1 * mag2
            p.angle = math.degrees(math.acos(numer / denom))

    for p in ff.p:
        d1 = (ff.x[p.i] - ff.center[0])**2 + \
             (ff.y[p.i] - ff.center[1])**2 + \
             (ff.z[p.i] - ff.center[2])**2
        d2 = ((ff.x[p.i] + p.norm[0]) - ff.center[0])**2 + \
             ((ff.y[p.i] + p.norm[1]) - ff.center[1])**2 + \
             ((ff.z[p.i] + p.norm[2]) - ff.center[2])**2
        if d1 - d2 <= 0:
            pass
        elif p.angle > THRESH_ANGLE:
            for t in p.t:
                tri.append(t)
        else:
            c1 = 0.
            c2 = 0.
            c3 = 0.
            c4 = 0.
            t1 = 0
            t2 = 0
            t3 = 0
            t4 = 0
            for n in p.n:
                if ff.p[n].angle > THRESH_ANGLE - 45:
                    c1 += 1
                for nn in ff.p[n].n:
                    if ff.p[nn].angle > THRESH_ANGLE - 35:
                        c2 += 1
                    for nnn in ff.p[nn].n:
                        if ff.p[nnn].angle > THRESH_ANGLE - 25:
                            c3 += 1
                        for nnnn in ff.p[nnn].n:
                            if ff.p[nnnn].angle > THRESH_ANGLE - 15:
                                c4 += 1
                        t4 += len(ff.p[nnn].n)
                    t3 += len(ff.p[nn].n)
                t2 += len(ff.p[n].n)
            t1 += len(p.n)

            weight = (.2 * (c1 / t1) + .3 * (c2 / t2)) + \
                     (.9 * (c3 / t3) + .9 * (c4 / t4))
            if (weight > PERCENTAGE):
                # print ("w: %.2f, 1: %.2f, 2: %.2f, 3: %.2f, 4: %.2f" % (weight, c1 / t1, c2 / t2, c3 / t3, c4 / t4))
                for t in p.t:
                    tri.append(t)

    for t in tri:
        f.write("f {0} {1} {2}\n".format(ff.t[t][0] + 1,
                                         ff.t[t][1] + 1,
                                         ff.t[t][2] + 1))

    f.write('\n')
    f.close()
    mc = MeshCleaner('test/seg.obj', 'test/seg_clean.obj',
                     out_color='1.0 0.0 0.0')
    mc.clean()

if __name__ == "__main__":
    main()
