"""
OBJ to XMODEL by Scobalula

Converts Wavefront OBJ files to XMODEL_EXPORT
"""
import struct
import os
import sys
import traceback
import time
import ConfigParser
import lz4
import pycod.xmodel as xModel
import numpy


class WavefrontOBJ(object):
    """
    Class to handle OBJ Files
    """
    # Needs some work since there's like 50000
    # variations of OBJ from what I have gathered..
    # NOTE: Only supports Quads/Tris. any more than
    # 4 indices and import won't work.
    def __init__(self, obj_file=None, triangluate=True):
        """
        Init. OBJ Data

        obj_file = filename (String) - File to import
        triangulate = True/False (Boolean) - Triangulate Quads

        """
        self.file = obj_file
        self.verts = []
        self.uvs = []
        self.normals = []
        self.materials = []
        self.faces = []
        self.triangluate = triangluate
        # Load OBJ
        if obj_file is not None:
            self.load_obj()

    def load_obj(self):
        """
        Loads OBJ
        """
        name = os.path.basename(os.path.splitext(self.file)[0])
        with open(self.file, 'r') as obj:
            # Set the initial material index.
            material_index = -1
            for line in obj:
                # Split lines
                data = line.replace("\n", "").split(" ")
                # Retards storing random spaces
                data[:] = [item for item in data if item != '']
                # Empty lines
                if not data:
                    continue
                # Vertex Offsets
                if data[0] == "v":
                    x = float(data[1])
                    y = float(data[2])
                    z = float(data[3])
                    self.verts.append((x,
                                       y,
                                       z))
                # UV Data
                elif data[0] == "vt":
                    u = float(data[1])
                    v = float(data[2])
                    self.uvs.append((u,
                                     v))
                # Vertex Normals
                elif data[0] == "vn":
                    x = float(data[1])
                    y = float(data[2])
                    z = float(data[3])
                    self.normals.append((x,
                                         y,
                                         z))
                # Materials
                elif data[0] == "usemtl" or data[0] == "g":
                    material_index += 1
                    try:
                        self.materials.append(data[1])
                    except:
                        self.material.append("mtl_%i" % material_index)
                # Faces
                elif data[0] == "f":
                    mtl = material_index
                    face_info = []
                    for i in range(1, len(data)):
                        face_info.append(data[i].strip().split("/"))
                    if len(data) == 5 and self.triangluate:
                        arrangement = []
                        # Index arrangement for new faces
                        arrangement.append([0, 1, 3])
                        arrangement.append([1, 2, 3])
                        # Create 2 new faces
                        for i in range(2):
                            face_data = []
                            # Add 3 new indices
                            for j in range(3):
                                # Init. Vars
                                texture = None
                                uv = None
                                normal = None
                                # Vertex
                                if len(info) == 1:
                                    vertex = int(face_info[arrangement[i][j]][0])
                                # Texture
                                if len(info) > 1:
                                    try:
                                        texture = int(face_info[arrangement[i][j]][1])
                                    except:
                                        texture = None
                                # Normal
                                if len(info) > 2:
                                    try:
                                        normal = int(face_info[arrangement[i][j]][2])
                                    except:
                                        normal = None

                                face_data.append([vertex, texture, normal])
                            # No UV info provided, skip.
                            if texture is None:
                                continue
                            self.faces.append([face_data, 0, mtl])
                    else:
                        face_data = []
                        for i in range(1, len(data)):
                            info = face_info[i - 1]
                            # Init. Vars
                            texture = None
                            uv = None
                            normal = None
                            # Vertex
                            if len(info) > 0:
                                vertex = int(info[0])
                            # UV
                            if len(info) > 1:
                                try:
                                    texture = int(info[1])
                                except:
                                    texture = None
                            # Normal
                            if len(info) > 2:
                                try:
                                    normal = int(info[2])
                                except:
                                    normal = None
                            face_data.append([vertex, texture, normal])
                        # No UV info provided, skip.
                        if texture is None:
                            continue
                        self.faces.append([face_data, 0, mtl])

        if not self.normals:
            self.normal = False


    def write_obj(self, file):
        file_path = os.path.splitext(file)[0]
        i = 0
        # Check if it exists already.
        while os.path.exists(file_path + ".obj"):
            file_path = file_path + str(i) + ".obj"
            i += 1
        
        with open(file, 'w') as obj:
            obj.write("# File written by Scobalula\'s OBJ Converter\n")
            obj.write("# Source file %s\n\n" % file)
            obj.write("g default\n")
            # Used materials
            used_materials = []
            # Write Verts/Offsets
            for vert in self.verts:
                obj.write("v %f %f %f\n" % (vert[0],
                                            vert[1],
                                            vert[2]))
            # Write UVs
            for uv in self.uvs:
                obj.write("vt %f %f\n" % (uv[0],
                                         uv[1]))
            # Write Normals
            for normal in self.normals:
                obj.write("vn %f %f %f\n" % (normal[0],
                                             normal[1],
                                             normal[2]))
            # Write Faces
            currentmtl = 0
            for face in self.faces:
                face_data = face[0]
                mtl_index = face[2]
                if self.materials[mtl_index] not in used_materials:
                    used_materials.append(self.materials[mtl_index])
                
                if mtl_index != currentmtl:
                    obj.write("usemtl %s\n" % self.materials[mtl_index])
                    currentmtl = mtl_index
                

                obj.write("f")
                for i in range(len(face_data)):
                    # at least have vert data
                    vert = str(face_data[i][0])
                    try:
                        uv = str(face_data[i][1])
                    except:
                        uv = ""
                    try:
                        normal = str(face_data[i][2])
                    except:
                        normal = ""
                    obj.write(" %s/%s/%s" % (vert, normal, uv))
                obj.write("\n")




def normalize_v3(arr):
    ''' Normalize a numpy array of 3 component vectors shape=(n,3) '''
    lens = numpy.sqrt( arr[:,0]**2 + arr[:,1]**2 + arr[:,2]**2 )
    arr[:,0] /= lens
    arr[:,1] /= lens
    arr[:,2] /= lens                
    return arr

def calculate_normal(faces, verts):
    """
    Recalculates Normals
    https://sites.google.com/site/dlampetest/python/calculating-normals-of-a-triangle-mesh-using-numpy
    """
    face_info = []
    for face in faces:
        face_data = []
        for i in range(3):
            face_data.append(face[0][i][0] - 1)
        face_info.append(face_data)
    
    faces = numpy.array(face_info)
    verts = numpy.array(verts)
    norm = numpy.zeros( verts.shape, dtype=verts.dtype )
    tris = verts[faces]
    n = numpy.cross( tris[::,1 ] - tris[::,0], tris[::,2 ] - tris[::,0] )
    normalize_v3(n)
    norm[ faces[:,0] ] += n
    norm[ faces[:,2] ] += n
    norm[ faces[:,1] ] += n
    normalize_v3(norm)
    no = norm[ faces ]
    return no

def export_xmodel(data,
                  xmodel_name,
                  upaxis="Y",
                  scale=1,
                  filemtlname=False,
                  fileformat="BOTH",
                  recalculate_normals=False):
    """
    Passes info to PyCod and writes
    """
    model = xModel.Model()
    # Create Tag_Origin
    tag_origin = xModel.Bone("tag_origin", -1)
    tag_origin.offset = (0, 0, 0)
    tag_origin.matrix = [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
    model.bones.append(tag_origin)
    mesh = xModel.Mesh("name")
    model.meshes.append(mesh)
    recalculated = False
    if not data.normals or recalculate_normals:
        print("<--- Calculating normals...")
        recalculated = True
        data.normals = calculate_normal(data.faces, data.verts).tolist()
    print("<--- Exporting data...")
    # Add verts
    for vert_offset in data.verts:
        vert = xModel.Vertex()
        # Flip the axis
        if upaxis == "Z":
            # Scale it
            x = vert_offset[0] * float(scale)
            y = (vert_offset[2] * (-1)) * float(scale)
            z = vert_offset[1] * float(scale)
        # Maintain axis.
        else:
            # Scale it.
            x = vert_offset[0] * float(scale)
            y = vert_offset[1] * float(scale)
            z = vert_offset[2] * float(scale)
        vert.offset = (x, y, z)
        # Set weights to tag_origin
        vert.weights.append((0, 1.00000))
        mesh.verts.append(vert)
    # Materials we've used.
    used_materials = []
    for j, tri in enumerate(data.faces):
        face = xModel.Face(0, 0)
        # Check -1 material index.
        if tri[2] < 0:
            tri[2] = 0
        # Check against current materials
        # if it's a new mtl, add it.
        if tri[2] not in used_materials:
            used_materials.append(tri[2])
        # Set the ID based off index in list.
        face.material_id = used_materials.index(tri[2])
        # Set mesh ID (tie it all to 0)
        face.mesh_id = 0
        verts = []
        # if not data.normals:
        #     calculated_normal = calculate_normal(data.verts[int(tri[0][0][0]) - 1],
        #                                          data.verts[int(tri[0][1][0]) - 1],
        #                                          data.verts[int(tri[0][2][0]) - 1])
        for i in range(3):
            vert = int(tri[0][i][0]) - 1
            uvs = data.uvs[int(tri[0][i][1]) - 1]
            # 1 - V value for CoD
            uv = (uvs[0], 1 - uvs[1])
            normal_index = tri[0][i][2]
            # Recalculated normals' index will
            # be same as face index.
            if recalculated:
                normal = data.normals[j][i]
            else:
                normal = data.normals[normal_index - 1]
            if upaxis == "Z":
                x = normal[0]
                y = (normal[2] * (-1))
                z = normal[1]
            # Keep them
            else:
                x = normal[0]
                y = normal[1]
                z = normal[2]
            normal = (x, y, z)
            color = (1, 1, 1, 1)
            faceVertex = xModel.FaceVertex(vert, normal, color, uv)
            face.indices[i] = faceVertex
        # 0 2 1 for Call of Duty
        # Rearragne the verts for normal up axis.
        v2 = face.indices[2]
        face.indices[2] = face.indices[1]
        face.indices[1] = v2
        # Add it to mesh.
        model.meshes[0].faces.append(face)
    # Export only used materials.
    for i, material in enumerate(used_materials):
        imgs = {"color": "", "normal": "", "s": ""}
        name = name = os.path.basename(os.path.splitext(data.file)[0])
        if filemtlname:
            name = ("mtl_%s_%i" % (name.lower(), i))
        else:
            name = data.materials[material].lower()
        mtl = xModel.Material(name,
                              "Lambert",
                              imgs)
        model.materials.append(mtl)
    # Write it.
    header = "// Exported via OBJ2XModel by Scobalula using PyCod\n"
    if fileformat == "BIN":
        model.WriteFile_Bin(xmodel_name.lower() + ".XMODEL_BIN")
    elif fileformat == "EXPORT":
        model.WriteFile_Raw(xmodel_name.lower() + ".XMODEL_EXPORT", version=6, header_message=header)
    if fileformat == "BOTH":
        model.WriteFile_Raw(xmodel_name.lower() + ".XMODEL_EXPORT", version=6, header_message=header)
        model.WriteFile_Bin(xmodel_name.lower() + ".XMODEL_BIN")

def load_xmodel(file, extension):
    """
    Loads xmodel from PyCod and adds info to OBJ
    """
    xmodel = xModel.Model()
    # Combine meshes for now.
    if extension.lower() == ".xmodel_bin":
        xmodel.LoadFile_Bin(file, split_meshes=False)
    else:
        xmodel.LoadFile_Raw(file, split_meshes=False)
    # Initiate OBJ
    obj = WavefrontOBJ()
    # Loop over each mesh
    for mesh in xmodel.meshes:
        # Add verts.
        for verts in mesh.verts:
            obj.verts.append(verts.offset)
        # Index for UV/Normals
        index = 1
        # Loop over faces.
        for faces in mesh.faces:
            # Face order fix
            v2 = faces.indices[2]
            faces.indices[2] = faces.indices[1]
            faces.indices[1] = v2
            # Face data.
            face_data = []
            for v in faces.indices:
                # Normals
                obj.normals.append(v.normal)
                # UV, 1 - V value from CoD
                uv = (v.uv[0], 1 - v.uv[1])
                obj.uvs.append((uv))
                # Add the face data
                face_data.append([v.vertex + 1, index, index])
                index += 1
            # Add the face information
            obj.faces.append([face_data, 0, faces.material_id])
    # Add materials.
    for material in xmodel.materials:
        obj.materials.append(material.name)
    # Return the obj object.
    return obj
        
def getsetting(setting,
               settingsFile=os.path.join(os.path.dirname(
                                           os.path.realpath(
                                               sys.argv[0])),
                                         "settings.ini")):
    """
    Gets setting from settngs file
    """
    # Check if our settings file exists.
    if not os.path.isfile(settingsFile):
        return None

    # Set config object.
    Config = ConfigParser.ConfigParser()

    # Read config file.
    Config.read(settingsFile)

    # Try read the setting.
    try:
        # Get it.
        value = Config.get("OBJConvert", setting)
        # Return userdefined.
        return value
    except Exception as e:
        return None


def main(files):
    """
    Main Function to handle passing files
    """
    print("<--- OBJ to XMODEL_EXPORT by Scobalula\n")

    # Get Scale
    scale = getsetting("scale")
    # Didn't get expected.
    if scale is None:
        scale = 1
    # Can't be casted to float.
    try:
        scale = float(scale)
    except:
        scale =1

    # Get upaxis
    upaxis = getsetting("upaxis")
    if upaxis is None:
        upaxis = "Z"
    if upaxis.upper() not in ["Y", "Z", "DEBUG"]:
        upaxis = "Z"

    # Get use file name for MTL
    filemtl = getsetting("filemtl")
    try:
        filemtl = int(filemtl)
    except:
        filemtl = 0
    if filemtl is None:
        filemtl = 0
    if filemtl > 1 and filemtl < 0:
        filemtl = 0
    
        # Get use file name for MTL
    recalculate = getsetting("recalculate")
    try:
        recalculate = int(recalculate)
    except:
        recalculate = 0
    if recalculate is None:
        recalculate = 0
    if recalculate > 1 and recalculate < 0:
        recalculate = 0
        
    # Get file format
    fileformat = getsetting("format")
    if fileformat is None:
        fileformat = "BOTH"
    if fileformat not in ["BIN", "EXPORT", "BOTH"]:
        fileformat = "BOTH"
    
    # Delete program's entry
    del files[0]
    # No files given
    if not files:
        print("<--- No files given.")
    # Loop over each file.
    for file in files:
        # Split file name
        name = os.path.splitext(file)
        # Check extension
        if name[-1].lower() not in [".obj",
                                    ".xmodel_export",
                                    ".xmodel_bin"]:
            print("<--- %s unrecognized extension, skipping."
                  % file)
            continue
        # Loading
        print("<--- Loading %s...." % file)
        # Start time
        begin = time.time()
        is_xmodel = False
        # Try load it.
        try:
            # Check which to load.
            if name[-1].lower() in [".xmodel_export", ".xmodel_bin"]:
                # XMODEL to OBJ
                is_xmodel = True
                data = load_xmodel(file, name[-1])
            else:
                # OBJ to XMODEL
                data = WavefrontOBJ(file)
        except Exception as e:
            print("<--- An error occured loading %s, skipping."
                  % file)
            print("<--- Error:\n %s"
                  % str(str(traceback.format_exc())))
            continue
        # Loaded the file
        print("<--- Loaded %i Verts, "
              "%i Faces "
              "and %i Materials in %f seconds.\n"
              % (len(data.verts), len(data.faces), len(data.materials),
                 time.time() - begin))
        # Begin time for export.
        begin = time.time()
        print("<--- Writing %s...." %
              (os.path.splitext(file)[0]))
        # Try write it.
        try:
            # Check which to write.
            if is_xmodel:
                # Write obj (xmodel was given)
                data.write_obj(name[0] + ".obj")
            else:
                # Write xmodel (obj was given)
                export_xmodel(data,
                              os.path.splitext(file)[0],
                              upaxis=upaxis,
                              scale=scale,
                              filemtlname=filemtl,
                              fileformat=fileformat.upper(),
                              recalculate_normals=recalculate)
        except Exception as e:
            print("<--- An error occured exporting %s, skipping."
                  % file)
            print("<--- Error:\n %s" % str(traceback.format_exc()))
            continue
        # Exported fine.
        print("<--- Wrote model in %f seconds.\n"
              % (time.time() - begin))

if __name__ == "__main__":
    main(sys.argv)

raw_input("Execution completed, press enter to exit.")
