import numpy as np
import numpy.linalg
import math
from chroma.make import rotate_extrude
from chroma.stl import mesh_from_stl
from chroma.geometry import *
from chroma.optics import *
from chroma.solids import build_8inch_pmt_with_lc
from chroma.transform import rotate, make_rotation_matrix
import os

def sno_vessel(sphere_radius, neck_radius, neck_top, nsteps=40):
    '''Compute the 2D coordinates of the profile of one side of a 
    SNO-style acrylic vessel.  The center of the sphere is at (0,0), with
    the neck extending along the positive y direction.
    
       sphere_radius: Radius of spherical part of profile
       neck_radius: Radius of neck part
       neck_top: y coordinate of top of neck
       nsteps: number of points around the circumference of the sphere

       Returns: Tuple of x and y coordinate numpy arrays.
    '''
    if neck_radius >= sphere_radius:
        raise ValueError('neck_radius must be less than sphere_radius')

    intersect_height = (sphere_radius**2 - neck_radius**2)**0.5
    max_angle = math.atan2(intersect_height, neck_radius)

    if neck_top < intersect_height:
        raise ValueError('neck_top must be greater than the y-value where the sphere and cylinder intersect')

    # Start with point at bottom
    angles = np.linspace(-math.pi/2, max_angle, nsteps)
    x = list(np.cos(angles) * sphere_radius)
    y = list(np.sin(angles) * sphere_radius)
    x[0] = 0.0 # Round-off error might make cos(-pi/2) not exactly zero
    
    # Neck intersection point
    x.append(neck_radius)
    y.append(intersect_height)

    # Top of neck
    x.append(neck_radius)
    y.append(neck_top)
    
    # Top of neck on axis
    x.append(0.0)
    y.append(neck_top)
    
    return x, y

##### SNO Parts

nsteps = 40
av_outside_profile = sno_vessel(sphere_radius=6.0604, neck_radius=0.79375, 
                                neck_top=10.50, nsteps=nsteps)
# For simplicity, cap the top of the AV with acrylic
av_inside_profile  = sno_vessel(sphere_radius=6.0053, neck_radius=0.72898,
                                neck_top=10.00, nsteps=nsteps)

av_outside_mesh = rotate_extrude(av_outside_profile[0], av_outside_profile[1],
                                 nsteps)
av_outside_mesh.vertices = rotate(av_outside_mesh.vertices, np.pi/2, (-1,0,0))

av_inside_mesh = rotate_extrude(av_inside_profile[0], av_inside_profile[1],
                                 nsteps)
av_inside_mesh.vertices = rotate(av_inside_mesh.vertices, np.pi/2, (-1,0,0))

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

def build_sno(real_av=False):
    import h5py
    pmtinfo = h5py.File(dir+'/sno_pmt_info.hdf5')

    pmt = build_8inch_pmt_with_lc()


    if real_av:
        geo = Geometry(water)
        real_av_mesh = mesh_from_stl(dir+'/sno_av_ascii.stl.bz2')
        real_av_mesh.vertices *= 0.0254 # inch -> meter
        geo.add_solid(Solid(real_av_mesh, glass, water, color=0xBBAAAAFF))
    else:
        geo = Geometry(water)
        geo.add_solid(Solid(av_outside_mesh, acrylic_sno, water, 
                            color=0xBBFFFFFF))
        geo.add_solid(Solid(av_inside_mesh, water, acrylic_sno, 
                            color=0xBB0000FF))
    geo.pmtids = []

    snoman_id_dict = {}

    pmt_offset = 122.0 - 2 * 0.01# max bucket height - 2 * facegap (mm)

    for position, direction, idnum, pmt_type in zip(pmtinfo['pos'], pmtinfo['dir'],
                                          pmtinfo['snoman_id'],
                                          pmtinfo['pmt_type']):

        # PMT type codes:
        #        1 = normal; 2 = OWL; 3 = Low Gain; 4 = BUTT; 5 = Neck
        #        6 = Calib channel;  10 = spare;  11 = invalid
        if pmt_type != 1:  # All the other PMTs have nonsense directions
            continue

        
        # Flip and renormalize
        direction *= -1.0/numpy.linalg.norm(direction)

        # Orient PMT that starts facing Y axis
        y_axis = np.array((0.0,1.0,0.0))
        axis = np.cross(direction, y_axis)
        angle = np.arccos(np.dot(y_axis, direction))
        rotation = make_rotation_matrix(angle, axis)

        # Values in database are for front face of concentrator.
        # need to shift so position is for equator of PMT
        if pmt_type == 1:
            displacement = position - direction * pmt_offset
        displacement /= 1000.0 # mm -> m
        chroma_id = geo.add_solid(pmt, rotation, displacement)
        snoman_id_dict[chroma_id] = idnum
        geo.pmtids.append(chroma_id)

    # Convert dict to numpy array for fast array remapping later
    chroma_ids = np.array(snoman_id_dict.keys())
    snoman_ids = np.array(snoman_id_dict.values())
    snoman_id_map = np.zeros(max(chroma_ids)+1) - 1 # Fill with -1 everywhere
    snoman_id_map[chroma_ids] = snoman_ids

    geo.snoman_id_map = snoman_id_map
        
    return geo