"""
Morfeus lab
The University of Texas
MD Anderson Cancer Center
Author - Caleb O'Connor
Email - csoconnor@mdanderson.org


Description:

Functions:

"""

import os

import zipfile
from PIL import ImageColor
import xml.etree.ElementTree as ET

import numpy as np
import pyvista as pv


class ThreeMfReader(object):
    """
    Converts 3mf file to pyvista polydata mesh.
    """
    def __init__(self, reader):
        self.reader = reader

    def input_files(self, files):
        self.reader.files['3mf'] = files

    def load(self):
        for file_path in self.reader.files['3mf']:
            self.read(file_path)

    def read(self, path):
        """
        Loads in the 3mf file, gets the vertices/vertice colors/triangles and creates a polydata 3D model using pyvista.

        Returns
        -------

        """
        namespace = {"3mf": "http://schemas.microsoft.com/3dmanufacturing/core/2015/02",
                     "m": "http://schemas.microsoft.com/3dmanufacturing/material/2015/02"}

        archive = zipfile.ZipFile(path, "r")
        root = ET.parse(archive.open("3D/3dmodel.model"))
        color_list = list()
        colors = root.findall('.//m:color', namespace)
        if colors:
            for color in colors:
                color_list.append(color.get("color", 0))

        obj = root.findall("./3mf:resources/3mf:object", namespace)[0]
        triangles = obj.findall(".//3mf:triangle", namespace)

        vertex_list = []
        for vertex in obj.findall(".//3mf:vertex", namespace):
            vertex_list.append([vertex.get("x"), vertex.get("y"), vertex.get("z")])

        triangle_list = np.empty((1, 4 * len(triangles)))
        vertices_color = np.zeros((len(vertex_list), 3))
        for ii, triangle in enumerate(triangles):
            v1 = int(triangle.get("v1"))
            v2 = int(triangle.get("v2"))
            v3 = int(triangle.get("v3"))
            tricolor = self.color_avg(color_list, (triangle.get("p1")), (triangle.get("p2")), (triangle.get("p3")))
            rgb_color = list(ImageColor.getcolor(tricolor, "RGB")[0:3])
            vertices_color[v1] = rgb_color
            vertices_color[v2] = rgb_color
            vertices_color[v3] = rgb_color
            triangle_list[0, ii * 4:(ii + 1) * 4] = [3, v1, v2, v3]

        self.reader.meshes += [pv.PolyData(np.float64(np.asarray(vertex_list)), triangle_list[0, :].astype(int))]
        self.reader.meshes[-1]['colors'] = np.abs(255-vertices_color)

    @staticmethod
    def color_avg(color_list, p1, p2, p3):
        """
        Get the average color from color list.

        Parameters
        ----------
        color_list
        p1
        p2
        p3

        Returns
        -------

        """
        p2rgb = None
        p3rgb = None

        p1hex = color_list[int(p1)]
        value = p1hex.lstrip('#')
        lv = len(value)
        p1rgb = tuple(int(value[i:i + lv // 3], 16) for i in range(0, lv, lv // 3))

        if isinstance(p2, int):
            p2hex = color_list[int(p2)]
            value = p2hex.lstrip('#')
            lv = len(value)
            p2rgb = tuple(int(value[i:i + lv // 3], 16) for i in range(0, lv, lv // 3))

        if isinstance(p3, int):
            p3hex = color_list[int(p3)]
            value = p3hex.lstrip('#')
            lv = len(value)
            p3rgb = tuple(int(value[i:i + lv // 3], 16) for i in range(0, lv, lv // 3))

        if p2rgb is not None and p3rgb is not None:
            rgbAvg = np.average(np.array(p1rgb), np.array(p2rgb), np.array(p3rgb))

        elif p2rgb is not None:
            rgbAvg = np.average(np.array(p1rgb), np.array(p2rgb))

        else:
            rgbAvg = p1rgb

        hexAvg = '#%02x%02x%02x' % (rgbAvg[0], rgbAvg[1], rgbAvg[2])

        return hexAvg
