# -*- coding: utf-8 -*-
# This file is part of the Horus Project

__author__ = 'Jesús Arroyo Torrens <jesus.arroyo@bq.com>'
__copyright__ = 'Copyright (C) 2014-2016 Mundo Reader S.L.'
__license__ = 'GNU General Public License v2 http://www.gnu.org/licenses/gpl2.html'

"""
PLY file point cloud loader.

    - Binary, which is easy and quick to read.
    - Ascii, which is harder to read, as can come with windows, mac and unix style newlines.

This module also contains a function to save objects as an PLY file.

http://en.wikipedia.org/wiki/PLY_(file_format)
"""

import struct
import numpy as np

from horus import __version__
from horus.util import model


def _load_ascii(mesh, stream, dtype, count):
    fields = dtype.fields

    v = 0
    c = 0

    if 'c' in fields:
        c += 3
    if 'n' in fields:
        c += 3

    i = 0
    while i < count:
        i += 1
        data = stream.readline().split(' ')
        if data is not None:
            mesh._add_vertex(data[v], data[v + 1], data[v + 2], data[c], data[c + 1], data[c + 2])


def _load_binary(mesh, stream, dtype, count):
    data = np.fromfile(stream, dtype=dtype, count=count)

    fields = dtype.fields
    mesh.vertex_count = count

    if 'v' in fields:
        mesh.vertexes = data['v']
    else:
        mesh.vertexes = np.zeros((count, 3))

    if 'n' in fields:
        mesh.normal = data['n']
    else:
        mesh.normal = np.zeros((count, 3))

    if 'c' in fields:
        mesh.colors = data['c']
    else:
        mesh.colors = 255 * np.ones((count, 3))


def load_scene(filename):
    obj = model.Model(filename, is_point_cloud=True)
    m = obj._add_mesh()
    with open(filename, "rb") as f:
        dtype = []
        count = 0
        format = None
        line = None
        header = ''

        while line != 'end_header\n' and line != '':
            line = f.readline()
            header += line
        # Discart faces
        header = header.split('element face ')[0].split('\n')

        if header[0] == 'ply':

            for line in header:
                if 'format ' in line:
                    format = line.split(' ')[1]
                    break

            if format is not None:
                if format == 'ascii':
                    fm = ''
                elif format == 'binary_big_endian':
                    fm = '>'
                elif format == 'binary_little_endian':
                    fm = '<'

            df = {'float': fm + 'f', 'uchar': fm + 'B'}
            dt = {'x': 'v', 'nx': 'n', 'red': 'c', 'alpha': 'a'}
            ds = {'x': 3, 'nx': 3, 'red': 3, 'alpha': 1}

            for line in header:
                if 'element vertex ' in line:
                    count = int(line.split('element vertex ')[1])
                elif 'property ' in line:
                    props = line.split(' ')
                    if props[2] in dt.keys():
                        dtype = dtype + [(dt[props[2]], df[props[1]], (ds[props[2]],))]

            dtype = np.dtype(dtype)

            if format is not None:
                if format == 'ascii':
                    m._prepare_vertex_count(count)
                    _load_ascii(m, f, dtype, count)
                elif format == 'binary_big_endian' or format == 'binary_little_endian':
                    _load_binary(m, f, dtype, count)
            obj._post_process_after_load()
            return obj

        else:
            print "Error: incorrect file format."
            return None


def save_scene(filename, _object):
    with open(filename, 'wb') as f:
        save_scene_stream(f, _object)


def save_scene_stream(stream, _object):
    m = _object._mesh

    binary = True

    if m is not None:
        frame = "ply\n"
        if binary:
            frame += "format binary_little_endian 1.0\n"
        else:
            frame += "format ascii 1.0\n"
        frame += "comment Generated by Horus {0}\n".format(__version__)
        frame += "element vertex {0}\n".format(m.vertex_count)
        frame += "property float x\n"
        frame += "property float y\n"
        frame += "property float z\n"
        frame += "property uchar red\n"
        frame += "property uchar green\n"
        frame += "property uchar blue\n"
        frame += "element face 0\n"
        frame += "property list uchar int vertex_indices\n"
        frame += "end_header\n"
        stream.write(frame)
        if m.vertex_count > 0:
            if binary:
                for i in xrange(m.vertex_count):
                    stream.write(struct.pack("<fffBBB",
                                             m.vertexes[i, 0], m.vertexes[i, 1], m.vertexes[i, 2],
                                             m.colors[i, 0], m.colors[i, 1], m.colors[i, 2]))
            else:
                for i in xrange(m.vertex_count):
                    stream.write("{0} {1} {2} {3} {4} {5}\n".format(
                                 m.vertexes[i, 0], m.vertexes[i, 1], m.vertexes[i, 2],
                                 m.colors[i, 0], m.colors[i, 1], m.colors[i, 2]))
