import numpy as np
from chroma.geometry import Solid, Geometry
import chroma.solids as solids
from chroma.optics import vacuum, water, water_wcsim, black_surface
from chroma.transform import rotate, make_rotation_matrix
from itertools import product
import chroma.make as make

def build_lbne(radius, height, nstrings, pmts_per_string, endcap_spacing, physical_model=True):
    if physical_model:
        pmt = solids.build_12inch_pmt_with_lc(outer_material=water_wcsim)
    else:
        pmt = solids.build_12inch_pmt_shell(outer_material=water_wcsim)

    lbne = Geometry()

    # outer cylinder
    cylinder_mesh = make.segmented_cylinder(radius, height+height/(pmts_per_string-1), nsteps=16*nstrings, nsegments=1200)
    cylinder_mesh.vertices = rotate(cylinder_mesh.vertices, np.pi/2, (-1,0,0))
    lbne.add_solid(Solid(cylinder_mesh, water_wcsim, vacuum, black_surface, 0xff0000ff))

    lbne.pmtids = []

    # 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))
            lbne.pmtids.append(lbne.add_solid(pmt, rotation, displacement))

    # 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 - endcap_spacing/2:
            rotation = make_rotation_matrix(+np.pi/2, (1,0,0))
            displacement = (x,y,+height/2+height/(pmts_per_string-1)/2)
            lbne.pmtids.append(lbne.add_solid(pmt, rotation, displacement))

    # 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 - endcap_spacing/2:
            rotation = make_rotation_matrix(-np.pi/2, (1,0,0))
            displacement = (x,y,-height/2-height/(pmts_per_string-1)/2)
            lbne.pmtids.append(lbne.add_solid(pmt, rotation, displacement))

    return lbne