import sys
import numpy as np
from copy import deepcopy
from chroma import layout
from chroma.stl import read_stl
from chroma.transform import rotate
from chroma.geometry import Geometry, Solid
from chroma.materials import glass, h2o
from itertools import product

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__()

        pmt_mesh = read_stl(layout.models + '/hamamatsu_12inch.stl')/1000.0
        pmt_solid = Solid(pmt_mesh, glass, h2o)

        # construct the barrel
        for i in range(pmts_per_string):
            for j in range(nstrings):
                pmt = deepcopy(pmt_solid)
                pmt.mesh += (-radius,0,-height/2+i*height/(pmts_per_string-1))
                pmt.mesh = rotate(pmt.mesh, j*2*np.pi/nstrings, (0,0,1))
                self.add_solid(pmt)

        # 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:
                pmt = deepcopy(pmt_solid)
                pmt.mesh = rotate(pmt.mesh, -np.pi/2, (0,1,0))
                pmt.mesh += (x,y,+height/2+height/(pmts_per_string-1)/2)
                self.add_solid(pmt)

        # 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:
                pmt = deepcopy(pmt_solid)
                pmt.mesh = rotate(pmt.mesh, +np.pi/2, (0,1,0))
                pmt.mesh += (x,y,-height/2-height/(pmts_per_string-1)/2)
                self.add_solid(pmt)