import numpy as np
import numpy.linalg
import math
from chroma.make import linear_extrude, sphere
from chroma.geometry import *
from chroma.optics import *
from chroma.transform import rotate, make_rotation_matrix
from chroma.itertoolset import grouper
import os


dir = os.path.split(os.path.realpath(__file__))[0]

def read_polygons(filename):
    polygons = {}
    with open(filename) as f:
        for line in f:
            parts = line.split()
            idnum = int(parts[0])
            npoints = int(parts[1])
            points = np.zeros(shape=(npoints,2), dtype=np.float32)
            for i, (x,y) in enumerate(grouper(2, parts[2:])):
                points[i] = [float(x), float(y)]
            polygons[idnum] = points
    return polygons

def read_cassettes(filename):
    cassettes = []
    with open(filename) as f:
        for line in f:
            parts = line.split()
            cassette_type = int(parts[0])
            rotation = np.transpose(np.array(map(float, parts[1:10]), dtype=np.float32).reshape((3,3)))
            displacement = np.array(map(float, parts[10:13]), dtype=np.float32)
            
            cassettes.append({'type' : cassette_type, 'rotation' : rotation,
                              'displacement' : displacement})

    return cassettes

def build_miniclean(real_av=False):
    geo = Geometry()

    simple_iv = sphere(0.818)
    geo.add_solid(Solid(simple_iv, liquid_argon, vacuum, color=0xffFF0000))

    polygons = read_polygons(os.path.join(dir, 'miniclean_polygons.txt'))

    height = 0.30
    polygon_types = {}
    for polygon_id, polygon_points in polygons.items():
        mesh = linear_extrude(polygon_points[::-1,0]/1000.0,
                              polygon_points[::-1,1]/1000.0,
                              height=height,
                              center=[0,0,height/2])
        colors = np.zeros(len(mesh.triangles), dtype=np.uint32)
        colors.fill(0x03CCCCFF)
        # Color the faces on one end
        triangles = mesh.assemble(group=True)
        triangle_centroids = triangles.mean(axis=1)
        face = triangle_centroids[:,2] < 0.1
        colors[face] = 0xffffff

        polygon_types[polygon_id] = (mesh, colors)

    cassettes = read_cassettes(os.path.join(dir, 'miniclean_cassettes.txt'))

    geo.pmtids = []
    geo.pmt_id_map = {}
    for i, cassette in enumerate(cassettes):
        polygon_mesh, polygon_colors = polygon_types[cassette['type']]
        solid = Solid(polygon_mesh, liquid_argon, liquid_argon, surface=shiny_surface,
                      color=polygon_colors)

        chroma_id = geo.add_solid(solid, cassette['rotation'], 
                                  cassette['displacement']/1000.0)
        geo.pmt_id_map[chroma_id] = i
        geo.pmtids.append(chroma_id)
        
    return geo