//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*
//   Program shower
//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*

#include "cip.h"
#include "nxgraph.h"

//---- definition of fundermental constants
//---- physical settings
const double N_A         = 6.02e23;   // Avogadro's number
const double R_E         = 2.814e-15; // classical electron radius [m]
const double M_ELE       = 0.510999066; // mass of electron [MeV]

//---- air settings
const double Z_AIR       = 7.2;       // mean atomic number of atom in air
const double LN_Z_AIR    = 2.00;      // log(Z_AIR)
const double A_AIR       = 14.4;      // mean mass number of atom in air
const double RHO_AIR     = 1.2;       // density of air [g/little]==[kg/m^3]
const double C_AIR       = 1.4e-4;    // factor of rho decresing [1/m]
const double SIGMA_O     = ALPHA*Z_AIR*Z_AIR*R_E*R_E; // don't edit
const double FACT_Z      = 5.21-LN_Z_AIR/3.0;         // don't edit

//---- experimental settings
const double INIT_PHOTON_HEIGHT= 5e4; // [m]
const double INIT_PHOTON_ENERGY= 1e3; // [MeV]
const double ENERGY_CUTOFF     = 1e2; // least energy [MeV], below thie energy, particles are treated as dead
const double BREMS_CUTOFF = 0.1;      // least photon energy radiated by bremsstrahlung [MeV]
const double FAR_AWAY     = 1e5;       // long distance [m]

//---- particle type indicators
const int    PHOTON       = 0;
const int    ELECTRON     = 1;
const int    POSITRON     = 2;

//---- reaction indicators
const int    FATE_NON     = 0;
const int    FATE_BRE     = 1;
const int    FATE_PRO     = 2;
const int    FATE_ANN     = 3;
 
//---- graphics settings
const int    WIN_WIDTH    = 256;
const int    WIN_HEIGHT   = 256;
const int    GROUND_HEIGHT= 16;
const int    XO           = WIN_WIDTH/2;
const int    YO           = WIN_HEIGHT-GROUND_HEIGHT;
const double YMAG         = (WIN_HEIGHT-GROUND_HEIGHT)/INIT_PHOTON_HEIGHT;
const double XMAG         = YMAG*128;



class Position
{
public:
  double x, y, z;
public:
  inline Position( void ){}
  inline Position( const double& _x, const double& _y, const double& _z ){
    x = _x;    y = _y;    z = _z;
  }
  inline void operator+=( const Position& _r ){
    x += _r.x;    y += _r.y;    z += _r.z;
  }
};

class Direction
{
public:
  double norm;
  double rot[3][3];
public:
  Direction( void ){};
  Direction( double _norm, double _theta, double _phi ){
    norm  = _norm;
    double ct = cos(_theta),    st = sin(_theta);
    double cp = cos(_phi),      sp = sin(_phi);
    rot[0][0] = +ct*cp;    rot[0][1] = -sp;    rot[0][2] = st*cp;
    rot[1][0] = +ct*sp;    rot[1][1] = +cp;    rot[1][2] = st*sp;
    rot[2][0] = -st   ;    rot[2][1] =   0;    rot[2][2] = ct   ;
  }

  inline Position operator*( const double& k ){
    return Position( rot[0][2]*k, rot[1][2]*k, rot[2][2]*k );
  }
  Direction& Turn( const Direction& _dir ){
    double t[3][3];
    int i, j, k;
      // multiply matrix rot[][] by given _dir's matrix
    for( i=0 ; i<3 ; i++ ){
      for( j=0 ; j<3 ; j++ ){
        t[i][j] = 0.0;
        for( k=0 ; k<3 ; k++ ){
          t[i][j] += _dir.rot[i][k] * rot[k][j];
        }
      }
    }
       // update matrix rot[][]
    for( i=0 ; i<3 ; i++ ){
      for( j=0 ; j<3 ; j++ ){
        rot[i][j] = t[i][j];
      }
    }
    return *this;
  }

};




static double CV, SV;

//---- Density of the air at given height
inline double N_AIR( double h )
{
  return (RHO_AIR*1000/A_AIR*N_A)*exp(-h*C_AIR);
}

//---- transform real coordinates into window coordinates
inline int Wx( double x, double y, double z )
{
  return XO + int(XMAG*(CV*x-SV*y));
}
inline int Wy( double x, double y, double z )
{
  return YO - int(YMAG*z) + int((XMAG/8)*(SV*x+CV*y));
}

//---- declaration of class Particle

class Particle
{
protected:
  char      type;                    // particle name
  char      fate;                    // reaction
  u_long    color;                   // color
  
  Position  ri, rf;                  // initial and final position
  Direction dir;                     // direction and energy
  double    free_path;               // length of its free path

  Particle* next;                    // pointer to next particle
public:
       Particle( const Position& _r, const Direction& _dir, Particle* _next=NULL );
  virtual ~Particle( void );

  void Evolve( void );
  void Draw( void );

protected:
  double FreePath( double sigma );
  virtual void Reaction( void )=0;

private:
  void CheckAlive( void );
};

//---- declaration of class Photon

class Photon : public Particle
{
public:
  Photon( const Position& _r, const Direction& _dir, Particle* _next=NULL );

private:
  void Reaction( void );
  void PairProduction( void );
};

//---- declaration of class Electron

class Electron : public Particle
{
public:
  Electron( const Position& _r, const Direction& _dir, Particle* _next=NULL );

private:
  void Reaction( void );
  void Bremsstrahlung( void );
};

//---- declaration of class Positron

class Positron : public Particle
{
public:
  Positron( const Position& _r, const Direction& _dir, Particle* _next=NULL );

private:
  void Reaction( void );
  void PairAnnihilation( void );
  void Bremsstrahlung( void );
};

//---- definition of methods under class Particle

//---- base constructor
Particle::Particle( const Position& _r, const Direction& _dir, Particle* _next )
{
  rf = ri = _r;    dir = _dir;    next = _next;
}

//---- recursive destructor
Particle::~Particle( void )
{
  if( next ) delete next;
}

//---- recursive evolve
void Particle::Evolve( void )
{
  rf += dir*free_path;
  CheckAlive();
  Reaction();
  if( next ) next->Evolve();
}

//---- recursive drawing
void Particle::Draw( void )
{
  int xi = Wx( ri.x, ri.y, ri.z );
  int yi = Wy( ri.x, ri.y, ri.z );
  int xf = Wx( rf.x, rf.y, rf.z );
  int yf = Wy( rf.x, rf.y, rf.z );

  NXSetColor( color );
  NXDrawLine(   xi, yi, xf, yf );
  NXDrawCircle( xf, yf, 4 );

  if( next ) next->Draw();
}

//---- check whether this particle is still alive
void Particle::CheckAlive( void )
{
  if( rf.z < 0 || dir.norm < ENERGY_CUTOFF ) fate = FATE_NON;
} 

//---- draw trace of particle from birth point to reaction point

//---- calculate length of free path
double Particle::FreePath( double sigma )
{
    // factor of decreasing density toward this direction.
  const double a = C_AIR*dir.rot[2][2];  // rot[2][2] == cos(theta)
    // number of targets
  const double n = -log(Urand0())/sigma;
    // temporal value
  const double t = 1-n*a/N_AIR(ri.z);

    // t<0 means the particle goes back to the universe
  if( t<0 ) return FAR_AWAY;

  return -log(t)/a;
}



//---- definition of methods under class Photon

//---- constructor for photon
Photon::Photon( const Position& _r, const Direction& _dir, Particle* _next )
:Particle( _r, _dir, _next )
{
  type      = PHOTON;
  color     = NX_ORANGE;

  free_path = FreePath( (28.0/9.0)*SIGMA_O*FACT_Z );
  fate      = FATE_PRO;
}

//---- select reaction
void Photon::Reaction( void )
{
  switch( fate ){
    case FATE_PRO : PairProduction(); break;
  }
}


//---- let photon react pair production
void Photon::PairProduction( void )
{
  double eps;                    // friction of energy carried by radiated electron
  double Ene_g  = dir.norm;      // energy of parent photon
  double Ene_1, Ene_2;           // energys of electron and positron
  double the_1, the_2;           // polar angles of each particles
  double phi_1, phi_2;           // azimuthal angles of each particles

  // get a sample of radiation energy
  double epso   = M_ELE/Ene_g;
  double lnepso = log(epso);
  do{
    eps = epso*exp(-Urand()*lnepso);
  }while( !(Urand() < eps-1+3.0/(4.0*eps) ) );
  Ene_1 = Ene_g*eps;
  
  // get a sample of radiation angle
  double A = 1.6*M_ELE/Ene_1;
  double D = 3.0;
  do{
    if( Urand()<1.0/(1.0+D) ){
      the_1 = -log(Urand0()*Urand0())*A;
    }else{
      the_1 = -log(Urand0()*Urand0())*A/3;
    }
  }while( !(the_1 < M_PI) );
  phi_1 = 2*M_PI*Urand();
  
  // calculate about another particle
  Ene_2 = hypot( Ene_1*sin(the_1), Ene_g-Ene_1*cos(the_1) );
  the_2 = atan2( Ene_1*sin(the_1), Ene_g-Ene_1*cos(the_1) );
  phi_2 = phi_1-M_PI;

  // prepare new direction
  Direction dir_1( Ene_1, the_1, phi_1 );
  Direction dir_2( Ene_2, the_2, phi_2 );

  // make child particles
  if( Urand() < 0.5 ){  // determine whether No.1 particle is electron or positron
    next = new Electron( rf, dir_1.Turn( dir ), next );
    next = new Positron( rf, dir_2.Turn( dir ), next );
  }else{
    next = new Positron( rf, dir_1.Turn( dir ), next );
    next = new Electron( rf, dir_2.Turn( dir ), next );
  }
}


//---- definition of methods under class Electron

//---- constructor for electron
Electron::Electron( const Position& _r, const Direction& _dir, Particle* _next )
:Particle( _r, _dir, _next )
{
  type      = ELECTRON;
  color     = NX_GREEN;


  double sigma_bre = (16.0/3.0)*SIGMA_O*FACT_Z*
    (log(dir.norm/BREMS_CUTOFF)-0.625);
  if( sigma_bre > 0.0 ){
    free_path = FreePath(sigma_bre);
    fate      = FATE_BRE;
  }else{
    free_path = FAR_AWAY;
    fate      = FATE_NON;
  }

}

//---- select reaction
void Electron::Reaction( void )
{
  switch( fate ){
    case FATE_BRE: Bremsstrahlung(); break;
  }
}

//---- let electron react bremsstrahlung
void Electron::Bremsstrahlung( void )
{
  double eps;                    // friction of energy carried by radiated photon
  double Ene_e  = dir.norm;      // energy of parent electron
  double Ene_g, Ene_d;           // energys of photon and decerated electron
  double the_g, the_d;           // polar angles of each particles
  double phi_g, phi_d;           // azimuthal angles of each particles

  // get a sample of radiation energy
  double epso   = BREMS_CUTOFF/Ene_e;
  double lnepso = log(epso);
  do{
    eps = epso*exp(-Urand()*lnepso);
  }while( !(Urand() < (4.0/3.0)*eps*eps - eps + 1.0 ) );
  Ene_g = Ene_e*eps;

  // get a sample of radiation angle
  double A = 1.6*M_ELE/Ene_e;
  double D = (1.16+1.88/Z_AIR)*(1+eps);
  do{
    if( Urand()<1.0/(1.0+D) ){
      the_g = -log(Urand0()*Urand0())*A;
    }else{
      the_g = -log(Urand0()*Urand0())*A/3;
    }
  }while( !(the_g < M_PI) );
  phi_g = 2*M_PI*Urand();
  
  // calculate about another particle
  Ene_d = hypot( Ene_g*sin(the_g), Ene_e-Ene_g*cos(the_g) );
  the_d = atan2( Ene_g*sin(the_g), Ene_e-Ene_g*cos(the_g) );
  phi_d = phi_g-M_PI;
  
  // prepare new direction
  Direction dir_d( Ene_d, the_d, phi_d );
  Direction dir_g( Ene_g, the_g, phi_g );

  // make child particles
  next = new Electron( rf, dir_d.Turn( dir ), next );
  next = new Photon(   rf, dir_g.Turn( dir ), next );
}



//---- definition of methods under class Positron

//---- constructor for positron
Positron::Positron( const Position& _r, const Direction& _dir, Particle* _next )
:Particle( _r, _dir, _next )
{
  double gamma, sigma_ann, sigma_bre, path_ann, path_bre;

  type     = POSITRON;
  color    = NX_RED;
  gamma    = dir.norm/M_ELE;
  sigma_ann = M_PI*Z_AIR*R_E*R_E*(log(2*gamma)-1)/gamma;
  sigma_bre = (16.0/3.0)*SIGMA_O*FACT_Z*
    (log(dir.norm/BREMS_CUTOFF)-0.625);

  if( sigma_ann > 0.0 ){
    path_ann = FreePath( sigma_ann );
  }else{
    path_ann = FAR_AWAY;
  }
  if( sigma_bre > 0.0 ){
    path_bre = FreePath( sigma_bre );
  }else{
    path_bre = FAR_AWAY;
  }

  if( path_ann == FAR_AWAY && path_bre == FAR_AWAY ){
    free_path = FAR_AWAY;
    fate      = FATE_NON;
  }else if( path_ann < path_bre ){
    free_path = path_ann;
    fate      = FATE_ANN;
  }else{
    free_path = path_bre;
    fate      = FATE_BRE;
  }
}

//---- select reaction
void Positron::Reaction( void )
{
  switch( fate ){
    case FATE_BRE: Bremsstrahlung(); break;
    case FATE_ANN: PairAnnihilation(); break;
  }
}

//---- let electron react bremsstrahlung
void Positron::Bremsstrahlung( void )
{
  double eps;                    // friction of energy carried by radiated photon
  double Ene_p  = dir.norm;      // energy of parent positron
  double Ene_g, Ene_d;           // energys of photon and decerated positron
  double the_g, the_d;           // polar angles of each particles
  double phi_g, phi_d;           // azimuthal angles of each particles
  
  // get a sample of radiation energy
  double epso   = BREMS_CUTOFF/Ene_p;
  double lnepso = log(epso);
  do{
    eps = epso*exp(-Urand()*lnepso);
  }while( !(Urand() < (4.0/3.0)*eps*eps - eps + 1.0 ) );
  Ene_g = Ene_p*eps;

  // get a sample of radiation angle
  double A = 1.6*M_ELE/Ene_p;
  double D = (1.16+1.88/Z_AIR)*(1+eps);
  do{
    if( Urand()<1.0/(1.0+D) ){
      the_g = -log(Urand0()*Urand0())*A;
    }else{
      the_g = -log(Urand0()*Urand0())*A/3;
    }
  }while( !(the_g < M_PI) );
  phi_g = 2*M_PI*Urand();
  
  // calculate about another particle
  Ene_d = hypot( Ene_g*sin(the_g), Ene_p-Ene_g*cos(the_g) );
  the_d = atan2( Ene_g*sin(the_g), Ene_p-Ene_g*cos(the_g) );
  phi_d = phi_g-M_PI;

  // prepare new direction
  Direction dir_d( Ene_d, the_d, phi_d );
  Direction dir_g( Ene_g, the_g, phi_g );
  // make child particles
  next = new Positron( rf, dir_d.Turn( dir ), next );
  next = new Photon(   rf, dir_g.Turn( dir ), next );
}

//---- let positron react pair annhilation
void Positron::PairAnnihilation( void )
{
  double eps;                    // friction of energy carried by one photon
  double Ene_p   = dir.norm;     // energy of parent positron
  double Ene_g1, Ene_g2;         // energys of radiated photons
  double the_g1, the_g2;         // polar angles of radiated photons
  double phi_g1, phi_g2;         // azimuthal angles of radiated photons

  // get a sample of radiation energy
  double epso   = M_ELE/(2*Ene_p);
  double lnepso = log(epso);
  do{
    eps = epso*exp(-Urand()*lnepso);
  }while( !(Urand() < 1-eps) );
  Ene_g1 = Ene_p*(eps);
  the_g1 = acos( 1.0 - (1.0-eps)/((eps)*Ene_p/M_ELE) );
  phi_g1 = 2*M_PI*Urand();

  // calculate about another particle
  Ene_g2 = Ene_p*(1-eps);
  the_g2 = acos( 1.0 - (eps)/((1.0-eps)*Ene_p/M_ELE) );
  phi_g2 = phi_g1-M_PI;

  // prepare new direction
  Direction dir_g1( Ene_g1, the_g1, phi_g1 );
  Direction dir_g2( Ene_g2, the_g2, phi_g2 );
  // make child particles
  next = new Photon( rf, dir_g1.Turn( dir ), next );
  next = new Photon( rf, dir_g2.Turn( dir ), next );
}

int main(void)
{
  XEvent   ev;
  NXOpenWindow( "Elemag Shower", WIN_WIDTH, WIN_HEIGHT );
  NXEnable( NX_DOUBLEBUFFER );

  Position   r( 0, 0, INIT_PHOTON_HEIGHT );
  Direction  dir( INIT_PHOTON_ENERGY, M_PI, 0 );

  do{
    Particle* parent = new Photon( r, dir );
    parent->Evolve();

    for( double the=0.0 ; the < 2*M_PI ; the+=0.0625 ){
      CV = cos(the);    SV = sin(the);
      NXClearWindow();
      parent->Draw();
      NXFlush();
      
      NXCheckEvent( NX_NOWAIT, ev );
      if( ev.type == KeyPress ){
        if( ev.xkey.keycode == 'q' ) goto LOOPOUT;
        if( ev.xkey.keycode == 'n' ) break;
      }
    }

    delete parent;
  }while(1);

 LOOPOUT:;

  NXCloseWindow();
  return 0;
}