#!/usr/bin/env python

# /*
# (c) 2012 +++ Filip Stoklas, aka FipS, http://www.4FipS.com +++
# THIS CODE IS FREE - LICENSED UNDER THE MIT LICENSE
# ARTICLE URL: http://forums.4fips.com/viewtopic.php?f=3&t=1047
# */

import sys
sys.path.append("../imports")
sys.path.append("../messages")

import geometry_pb2
import collada
import struct

def expand_positions(prim):
    num_tris = len(prim.vertex_index)
    positions = []
    if len(prim.vertex):
        for i in range(0, num_tris):
            for v in range (0, 3):
                pos = prim.vertex[prim.vertex_index[i]][v]
                positions.append(struct.pack("fff", *pos))
    return positions

def expand_normals(prim):
    num_tris = len(prim.normal_index)
    normals = []
    if len(prim.normal):
        for i in range(0, num_tris):
            for v in range (0, 3):
                norm = prim.normal[prim.normal_index[i]][v]
                normals.append(struct.pack("fff", *norm))
    return normals

def expand_texcoords(prim, index):
    if index >= len(prim.texcoord_indexset):
        return []
    num_tris = len(prim.texcoord_indexset[index])
    texcoords = []
    if len(prim.texcoordset[index]):
        for i in range(0, num_tris):
            for v in range (0, 3):
                texcoord = prim.texcoordset[index][prim.texcoord_indexset[index][i]][v]
                texcoords.append(struct.pack("ff", *texcoord))
    return texcoords

def create_vertex_data(positions, normals, texcoords0, texcoords1):
    vertex_data = ""
    for i in range(0, len(positions)):
        vertex_data += positions[i] if positions else ""
        vertex_data += normals[i] if normals else ""
        vertex_data += texcoords0[i] if texcoords0 else ""
        vertex_data += texcoords1[i] if texcoords1 else ""
    return vertex_data

def export_geometry(vertex_data, has_positions, has_normals, has_texcoords0, has_texcoords1):
    geom = geometry_pb2.Geometry()
    if has_positions:
        e = geom.format.elements.add()
        e.type = geometry_pb2.VET_Float_3
        e.semantics = geometry_pb2.VES_Position
    if has_normals:
        e = geom.format.elements.add()
        e.type = geometry_pb2.VET_Float_3
        e.semantics = geometry_pb2.VES_Normal
    if has_texcoords0:
        e = geom.format.elements.add()
        e.type = geometry_pb2.VET_Float_2
        e.semantics = geometry_pb2.VES_Texcoord
    geom.data = vertex_data
    with open("geometry.bin", "wb") as f:
        f.write(geom.SerializeToString())

def inspect_geometry(obj):
    print "GEOMETRY: '%s'" % obj.original.name
    for prim in obj.primitives():
        print "- PRIMITIVE: type: '%s', triangles: %d, vertices: %d, ..." % (type(prim).__name__, len(prim), len(prim.vertex))
        if type(prim) is not collada.triangleset.BoundTriangleSet:
            continue
        num_tris = len(prim.vertex_index)
        positions = expand_positions(prim)
        normals = expand_normals(prim)
        texcoords0 = expand_texcoords(prim, 0)
        texcoords1 = expand_texcoords(prim, 1)
        print "- SOUP: triangles: %d (vertices: %d, ...)" % (num_tris, len(positions))
        vertex_data = create_vertex_data(positions, normals, texcoords0, texcoords1)
        export_geometry(vertex_data, bool(positions), bool(normals), bool(texcoords0), bool(texcoords1))

def inspect_collada(col):
    if col.scene:
        for geom in col.scene.objects("geometry"):
            inspect_geometry(geom)

if __name__ == "__main__":
    filename = sys.argv[1]
    col = collada.Collada(filename, ignore=[collada.DaeUnsupportedError, collada.DaeBrokenRefError])
    inspect_collada(col)