#include "TROOT.h"
#include "TClass.h"
#include "TObject.h"
#include "TNamed.h"
#include "TMath.h"
#include "TRandom.h"
#include "TRandom3.h"
#include "TRotation.h"
#include "TVector3.h"
#include "TGeoManager.h"
#include "TGeoMatrix.h"
#include "TGeoMedium.h"
#include "TGeoVolume.h"
#include "TGeoNode.h"
#include "TGeoShape.h"
#include "TGeoBBox.h"
#include "TGeoHype.h"
#include "TGeoCompositeShape.h"

const Double_t pi = TMath::Pi();               // Pi
const Double_t h = TMath::H();                  // Planck's constant
Double_t lam = 780 * 1e-9;                  // Wavelength ( 780 nm )
Double_t molcham_rmin = 5.870 * 0.0254 / 2;            // Inner radius of the molasses chamber
Double_t molcham_rmax = 6 * 0.0254 / 2;               // Outer radius of the molasses chamber
Double_t w0lens[4] = { 1.16e-6 , 1.16e-6 , 1.16e-6 , 1.16e-6 };      // Waist of laser beams
Double_t zrlens[4];                     // Rayleigh length per beam
Double_t beamdiverlens[4];                  // Beam divergence
Double_t lenswidth;                     // Average spot size of beams at center
TVector3 clens( -0.0001075 , 0.00 , 0.01 );             // Position lens center
TVector3 w0lenspos[4];                     // Position of waist of laser beams
TVector3 klens[4];                     // Wavevector of laser beams

void chamber_geom()
{
   //------Beam geometry (all lengths must be in centimeters)------//

   for ( Int_t i = 0 ; i < 4 ; i ++ )
   {
      klens[i].SetX( 1 );

      klens[i].SetMag( 1 / lam );

      klens[i].SetTheta( pi / 2 );

      if ( i < 2 )

         klens[i].SetPhi( i * pi );

      else

         klens[i].SetPhi( ( 2 * i - 3 ) * pi / 2 );


      w0lenspos[i] = - klens[i].Unit();

      w0lenspos[i].SetMag( 0.035 );

      zrlens[i] = pi * pow ( w0lens[i] , 2 ) / lam;

      lenswidth += w0lens[i] * sqrt( 1 + pow( w0lenspos[i].Mag() / zrlens[i] , 2 ) ) / 4;

      beamdiverlens[i] = atan( lam / ( pi * w0lens[i] ) ) * 180 / pi;
   }



   //------Definition of top volume corresponding to a cube circumscribing the internal region of the molasses vacuum chamber------//

   TGeoManager *mgr = new TGeoManager( "chamber" , "Simple geometry" );

   TGeoMedium *medium = 0;

   TGeoVolume *top = mgr->MakeBox( "top" , medium , 100 * molcham_rmin , 100 * molcham_rmin , 100 * molcham_rmin );

   mgr->SetTopVolume( top );


   //------Rotations and translations to be performed on hyperboloid base------//


   //------Rotations are in GEANT3 style, 3 pairs of angles, corresponding to angular position of the new axes------//

   TGeoRotation *rot1 = new TGeoRotation( "rot1", 90. , 90. , 0. , 0. , 90. , 0. );

   TGeoRotation *rot2 = new TGeoRotation( "rot2", 90. , 270. , 180. , 0. , 90. , 0. );

   TGeoRotation *rot3 = new TGeoRotation( "rot3", 90. , 0. , 180. , 0. , 90. , 90. );

   TGeoRotation *rot4 = new TGeoRotation( "rot4", 90. , 0. , 0. , 0. , 90. , 270. );


   //------Combination of translations and rotations (laser lens)------//

   Double_t centerlens[3] = { clens.X() * 100 , clens.Y() * 100 , clens.Z() * 100 };

   TGeoCombiTrans *combi11 = new TGeoCombiTrans( "combi11" , w0lenspos[0].X() * 100 + centerlens[0] , w0lenspos[0].Y() * 100 + centerlens[1] , w0lenspos[0].Z() * 100 + centerlens[2] , rot1 );

   TGeoCombiTrans *combi21 = new TGeoCombiTrans( "combi21" , w0lenspos[1].X() * 100 + centerlens[0] , w0lenspos[1].Y() * 100 + centerlens[1] , w0lenspos[1].Z() * 100 + centerlens[2] , rot2 );

   TGeoCombiTrans *combi31 = new TGeoCombiTrans( "combi31" , w0lenspos[2].X() * 100 + centerlens[0] , w0lenspos[2].Y() * 100 + centerlens[1] , w0lenspos[2].Z() * 100 + centerlens[2] , rot3 );

   TGeoCombiTrans *combi41 = new TGeoCombiTrans( "combi41" , w0lenspos[3].X() * 100 + centerlens[0] , w0lenspos[3].Y() * 100 + centerlens[1] , w0lenspos[3].Z() * 100 + centerlens[2] , rot4 );

   combi11->RegisterYourself();

   combi21->RegisterYourself();

   combi31->RegisterYourself();

   combi41->RegisterYourself();


   //------Basic shapes: a hyperboloid for each beam and a box as constraint (laser lens)------//

   TGeoBBox *lensbox = new TGeoBBox( "lensbox" , ( w0lenspos[0].Mag() + w0lenspos[1].Mag() ) * 100 / 2 , ( w0lenspos[2].Mag() + w0lenspos[3].Mag() ) * 100 / 2 , lenswidth * 100 , centerlens );

   TGeoHype *beam11 = new TGeoHype( "beam11" , 0 , 0 , w0lens[0] * 100 , beamdiverlens[0] , w0lenspos[0].Mag() * 100 );

   TGeoHype *beam21 = new TGeoHype( "beam21" , 0 , 0 , w0lens[1] * 100 , beamdiverlens[1] , w0lenspos[1].Mag() * 100 );

   TGeoHype *beam31 = new TGeoHype( "beam31" , 0 , 0 , w0lens[2] * 100 , beamdiverlens[2] , w0lenspos[2].Mag() * 100 );

   TGeoHype *beam41 = new TGeoHype( "beam41" , 0 , 0 , w0lens[3] * 100 , beamdiverlens[3] , w0lenspos[3].Mag() * 100 );



   //------Composite shape: translation + rotation of each hyperboloid intersected with box (laser lens)------//

   TGeoCompositeShape *lensbeams = new TGeoCompositeShape( "lensbeams" , "( beam11:combi11 + beam21:combi21 + beam31:combi31 + beam41:combi41 ) * lensbox" );

   TGeoVolume *laserlens = new TGeoVolume( "laserlens" , lensbeams );

   laserlens->SetLineColor( kRed );

   top->AddNode( laserlens , 1 );


   mgr->CloseGeometry();
   mgr->SetNsegments(120);
   //top->Raytrace();

   top->Draw();
}