import io
import copy
import uuid

import numpy as np

from .. import util
from .. import visual

from ..constants import log


def load_collada(file_obj, resolver=None, **kwargs):
    """
    Load a COLLADA (.dae) file into a list of trimesh kwargs.

    Parameters
    ----------
    file_obj : file object
      Containing a COLLADA file
    resolver : trimesh.visual.Resolver or None
      For loading referenced files, like texture images
    kwargs : **
      Passed to trimesh.Trimesh.__init__

    Returns
    -------
    loaded : list of dict
      kwargs for Trimesh constructor
    """
    import collada

    # load scene using pycollada
    c = collada.Collada(file_obj)

    # Create material map from Material ID to trimesh material
    material_map = {}
    for m in c.materials:
        effect = m.effect
        material_map[m.id] = _parse_material(effect, resolver)

    # name : kwargs
    meshes = {}
    # list of dict
    graph = []
    for node in c.scene.nodes:
        _parse_node(node=node,
                    parent_matrix=np.eye(4),
                    material_map=material_map,
                    meshes=meshes,
                    graph=graph,
                    resolver=resolver)

    # create kwargs for load_kwargs
    result = {'class': 'Scene',
              'graph': graph,
              'geometry': meshes}

    return result


def export_collada(mesh, **kwargs):
    """
    Export a mesh or a list of meshes as a COLLADA .dae file.

    Parameters
    -----------
    mesh: Trimesh object or list of Trimesh objects
        The mesh(es) to export.

    Returns
    -----------
    export: str, string of COLLADA format output
    """
    import collada

    meshes = mesh
    if not isinstance(mesh, (list, tuple, set, np.ndarray)):
        meshes = [mesh]

    c = collada.Collada()
    nodes = []
    for i, m in enumerate(meshes):

        # Load uv, colors, materials
        uv = None
        colors = None
        mat = _unparse_material(None)
        if m.visual.defined:
            if m.visual.kind == 'texture':
                mat = _unparse_material(m.visual.material)
                uv = m.visual.uv
            elif m.visual.kind == 'vertex':
                colors = (m.visual.vertex_colors / 255.0)[:, :3]
        c.effects.append(mat.effect)
        c.materials.append(mat)

        # Create geometry object
        vertices = collada.source.FloatSource(
            'verts-array', m.vertices.flatten(), ('X', 'Y', 'Z'))
        normals = collada.source.FloatSource(
            'normals-array', m.vertex_normals.flatten(), ('X', 'Y', 'Z'))
        input_list = collada.source.InputList()
        input_list.addInput(0, 'VERTEX', '#verts-array')
        input_list.addInput(1, 'NORMAL', '#normals-array')
        arrays = [vertices, normals]
        if uv is not None:
            texcoords = collada.source.FloatSource(
                'texcoords-array', uv.flatten(), ('U', 'V'))
            input_list.addInput(2, 'TEXCOORD', '#texcoords-array')
            arrays.append(texcoords)
        if colors is not None:
            idx = 2
            if uv:
                idx = 3
            colors = collada.source.FloatSource('colors-array',
                                                colors.flatten(), ('R', 'G', 'B'))
            input_list.addInput(idx, 'COLOR', '#colors-array')
            arrays.append(colors)
        geom = collada.geometry.Geometry(
            c, uuid.uuid4().hex, uuid.uuid4().hex, arrays
        )
        indices = np.repeat(m.faces.flatten(), len(arrays))

        matref = u'material{}'.format(i)
        triset = geom.createTriangleSet(indices, input_list, matref)
        geom.primitives.append(triset)
        c.geometries.append(geom)

        matnode = collada.scene.MaterialNode(matref, mat, inputs=[])
        geomnode = collada.scene.GeometryNode(geom, [matnode])
        node = collada.scene.Node(u'node{}'.format(i), children=[geomnode])
        nodes.append(node)
    scene = collada.scene.Scene('scene', nodes)
    c.scenes.append(scene)
    c.scene = scene

    b = io.BytesIO()
    c.write(b)
    b.seek(0)
    return b.read()


def _parse_node(node,
                parent_matrix,
                material_map,
                meshes,
                graph,
                resolver=None):
    """
    Recursively parse COLLADA scene nodes.
    """
    import collada

    # Parse mesh node
    if isinstance(node, collada.scene.GeometryNode):
        geometry = node.geometry

        # Create local material map from material symbol to actual material
        local_material_map = {}
        for mn in node.materials:
            symbol = mn.symbol
            m = mn.target
            if m.id in material_map:
                local_material_map[symbol] = material_map[m.id]
            else:
                local_material_map[symbol] = _parse_material(m, resolver)

        # Iterate over primitives of geometry
        for i, primitive in enumerate(geometry.primitives):
            if isinstance(primitive, collada.polylist.Polylist):
                primitive = primitive.triangleset()
            if isinstance(primitive, collada.triangleset.TriangleSet):
                vertex = primitive.vertex
                vertex_index = primitive.vertex_index
                vertices = vertex[vertex_index].reshape(
                    len(vertex_index) * 3, 3)

                # Get normals if present
                normals = None
                if primitive.normal is not None:
                    normal = primitive.normal
                    normal_index = primitive.normal_index
                    normals = normal[normal_index].reshape(
                        len(normal_index) * 3, 3)

                # Get colors if present
                colors = None
                s = primitive.sources
                if ('COLOR' in s and len(s['COLOR'])
                        > 0 and len(primitive.index) > 0):
                    color = s['COLOR'][0][4].data
                    color_index = primitive.index[:, :, s['COLOR'][0][0]]
                    colors = color[color_index].reshape(
                        len(color_index) * 3, 3)

                faces = np.arange(
                    vertices.shape[0]).reshape(
                    vertices.shape[0] // 3, 3)

                # Get UV coordinates if possible
                vis = None
                if primitive.material in local_material_map:
                    material = copy.copy(
                        local_material_map[primitive.material])
                    uv = None
                    if len(primitive.texcoordset) > 0:
                        texcoord = primitive.texcoordset[0]
                        texcoord_index = primitive.texcoord_indexset[0]
                        uv = texcoord[texcoord_index].reshape(
                            (len(texcoord_index) * 3, 2))
                    vis = visual.texture.TextureVisuals(
                        uv=uv, material=material)

                primid = u'{}.{}'.format(geometry.id, i)
                meshes[primid] = {
                    'vertices': vertices,
                    'faces': faces,
                    'vertex_normals': normals,
                    'vertex_colors': colors,
                    'visual': vis}

                graph.append({'frame_to': primid,
                              'matrix': parent_matrix,
                              'geometry': primid})

    # recurse down tree for nodes with children
    elif isinstance(node, collada.scene.Node):
        if node.children is not None:
            for child in node.children:
                # create the new matrix
                matrix = np.dot(parent_matrix, node.matrix)
                # parse the child node
                _parse_node(
                    node=child,
                    parent_matrix=matrix,
                    material_map=material_map,
                    meshes=meshes,
                    graph=graph,
                    resolver=resolver)

    elif isinstance(node, collada.scene.CameraNode):
        # TODO: convert collada cameras to trimesh cameras
        pass
    elif isinstance(node, collada.scene.LightNode):
        # TODO: convert collada lights to trimesh lights
        pass


def _load_texture(file_name, resolver):
    """
    Load a texture from a file into a PIL image.
    """
    from PIL import Image

    file_data = resolver.get(file_name)
    image = Image.open(util.wrap_as_stream(file_data))
    return image


def _parse_material(effect, resolver):
    """
    Turn a COLLADA effect into a trimesh material.
    """
    import collada

    # Compute base color
    baseColorFactor = np.ones(4)
    baseColorTexture = None
    if isinstance(effect.diffuse, collada.material.Map):
        try:
            baseColorTexture = _load_texture(
                effect.diffuse.sampler.surface.image.path, resolver)
        except BaseException:
            log.warning('unable to load base texture',
                        exc_info=True)
    elif effect.diffuse is not None:
        baseColorFactor = effect.diffuse

    # Compute emission color
    emissiveFactor = np.zeros(3)
    emissiveTexture = None
    if isinstance(effect.emission, collada.material.Map):
        try:
            emissiveTexture = _load_texture(
                effect.diffuse.sampler.surface.image.path, resolver)
        except BaseException:
            log.warning('unable to load emissive texture',
                        exc_info=True)
    elif effect.emission is not None:
        emissiveFactor = effect.emission[:3]

    # Compute roughness
    roughnessFactor = 1.0
    if (not isinstance(effect.shininess, collada.material.Map)
            and effect.shininess is not None):
        roughnessFactor = np.sqrt(2.0 / (2.0 + effect.shininess))

    # Compute metallic factor
    metallicFactor = 0.0

    # Compute normal texture
    normalTexture = None
    if effect.bumpmap is not None:
        try:
            normalTexture = _load_texture(
                effect.bumpmap.sampler.surface.image.path, resolver)
        except BaseException:
            log.warning('unable to load bumpmap',
                        exc_info=True)

    # Compute opacity
    if (effect.transparent is not None
            and not isinstance(effect.transparent, collada.material.Map)):
        baseColorFactor = tuple(
            np.append(baseColorFactor[:3], float(effect.transparent[3])))

    return visual.material.PBRMaterial(
        emissiveFactor=emissiveFactor,
        emissiveTexture=emissiveTexture,
        normalTexture=normalTexture,
        baseColorTexture=baseColorTexture,
        baseColorFactor=baseColorFactor,
        metallicFactor=metallicFactor,
        roughnessFactor=roughnessFactor)


def _unparse_material(material):
    """
    Turn a trimesh material into a COLLADA material.
    """
    import collada

    # TODO EXPORT TEXTURES
    if isinstance(material, visual.material.PBRMaterial):
        diffuse = material.baseColorFactor
        if diffuse is not None:
            diffuse = list(diffuse)

        emission = material.emissiveFactor
        if emission is not None:
            emission = [float(emission[0]), float(emission[1]),
                        float(emission[2]), 1.0]

        shininess = material.roughnessFactor
        if shininess is not None:
            shininess = 2.0 / shininess**2 - 2.0

        effect = collada.material.Effect(
            uuid.uuid4().hex, params=[], shadingtype='phong',
            diffuse=diffuse, emission=emission,
            specular=[1.0, 1.0, 1.0, 1.0], shininess=float(shininess)
        )
        material = collada.material.Material(
            uuid.uuid4().hex, 'pbrmaterial', effect
        )
    else:
        effect = collada.material.Effect(
            uuid.uuid4().hex, params=[], shadingtype='phong'
        )
        material = collada.material.Material(
            uuid.uuid4().hex, 'defaultmaterial', effect
        )
    return material


def load_zae(file_obj, resolver=None, **kwargs):
    """
    Load a ZAE file, which is just a zipped DAE file.

    Parameters
    -------------
    file_obj : file object
      Contains ZAE data
    resolver : trimesh.visual.Resolver
      Resolver to load additional assets
    kwargs : dict
      Passed to load_collada

    Returns
    ------------
    loaded : dict
      Results of loading
    """

    # a dict, {file name : file object}
    archive = util.decompress(file_obj,
                              file_type='zip')

    # load the first file with a .dae extension
    file_name = next(i for i in archive.keys()
                     if i.lower().endswith('.dae'))

    # a resolver so the loader can load textures / etc
    resolver = visual.resolvers.ZipResolver(archive)

    # run the regular collada loader
    loaded = load_collada(archive[file_name],
                          resolver=resolver,
                          **kwargs)
    return loaded


# only provide loaders if `pycollada` is installed
_collada_loaders = {}
_collada_exporters = {}
if util.has_module('collada'):
    _collada_loaders['dae'] = load_collada
    _collada_loaders['zae'] = load_zae
    _collada_exporters['dae'] = export_collada
