import os
import sys
import numpy as np

dir = os.path.split(os.path.realpath(__file__))[0]
sys.path.append(dir + '/..')

import models
from mesh import mesh_from_stl
from solid import Solid
from geometry import Geometry
from transform import rotate, make_rotation_matrix
from itertools import product
from materials import *
from make import cylinder

endcap_spacing = .485

radius = 25.0/10.0
height = 50.0/10.0

nstrings = 324//10
pmts_per_string = 102//10

class LBNE(Geometry):
    def __init__(self):
        super(LBNE, self).__init__()

        # outer cylinder
        cylinder_solid = Solid(0, cylinder(radius, height+height/(pmts_per_string-1)), material1=lightwater_sno, material2=vacuum, surface=black_surface, color=0x0000ff)
        self.add_solid(cylinder_solid)

        pmt_inner_mesh = \
            mesh_from_stl(models.dir + '/hamamatsu_12inch_inner.stl')
        pmt_outer_mesh = \
            mesh_from_stl(models.dir + '/hamamatsu_12inch_outer.stl')

        photocathode_triangles = np.mean(pmt_inner_mesh[:], axis=1)[:,1] > 0

        inner_color = np.empty(len(pmt_inner_mesh.triangles), np.uint32)
        inner_color[photocathode_triangles] = 0xff0000
        inner_color[~photocathode_triangles] = 0x00ff00

        outer_color = np.empty(len(pmt_outer_mesh.triangles), np.uint32)
        outer_color[:] = 0xffffff

        inner_surface = np.empty(len(pmt_inner_mesh.triangles), np.object)
        inner_surface[photocathode_triangles] = black_surface
        inner_surface[~photocathode_triangles] = shiny_surface

        self.pmtids = [1]

        # construct the barrel
        for i in range(pmts_per_string):
            for j in range(nstrings):
                rotation = make_rotation_matrix(j*2*np.pi/nstrings, (0,0,1))
                displacement = rotate((0,-radius,-height/2+i*height/(pmts_per_string-1)), j*2*np.pi/nstrings, (0,0,1))
                self.add_solid(Solid(self.pmtids[-1], pmt_inner_mesh,
                                     rotation, displacement,
                                     material1=vacuum,
                                     material2=glass,
                                     surface=inner_surface,
                                     color=inner_color))
                self.add_solid(Solid(self.pmtids[-1], pmt_outer_mesh,
                                     rotation, displacement,
                                     material1=glass,
                                     material2=lightwater_sno,
                                     color=outer_color))
                self.pmtids += [self.pmtids[-1]+1]

        # construct the top endcap
        for x, y in np.array(tuple(product(\
                    np.arange(-radius+0.075, radius, endcap_spacing),
                    np.arange(-radius+0.075, radius, endcap_spacing)))):
            if np.sqrt(x**2 + y**2) <= radius:
                rotation = make_rotation_matrix(+np.pi/2, (1,0,0))
                displacement = (x,y,+height/2+height/(pmts_per_string-1)/2)
                self.add_solid(Solid(self.pmtids[-1], pmt_inner_mesh,
                                     rotation, displacement,
                                     material1=vacuum,
                                     material2=glass,
                                     surface=inner_surface,
                                     color=inner_color))
                self.add_solid(Solid(self.pmtids[-1], pmt_outer_mesh,
                                     rotation, displacement,
                                     material1=glass,
                                     material2=lightwater_sno,
                                     color=outer_color))
                self.pmtids += [self.pmtids[-1]+1]

        # construct the bottom endcap
        for x, y in np.array(tuple(product(\
                    np.arange(-radius+0.075, radius, endcap_spacing),
                    np.arange(-radius+0.075, radius, endcap_spacing)))):
            if np.sqrt(x**2 + y**2) <= radius:
                rotation = make_rotation_matrix(-np.pi/2, (1,0,0))
                displacement = (x,y,-height/2-height/(pmts_per_string-1)/2)
                self.add_solid(Solid(self.pmtids[-1], pmt_inner_mesh,
                                     rotation, displacement,
                                     material1=vacuum,
                                     material2=glass,
                                     surface=inner_surface,
                                     color=inner_color))
                self.add_solid(Solid(self.pmtids[-1], pmt_outer_mesh,
                                     rotation, displacement,
                                     material1=glass,
                                     material2=lightwater_sno,
                                     color=outer_color))
                self.pmtids += [self.pmtids[-1]+1]