//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=
//    Program rabi
//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=

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

//---- experimental settings
const double dT      = 1.0/4096;       // temporal step
const int    BASES   = 4;              // number of base waves

//---- physical settings
// setting about the well
const double WEL_LEN = 1.0;            // length of the well

// setting about the wave function
const int    WAVE_MODE = 0;            // initial mode

// macros for eigen_values
inline double eigen_moment( int mode ){
  return M_PI/WEL_LEN*(mode+1);
}
inline double eigen_energy( int mode ){
  return 0.5*sqr(M_PI/WEL_LEN*(mode+1));
}

// setting about the electromagnetic-wave
const double Eamp = 4.0;               // amplitude of elecro-wave
const double Omg  = eigen_energy( WAVE_MODE+1 )
                  - eigen_energy( WAVE_MODE   ); // angular frequency

//---- graphics settings
const int    WIN_WIDTH  = 512;
const int    WIN_HEIGHT = 256;
const int    BASE_Y     = 128;         // Y-coordinate of the origin
const double WAVE_MAG   = 64.0;        // magnify factor for wave
const double POTE_MAG   = 16.0;        // magnify factor for potential
const int    GRIDS      = 64+1;        // number of GRIDS for drawing


//---- declaration of class Wave

class Wave
{
    // property of each base functions
  Complex C[BASES];                // coefficiences
  double  E0[BASES], k0[BASES];    // eigen energy and momentum
  double  d[BASES][BASES];         // matrix elements of dipole
public:
       Wave( void );
  void Evolve( void );
  void Draw( void );
private:
  double Probability( void );
  Complex Psi( double x );
};

//---- declaration of class EleMag

class EleMag
{
  double Ele;                           // Electro-field
public:
       EleMag( void ){ Ele = 0.0; }
  inline double E( void ){ return Ele; }
  void Evolve( void );
  void Draw( void );
};


//---- defintion of class Wave

//---- initializes properties
Wave::Wave( void )
{
  int n, m;

  for( n=0 ; n<BASES ; n++ ){
    if( n == WAVE_MODE ) C[n] = 1.0; else C[n] = 0.0;
    k0[n] = eigen_moment(n);
    E0[n] = eigen_energy(n);
  }

  for( n=0 ; n<BASES ; n++ ){
    for( m=0 ; m<BASES ; m++ ){
      if( (n-m)%2 != 0 ){
	d[n][m] = (8.0*WEL_LEN/(M_PI*M_PI)) * ((double)(n+1)*(m+1)/sqr(sqr(n+1)-sqr(m+1)));
      }else{
	d[n][m] = 0.0;
      }
    }
  }
}

//---- evolves the wave function
void Wave::Evolve( void )
{
  extern EleMag F;       // assumes the existence of an object F
  int n, m;

  double D = F.E();      // gets the electric field

    // symplectic integration, part I.
  for( m=0 ; m<BASES ; m++ ){
    for( n=0 ; n<BASES ; n++ ){
      if( (n-m)%2 == 0 ) continue;
      C[n] += -(dT/2)*D*d[n][m]*CI*C[m];
    }
  }

    // symplectic integration, part II.
  for( n=0 ; n<BASES ; n++ ){
    C[n] *= polar( 1.0, -dT*E0[n] );
  }

    // symplectic integration, part III.
  for( m=BASES-1 ; m>=0 ; m-- ){
    for( n=0 ; n<BASES ; n++ ){
      if( (n-m)%2 == 0 ) continue;
      C[n] += -(dT/2)*D*d[n][m]*CI*C[m];
    }
  }
}

double Wave::Probability( void )
{
  double prob=0.0;
  for( int n=0; n<BASES ; n++ ){
    prob += norm(C[n]);
  }
  return prob;
}

//---- invert sin transform
Complex Wave::Psi( double x )
{
  Complex psi=0.0;
  for( int n=0; n<BASES ; n++ ){
    psi += C[n] * sin(k0[n]*x);
  }
  return psi;
}

//---- draws wave function
void Wave::Draw( void )
{
  int i;
  static Complex psi[GRIDS];
  const double dX = WEL_LEN/(GRIDS-1); // length of a cell

    // converts into real space
  for( i=0 ; i<GRIDS ; i++ ){
    psi[i] = Psi(dX*i);
  }

    // draws probability density
  NXSetColor( NX_BLUE );
  NXDrawMoveto();
  for( i=0 ; i<GRIDS ; i++ ){
    NXDrawLineto( int(i*dX*WIN_WIDTH/WEL_LEN),
		 BASE_Y - int(norm(psi[i])*WAVE_MAG) );
  }

    // draws real part of the wave function
  NXSetColor( NX_CYAN );
  NXDrawMoveto();
  for( i=0 ; i<GRIDS ; i++ ){
    NXDrawLineto( int(i*dX*WIN_WIDTH/WEL_LEN),
		 BASE_Y - int(psi[i].real*WAVE_MAG) );
  }

    // draws imaginaly part of the wave function
  NXSetColor( NX_GRASS );
  NXDrawMoveto();
  for( i=0 ; i<GRIDS ; i++ ){
    NXDrawLineto( int(i*dX*WIN_WIDTH/WEL_LEN),
		 BASE_Y - int(psi[i].imag*WAVE_MAG) );
  }

    // draws coefficiences as rectangle bars
  NXSetColor( NX_BLACK );
  for( int n=0 ; n<BASES ; n++ ){
    int h = int(128*norm(C[n]));
    NXFillRectangle( n*32+16, WIN_HEIGHT-1-h, 24, +h );
  }

    // prints some values
  NXDrawPrintf( 16, 16, "Probability = %lf", Probability() );
}





//---- definition of class EleMag

//---- updates electro-field
void EleMag::Evolve( void )
{
  extern double T;
  Ele = Eamp*sin(Omg*T);
}

//---- draws electro-field
void EleMag::Draw( void )
{
  NXSetColor( NX_BLACK );
  NXDrawLine( int(0*WIN_WIDTH/WEL_LEN),
              BASE_Y - int(Ele*0*POTE_MAG),
              int(WEL_LEN*WIN_WIDTH/WEL_LEN),
              BASE_Y - int(Ele*WEL_LEN*POTE_MAG) );
}


//---- definition of fundermental objects
double T = 0.0;
Wave   W;
EleMag F;


int main( void )
{
  NXOpenWindow( "Rabi oscillation", WIN_WIDTH, WIN_HEIGHT );
  NXEnable( NX_DOUBLEBUFFER );

  while( NXNoEvents() ){
    for( int i=0 ; i<32 ; i++ ){
      F.Evolve();
      W.Evolve();
      T += dT;
    }
    NXClearWindow();
    F.Draw();
    W.Draw();
    NXFlush();
  }

  NXCloseWindow();

  return 0;
}