//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=
//    Program orbital
//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=

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

//---- definition of fundermental constants
//---- physical settings
const double Z = 1.0;

//---- experimental settings
const double DX           = 0.0625;        // size of cell
const int    R_NUM        = 256;           // total number of cell
const int    R_FAR        = R_NUM-1;       // index of farthest cell
const double R_O          = 160;           // index of Bohr radius
const int    TH_NUM       = 128;           // number of polar cell
const double DTH          = M_PI/TH_NUM;   // step of polar angle
const int    PH_NUM       = 128;           // number of azimuthal cell
const double DPH          = 2*M_PI/PH_NUM; // step of azimuthal angle

//---- graphics settings
const int    PLOT_NUM     = 2048;
const int    WIN_HEIGHT   = 256;
const int    WIN_WIDTH    = 256;
const double MAG          = 8.0;

//---- prototypes of functions
void CalcRadial( int n, int l, double Frad[] );
void CalcPolar( int l, int m, double Fthe[] );
void CalcAzimuthal( int m, double Fphi[] );
void Accumulate( int n, double f[], double S[] );
int  Sample( int n, double S[] );

void DetectOrbital( short x[], short y[], short z[], double Srad[], double Sthe[], double Sphi[] );
void DrawOrbital( double the, short x[], short y[], short z[] );

//---- table of Radius
class RTable
{
  double r[R_NUM];
public:
  RTable( void ){
    for( int i=0 ; i<=R_FAR ; i++ ){
      r[i] = (1.0/Z)*exp(DX*(i-R_O));
    }
  }
  inline double operator [] ( int i ){
    return r[i];
  }
} R;

//---- main function
int main( void )
{
  int    n, l, m;
  double Frad[R_NUM], Fthe[TH_NUM], Fphi[PH_NUM];
  double Srad[R_NUM], Sthe[TH_NUM], Sphi[PH_NUM];
  short  x[PLOT_NUM], y[PLOT_NUM], z[PLOT_NUM];

  XEvent ev;
  NXOpenWindow("Orbital of Hydrogen Atom", WIN_WIDTH, WIN_HEIGHT );
  NXEnable( NX_DOUBLEBUFFER );

  for( l=0 ; l<=2 ; l++ ){
    n = l+1;
    for( m=+l ; m>=-l ; m-- ){
      printf("(n,l,m)=(% d,% d,% d)\n", n, l, m );

      CalcRadial( n, l, Frad );
      Accumulate( R_NUM,  Frad,  Srad );
      CalcPolar( l, m, Fthe );
      Accumulate( TH_NUM, Fthe, Sthe );
      CalcAzimuthal( m, Fphi );
      Accumulate( PH_NUM, Fphi, Sphi );
      DetectOrbital( x, y, z, Srad, Sthe, Sphi );
      
      for( double the=0.0 ; the < 2*M_PI ; the+=0.0625 ){
        DrawOrbital( the, x, y, z );
        NXCheckEvent( NX_NOWAIT, ev );
        if( ev.type == KeyPress ){
          if( ev.xkey.keycode == 'q' ) goto LOOPOUT;
          if( ev.xkey.keycode == 'n' ) break;
        }
      }
    }
  }
 LOOPOUT:;
  NXCloseWindow();

  return 0;
}


//---- calculate distribution of radial wave function
void CalcRadial( int n, int l, double Frad[] )
{
  if( n==1 ){
    for( int i=0 ; i<R_NUM ; i++ ){
      double r = R[i];
      Frad[i] = r*r*sqr(exp(-r/1));
    }
  }else if( n==2 ){
    for( int i=0 ; i<R_NUM ; i++ ){
      double r = R[i];
      Frad[i] = r*r*sqr(r*exp(-r/2));
    }
  }else if( n==3 ){
    for( int i=0 ; i<R_NUM ; i++ ){
      double r = R[i];
      Frad[i] = r*r*sqr(r*r*exp(-r/3));
    }
  }
}

//---- calculate distribution of polar wave function
void CalcPolar( int l, int m, double Fthe[] )
{
  int i, it;
  double z, z2, a, s;

  m = abs(m);
  for( it=0 ; it<TH_NUM ; it++ ){
    z  = cos(DTH*it);
    z2 = z*z;
    if( (l+m)%2 ){ // odd
      i = 3;
      a = z;
    }else{         // even
      i = 2;
      a = 1;
    }
    for( s=a ; i<=l-m ; s+=a, i+=2 ){
      a *= z2*(i-l+m-2)*(i+l+m-1)/(i*(i-1));
    }
    Fthe[it] = sin(DTH*it) * pow( 1-z2, m ) * s*s;
  }
}

//---- calculate distribution of azimuthal wave function
void CalcAzimuthal( int m, double Fphi[] )
{
  int i;
  for( i=0 ; i<PH_NUM ; i++ ){
    if( m == 0 ){
      Fphi[i] = 1;
    }else if( m > 0 ){
      Fphi[i] = sqr(cos(m*DPH*i));
    }else if( m < 0 ){
      Fphi[i] = sqr(sin(m*DPH*i));
    }
  }
}
  
//---- accumulate distribution
void Accumulate( int n, double f[], double S[] )
{
  int i;
  double sum, norm;
    // accumulate area
  for( i=0, sum=0.0 ; i<n ; i++ ) S[i] = (sum+=f[i]);
    // normalize area
  for( i=0, norm=1.0/sum ; i<n ; i++ ) S[i] *= norm;
}

//---- take a sample from given accumulate distribution
int  Sample( int n, double S[] )
{
  int i;
  double r = Urand();

  for( i=0 ; i<n ; i++ ) if( S[i] > r ) break;
  return i;
}

//---- detect electron many times
void DetectOrbital( short x[], short y[], short z[], double Srad[], double Sthe[], double Sphi[]  )
{
  for( int i=0 ; i<PLOT_NUM ; i++ ){
    double rad = R[ Sample( R_NUM, Srad  ) ];
    double the = DTH*Sample( TH_NUM, Sthe );
    double phi = DPH*Sample( PH_NUM, Sphi );
    
    x[i] = short(MAG*rad*sin(the)*cos(phi));
    y[i] = short(MAG*rad*sin(the)*sin(phi));
    z[i] = short(MAG*rad*cos(the)         );
  }
}

//---- rotation matrix
static double CV, SV;

//---- transform real coordinates into window coordinates
inline int Wx( short x, short y, short z )
{
  return (WIN_WIDTH/2) + int(CV*x-SV*y);
}
inline int Wy( short x, short y, short z )
{
  return (WIN_HEIGHT/2) - int(z) + int((1.0/4)*(SV*x+CV*y));
}

//---- draw orbital
void DrawOrbital( double the, short x[], short y[], short z[] )
{
  CV = cos(the);    SV = sin(the);
  NXClearWindow();
  for( int i=0 ; i<PLOT_NUM ; i++ ){
    NXDrawPoint( Wx(x[i],y[i],z[i]), Wy(x[i],y[i],z[i]) );
  }

  const int AXIS = WIN_WIDTH/2;
  NXDrawLine( Wx(-AXIS,0,0), Wy(-AXIS,0,0), Wx(+AXIS,0,0), Wy(+AXIS,0,0) );  // X-axis
  NXDrawLine( Wx(0,-AXIS,0), Wy(0,-AXIS,0), Wx(0,+AXIS,0), Wy(0,+AXIS,0) );  // Y-axis
  NXDrawLine( Wx(0,0,-AXIS), Wy(0,0,-AXIS), Wx(0,0,+AXIS), Wy(0,0,+AXIS) );  // Z-axis
  NXFlush();
}