'''
See more here: http://www.pymolwiki.org/index.php/wfmesh
###############################################
#  File:          WFObj.py
#  Author:        Dan Kulp
#  Creation Date: 5/13/05
#
#  Notes:
#  Create openGL objects from a wavefront (obj) file
###############################################
'''

from __future__ import print_function

import os
import re
import math
from OpenGL.GL import *
from pymol.callback import Callback
from pymol import cmd

# Wrapper Function, to create a given WFObj with a specific name (flip = 1 if OpenFX + Crossroads used)


def createWFObj(file, name, translate=(0, 0, 0), flip=0):
    obj = WFMesh(file, translate, flip)
    cmd.load_callback(obj, name)


# Class for Wavefront Mesh
class WFMesh(Callback):

    verts = []        # list of vertices
    polys = []        # list of poylgons
    pnorms = []       # list of polynomal normals
    vnorms = {}       # dict. of vertex normals
    vavenorms = []    # list of vertex normals, redundant -- I'm far froma python pro.
    sections = {}     # list of sections of mesh

    # Read mesh into memory
    def readOBJ(self, file):
        if os.path.exists(file):
            input = open(file, 'r')
            for line in input:
                dat = re.split("\s+", line)

                # Find vertex line
                if line[0] == 'v' and line[1] != 't' and line[1] != 'n':
                    self.verts.append([dat[1], dat[2], dat[3]])

                # Find polygon line
                if line[0] == 'f':
                    self.polys.append([dat[1], dat[2], dat[3]])

                # Find section line
                if line[0] == 'g':
                    self.sections[len(self.polys)] = dat[1]

    # Compute the normals for each polygon and each vertex
    def computeNorms(self):

        # Compute norms for each polygon
        for p in self.polys:
            v12 = [float(self.verts[int(p[1]) - 1][0]) - float(self.verts[int(p[0]) - 1][0]),
                   float(self.verts[int(p[1]) - 1][1]) - float(self.verts[int(p[0]) - 1][1]),
                   float(self.verts[int(p[1]) - 1][2]) - float(self.verts[int(p[0]) - 1][2])
                   ]

            v13 = [float(self.verts[int(p[2]) - 1][0]) - float(self.verts[int(p[0]) - 1][0]),
                   float(self.verts[int(p[2]) - 1][1]) - float(self.verts[int(p[0]) - 1][1]),
                   float(self.verts[int(p[2]) - 1][2]) - float(self.verts[int(p[0]) - 1][2])
                   ]

            # Compute poly normal
            polynorm = self.cross(v12, v13)
            norm = self.normalize(polynorm)

            # Files created by OpenFX, Crossroads combination need have their normals flipped
            if self.flip:
                norm[0] = -norm[0]
                norm[1] = -norm[1]
                norm[2] = -norm[2]

            # Append poly norm to polygonal norm array
            self.pnorms.append(norm)

            # Add norm to each vertexes norm..
            try:
                self.vnorms[int(p[0]) - 1] = [float(self.vnorms[int(p[0]) - 1][0]) + norm[0],
                                            float(self.vnorms[int(p[0]) - 1][1]) + norm[1],
                                            float(self.vnorms[int(p[0]) - 1][2]) + norm[2]
                                            ]
            except:
                self.vnorms[int(p[0]) - 1] = [norm[0], norm[1], norm[2]]

            try:
                self.vnorms[int(p[1]) - 1] = [float(self.vnorms[int(p[1]) - 1][0]) + norm[0],
                                             float(self.vnorms[int(p[1]) - 1][1]) + norm[1],
                                             float(self.vnorms[int(p[1]) - 1][2]) + norm[2]
                                             ]
            except:
                self.vnorms[int(p[1]) - 1] = [norm[0], norm[1], norm[2]]

            try:
                self.vnorms[int(p[2]) - 1] = [float(self.vnorms[int(p[1]) - 1][0]) + norm[0],
                                             float(self.vnorms[int(p[1]) - 1][1]) + norm[1],
                                             float(self.vnorms[int(p[1]) - 1][2]) + norm[2]
                                             ]
            except:
                self.vnorms[int(p[2]) - 1] = [norm[0], norm[1], norm[2]]

        # Average out each vnorm..
        index = 0
        for v in self.vnorms.values():
            self.vavenorms.append([v[0] / 4, v[1] / 4, v[2] / 4])
            index += 1

    # Utility function to normalize a given vector
    def normalize(self, v):
        mag = v[0] * v[0] + v[1] * v[1] + v[2] * v[2]
        if mag <= 0:
            mag = 1
        else:
            mag = math.sqrt(mag)

        return [v[0] / mag, v[1] / mag, v[2] / mag]

    # Utility cross product function
    def cross(self, v1, v2):
        x = 0
        y = 1
        z = 2

        return [v1[y] * v2[z] - v1[z] * v2[y],
                v1[z] * v2[x] - v1[x] * v2[z],
                v1[x] * v2[y] - v1[y] * v2[x]
                ]

    # Constructor
    def __init__(self, file, translate=(0, 0, 0), flip=0):
        self.verts = []
        self.polys = []
        self.pnorms = []
        self.vnorms = {}
        self.vavenorms = []
        self.translate = translate
        self.flip = flip

        print("Read in file: " + str(file))
        self.readOBJ(file)
        print("Done reading in WFMesh, now compute norms")
        self.computeNorms()
        print("Done computing norms, now display WFMesh")

    # Draw Function
    def __call__(self):

        glColorMaterial(GL_FRONT, GL_DIFFUSE)
        glEnable(GL_COLOR_MATERIAL)
        glShadeModel(GL_SMOOTH)

        # Color Everything grey
        glColor3f(0.5, 0.5, 0.5)

        glPushMatrix()
        glTranslated(self.translate[0], self.translate[1], self.translate[2])
        for index, p in enumerate(self.polys):
            glBegin(GL_POLYGON)
            glNormal3f(float(self.pnorms[index][0]), float(self.pnorms[index][1]), float(self.pnorms[index][2]))

            for x in p:
                glVertex3f(float(self.verts[int(x) - 1][0]), float(self.verts[int(x) - 1][1]), float(self.verts[int(x) - 1][2]))

                # Vertex Normals - not computed correctly, so commented out for now
#                       norm = self.vnorms[int(x)-1]
#                       glNormal3f(float(norm[0]),float(norm[1]),float(norm[2]))
            glEnd()
        glPopMatrix()
cmd.extend("createWFObj", createWFObj)