import os
import numpy as np
import pickle
from chroma import *

models_directory = os.path.split(os.path.realpath(__file__))[0] + '/../models'

strings = 10
pmts_per_string = 10
radius = 10.0
height = 20.0

def build_lbne(bits=4):
    lbne = geometry.Geometry()

    sphere_mesh = stl.read_stl(models_directory + '/hamamatsu_12inch.stl')
    sphere_mesh /= 1000.0

    solids = []
    for i in range(pmts_per_string):
        for j in range(strings):
            sphere = np.copy(sphere_mesh)
            sphere += (-radius,0,i*(height/pmts_per_string))
            sphere = transform.rotate(sphere, j*2*np.pi/strings, (0,0,1))
            lbne.add_solid(geometry.Solid(sphere, materials.vacuum, materials.h2o))

    lbne.build(bits)

    return lbne

def cache_lbne(filename):
    lbne = build_lbne(bits=8)
    f = open(filename, 'wb')
    pickle.dump(lbne, f)
    f.close()

def load_lbne(filename):
    f =  open(filename, 'rb')
    lbne = pickle.load(f)
    f.close()
    return lbne

if __name__ == '__main__':
    build_lbne()