#! /usr/bin/env python

"""
3D event viewer via WebGL/Three.js

TODO:
 * Non-BS spacetime units!!
 * Make B-field stepping pT-dependent: more steps for high pT
 * Particle labeling/interactive mouse-over
 * Line thickness/saturation/opacity proportional to pT or E
 * Better detector
 * Animation
"""

import argparse
ap = argparse.ArgumentParser()
ap.add_argument("HEPMCFILE", help="file to read events from")
ap.add_argument("--mom", dest="MOMENTUM_VIEW", action="store_true", default=False, help="show particles in momentum space (default=spacetime)")
ap.add_argument("--ptmin", dest="PT_CUTOFF", metavar="PT", type=float, default=0., help="hide particles with pT < PT (in GeV)")
args = ap.parse_args()

import hepmcio
reader = hepmcio.HepMCReader(args.HEPMCFILE)
evt = reader.next()
evt = reader.next()

import numpy as np

import pypdt
PDT = pypdt.PDT()

doc = """
<!DOCTYPE html>
<html>
	<head>
		<meta charset=utf-8>
		<title>My first three.js app</title>
		<style>
			body { margin: 0; }
			canvas { width: 100%; height: 100% }
		</style>
	</head>
	<body>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/93/three.js"></script>
    <!-- <script src="three.js-master/examples/js/controls/OrbitControls.js"></script> -->
    <script src="OrbitControls.js"></script>
<script>

var camera, controls, scene, renderer;
var material, geometry, line, light;
init();
animate();

function init() {
  renderer = new THREE.WebGLRenderer();
  renderer.setSize( window.innerWidth, window.innerHeight );
  document.body.appendChild( renderer.domElement );

  camera = new THREE.PerspectiveCamera( 45, window.innerWidth / window.innerHeight, 0.01, 5000 );
  camera.position.set( 400, 150, 700 );
  //camera.position.set( 0, 0, 500 );
  camera.lookAt( new THREE.Vector3( 0, 0, 0 ) );

  controls = new THREE.OrbitControls( camera, renderer.domElement );
  ////controls.addEventListener( 'change', render ); // call this only in static scenes (i.e., if there is no animation loop)
  //controls.enableDamping = true; // an animation loop is required when either damping or auto-rotation are enabled
  //controls.dampingFactor = 0.1;
  //controls.screenSpacePanning = false;
  //controls.minDistance = 100;
  //controls.maxDistance = 500
  //controls.maxPolarAngle = Math.PI / 2;

  scene = new THREE.Scene();

  light = new THREE.PointLight(0xffffff);
  light.position.set(200.0, 200.0, 500.0);
  scene.add( light );
  light = new THREE.AmbientLight(0x444444)
  scene.add( light );

  var material_det = new THREE.MeshLambertMaterial( { color: 0x9999ff, transparent: true, opacity: 0.4, emissive: 10, reflectivity: 1 } );
  var material_beam = new THREE.LineDashedMaterial( { color: 0x666666, linewidth: 5. } );
  var material_pho = new THREE.LineBasicMaterial( { color: 0x0099aa, linewidth: 2. } );
  var material_lep = new THREE.LineBasicMaterial( { color: 0xddbb00, linewidth: 4. } );
  var material_nu  = new THREE.LineBasicMaterial( { color: 0x333333, linewidth: 2. } );
  var material_had = new THREE.LineBasicMaterial( { color: 0x338833, linewidth: 2. } );

  geometry = new THREE.Geometry();
  geometry.vertices.push(new THREE.Vector3( 0, 0, -10000 ) );
  geometry.vertices.push(new THREE.Vector3( 0, 0,  10000 ) );
  line = new THREE.Line( geometry, material_beam );
  scene.add( line );

  var geometry = new THREE.CylinderGeometry( 115.0, 115.0, 700.0, 64 );
  var solenoid = new THREE.Mesh( geometry, material_det );
  solenoid.rotation.x += Math.PI / 2
  scene.add( solenoid );

"""

## Find stable particles
particles1 = []
for ip, p in evt.particles.items():
    if p.status != 1:
        continue
    if p.mom[0]**2 + p.mom[1]**2 < args.PT_CUTOFF**2:
        continue
    particles1.append(p)

## Find interesting ancestors of stable particles
def get_ancestors(p, dmin=1e-5):
    rtn = []
    v = p.vtx_start()
    if v:
        d2 = v.pos[0]**2 + v.pos[1]**2 + v.pos[2]**2
        if d2 > dmin**2:
            for pp in p.parents():
                # if pp.status != 2:
                #     continue
                rtn += get_ancestors(pp)
    rtn.append(p)
    return rtn
#
particles2 = []
for p1 in particles1:
    particles2 += get_ancestors(p1)[:-1]


## Render particles
for p in particles2 + particles1:
    pidname = "had"
    if p.pid == 22:
        pidname = "pho"
    elif abs(p.pid) in [11,13,15]:
        pidname = "lep"
    elif abs(p.pid) in [12,14,16]:
        pidname = "nu"

    DT = 5e-10 #< seconds
    CLIGHT = 3e10 # mm/s
    QE = 1.6022e-19 # mm/s
    SOLENOID_BVEC = np.array([0,0,4]) #< T
    NSTEP = 16

    verts = []
    if args.MOMENTUM_VIEW:
        verts += [(0,0,0), p.mom[:3]]
    else:
        E = p.mom[3]
        mass = 0.0
        if p.mom[0]**2 + p.mom[1]**2 + p.mom[2]**2 < E**2:
            mass = np.sqrt(E**2 - p.mom[0]**2 - p.mom[1]**2 - p.mom[2]**2)
        pos3 = np.array(p.vtx_start().pos[:3] if p.vtx_start() else [0.,0.,0.])
        mom3 = np.array(p.mom[:3])
        charge = PDT[p.pid].charge if PDT[p.pid] else 0.
        verts.append(pos3.tolist())
        for i in range(NSTEP):
            #print mom3, pos3
            vel3 = mom3/E * CLIGHT # p/E = beta, with c=1
            if charge != 0:
                pos3 += vel3*DT
                force = PDT[p.pid].charge * np.cross(vel3, SOLENOID_BVEC) #* QE
                mom3 += DT*force  / 500
                #print pos3, mom3, vel3, force
                verts.append(pos3.tolist())
            else:
                pos3 += vel3*NSTEP*DT
                verts.append(pos3.tolist())
                break

    if len(verts) == 2:
        doc += "\n  geometry = new THREE.Geometry();"
        for v in verts:
            doc += "\n  geometry.vertices.push(new THREE.Vector3( {v[0]:.2f}, {v[1]:.2f}, {v[2]:.2f} ) );".format(v=v)
    else:
        doc += "\n  var curve = new THREE.CatmullRomCurve3( ["
        for iv, v in enumerate(verts):
            doc += "\n    new THREE.Vector3( {v[0]:.2f}, {v[1]:.2f}, {v[2]:.2f} ) ".format(v=v)
            doc += " ] );" if iv == len(verts)-1 else ","
        doc += "\n  geometry = new THREE.BufferGeometry().setFromPoints( curve.getPoints(100) );"

    doc += "\n  line = new THREE.Line( geometry, material_{mat} );".format(mat=pidname)
    doc += "\n  scene.add( line );\n"

doc += """
}

function animate() {
    requestAnimationFrame( animate );
    //controls.update(); // only required if controls.enableDamping = true, or if controls.autoRotate = true
    render();
}

function render() {
    renderer.render( scene, camera );
}

</script>
    </body>
</html>
"""

print doc
