//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=
//    Program scatt
//=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=*=

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

//---- experimental settings
const double dT      = 1.0/4096;       // temporal step 
const int    GRIDS   = 256+1;          // number of GRIDS points

//---- physical settings
// settings about the system
const double SYS_LEN = 8.0;            // length of the system
const double dX      = SYS_LEN/(GRIDS-1); // length of a cell, dont edit

// settings about initial conditions of the Gaussian wave packet
const double PAC_POS = 2.5;            // initial position
const double PAC_DX  = 1.0/4;          // initial width
const double PAC_MOM = 24.0;           // initial momentum
const double PAC_ENE = sqr(PAC_MOM)/2; // kinetic energy, dont edit

// settings about the potential barrier
const double BAR_POS = 4.0;            // position
const double BAR_LEN = 1.0/8;          // length
const double BAR_HIG = PAC_ENE*1.25;   // height or barrier energy


//---- graphics settings
const int    WIN_WIDTH  = 512;
const int    WIN_HEIGHT = 256;
const int    BASE_Y     = 228;         // Y-coordinate of the origin
const double WAVE_MAG   = 64.0;        // magnify factor for wave
const double POTE_MAG   = 1.0/4;       // magnify factor for potential



//---- declaration of class Wave

class Wave
{
  Complex Psi[GRIDS];
public:
       Wave( void );
  void Evolve( void );
  void Draw( void );
private:
  void TriDiagEQ( Complex A[], Complex x[], Complex b[] );
  double Probability( void );
  double Energy( void );
};


//---- declaration of class Potential

class Potential
{
  double V[GRIDS];
public:
  Potential( void );
  inline double operator [] ( int i ){ return V[i]; }
  void Draw( void );
};

  

//---- defintion of class Wave

//---- sets Gaussian wave packet
Wave::Wave( void )
{
  const double Amp = pow( 2*M_PI*sqr(PAC_DX), -0.25 );
  for( int i=1 ; i<GRIDS-1 ; i++ ){
    double x = dX*i;
    Psi[i] = Amp * exp( -sqr((x-PAC_POS)/(2*PAC_DX)) ) * polar( 1.0, PAC_MOM*x );
  }
  Psi[0] = Psi[GRIDS-1] = 0.0;
}



//---- evolves the wave function by using Ceyley's form
void Wave::Evolve( void )
{
  extern Potential V; // assumes the existence of an object V

  static Complex b[GRIDS], A[GRIDS], B[GRIDS];

  for( int i=1 ; i<=GRIDS-2 ; i++ ){
    B[i] = CI*4*dX*dX/dT + (2 + V[i]*2*dX*dX)*C1;
    A[i] = CI*4*dX*dX/dT - (2 + V[i]*2*dX*dX)*C1;
    b[i] = B[i]*Psi[i] - Psi[i+1] - Psi[i-1];
  }

  TriDiagEQ( A, Psi, b );
}

//---- soleves equations of special tri-diagonal matrix
void Wave::TriDiagEQ( Complex A[], Complex x[], Complex b[] )
{
  static Complex u[GRIDS];
  int i;

  u[1] = 1.0/A[1];
  x[1] = b[1]*u[1];
  for( i=2 ; i<=GRIDS-2 ; i++ ){
    u[i] = 1.0/(A[i]-u[i-1]);
    x[i] = (b[i]-x[i-1])*u[i];
  }
  for( i=GRIDS-3 ; i>=1 ; i-- ){
    x[i] -= x[i+1]*u[i];
  }
}

//---- calculates integral of probability density
double Wave::Probability( void )
{
  double prob=0.0;
  for( int i=0 ; i<GRIDS ; i++ ){
    prob += norm(Psi[i])*dX;
  }
  return prob;
}

//---- calculates expectational energy
double Wave::Energy( void )
{
  extern Potential V;
  double energy=0.0;

  for( int i=1 ; i<GRIDS-1 ; i++ ){
    energy += real( conj(Psi[i]) * (
      (2*Psi[i] - Psi[i+1] - Psi[i-1])/(2*dX*dX) + V[i]*Psi[i] ) )*dX;
  }
  return energy;
}


//---- draws wave function
void Wave::Draw( void )
{
  int i;

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

    // prints some values
  NXSetColor( NX_BLACK );
  NXDrawPrintf( 16, WIN_HEIGHT-16, "Probability = %lf", Probability() );
  NXDrawPrintf( 16, WIN_HEIGHT-2,  "Energy = %lf", Energy() );
}



//---- definition of class Potential

//---- sets potential
Potential::Potential( void )
{
  for( int i=0 ; i<GRIDS ; i++ ){
    if( BAR_POS - BAR_LEN/2 < dX*i && dX*i < BAR_POS + BAR_LEN/2 ){
      V[i] = BAR_HIG;
    }else{
      V[i] = 0.0;
    }
  }
}

//---- draws potential
void Potential::Draw( void )
{
  NXSetColor( NX_GREEN );
  NXDrawMoveto();
  for( int i=0 ; i<GRIDS ; i++ ){
    NXDrawLineto( int(i*dX*WIN_WIDTH/SYS_LEN),
		 BASE_Y - int(V[i]*POTE_MAG) );
  }
}


//---- definition of fundermental objects
double    T=0.0;
Potential V;
Wave      W;

int main( void )
{
  NXOpenWindow( "Scattering of wave packets", WIN_WIDTH, WIN_HEIGHT );
  NXEnable( NX_DOUBLEBUFFER );

  while( NXNoEvents() ){
    for( int i=0 ; i<8 ; i++ ){
      W.Evolve();
      T += dT;
    }

    NXClearWindow();
    W.Draw();
    V.Draw();
    NXFlush();
  }

  NXCloseWindow();

  return 0;
}