// MakeWaves.cc
//   features:
//   * integer and double vectors and matrices

#include <iostream>
#include <fstream>  
#include <cmath>   
using namespace std;
#include "utils.h"

#define PI 3.14159265358979323846

// functions defined in this file
void MakeSinusoids( double **M , int K , int N , double *k ) ;
void DispersionRelation( double*k , double *omega , int K , int type ) ;
void Superpose( double t , double *yt, double *vt, double **M ,
		double *omega ,
		int N, int K ,  double *yk, double *vk );

void Superpose( double t , double *yt, double *vt, double **M ,
		double *omega ,
		int N, int K ,  double *yk, double *vk ){
  // Explanation:
  // turn fourier state into state at time t by superposition
  // For a single fourier coefficient, 
  // yt    =  cos(omega*t)*A       + sin(omega*t)*B   implies
  // dy/dt = -sin(omega*t)*A*omega + cos(omega*t)*B*omega
  // at time t=0, 
  // yt    =  A        (defined to be yk)   and
  // dy/dt =  B*omega  (defined to be vk)
  // so rearranging, we obtain the future sinusoid from the state
  // at time zero thus:
  // yt    =  cos(omega*t)*yk[m]       + sin(omega*t)*vk[m]/omega  
  // dy/dt = -sin(omega*t)*yk[m]*omega + cos(omega*t)*vk[m] 

  constantdvector( yt , 1 , N , 0.0 ) ;
  constantdvector( vt , 1 , N , 0.0 ) ;
  for ( int m = 1 ; m <= K ; m ++ ) {
    incrementdvector( yt ,
      cos(omega[m]*t)*yk[m] + sin(omega[m]*t) * vk[m]/omega[m] ,
		      M[m] , 1 , N );
    incrementdvector( vt ,
      cos(omega[m]*t) * vk[m] - sin(omega[m]*t)*yk[m]*omega[m] ,
		      M[m] , 1 , N );
  }
}

#define GRAVITY 10.0
// average depth of ocean in metres
#define DEPTH   3800.0
// density of water
#define RHO     1000.0
// surface tension of water 
#define SIGMA 0.075 
// 0.075 J per (m**2)
double dispersionRelation( double k , int type) {
  // returns the frequency omega associated with a wavenumber k
  return (type==1) ? k
    : (type==2) ? sqrt( k )
    : (type==3) ? sqrt( GRAVITY * k * tanh( k * DEPTH ) )
    :             sqrt( GRAVITY * k + SIGMA * k * k * k / RHO )
    ;
}
// some of these assume that
// k is in SI units, i.e. (1/m).

// The spacing of two points separated by "1" unit in the
// x-direction (n=1..N):
#define SPACING_IN_METRES 10.0
// note MakeDrop assumes a width of 5 units when the drop is made
void DispersionRelation( double*k , double *omega , int K , int type) {
  for ( int m = 1 ; m <= K ; m ++ ) {
    // omega = c k -- non-dispersive
    // omega**2 = g k   -- deep water
    // omega**2 = g k tanh(kd)   -- reasonably deep water
    // omega**2 = gk + sigma k**3/rho   -- deep water and ripples
    omega[m] = dispersionRelation( k[m] , type ) ;
  }  
}

double square(double x ) {
  return x*x;
}

// sets the velocity profile for a pair of bumps of equal and opposite
// momentum hitting the surface
void   MakeDropV( double *v , int N  ) {
  for ( int n = 1 ; n <= N ; n ++ ) {
    v[n] = ( n < N/2 - 50 || n > N/2+50 ) ? 0.0
      : -((double)(n)-(double)(N/2)) *
      exp(-square(((double)(n)-(double)(N/2)))/500.0) ;
  }  
}

// sets the surface and velocity profile for a wavepacket (or pair)
// sets the velocity only, so makes a pair of packets
void   MakePacketV( double *y , double *v , int N  ) {
  for ( int n = 1 ; n <= N ; n ++ ) {
    v[n] = ( n < N/2 - 250 || n > N/2+250 ) ? 0.0
      : cos( ((double)(n)-(double)(N/2)) *2*PI / 10.0 ) *
      exp(-square(((double)(n)-(double)(N/2)))/2500.0) ;
    y[n] = 0.0 ;
  }  
}

#define WAVELENGTH 20.0
#define WAVENUMBER (2*PI/WAVELENGTH)
// sets the surface and velocity profile for a wavepacket (or pair)
// sets the position only, so makes a pair of packets
void   MakePacketY( double *y , double *v , int N  ) {
  for ( int n = 1 ; n <= N ; n ++ ) {
    y[n] = ( n < N/4 - 250 || n > N/4+250 ) ? 0.0
      : cos( ((double)(n)-(double)(N/2)) *2*PI / WAVELENGTH ) *
      exp(-square(((double)(n)-(double)(N/4)))/2500.0) ;
    v[n] = 0.0 ;
  }  
}

// sets the surface and velocity profile for a wavepacket 
void   MakePacket( double *y , double *v , int N , int type ) {
  double A = dispersionRelation( WAVENUMBER, type ) ;
  // see  Explanation 
  for ( int n = 1 ; n <= N ; n ++ ) {
    y[n] = ( n < N/4 - 250 || n > N/4+250 ) ? 0.0
      : cos( ((double)(n)-(double)(N/2)) *WAVENUMBER ) *
      exp(-square(((double)(n)-(double)(N/4)))/2500.0) ;
    v[n] = ( n < N/4 - 250 || n > N/4+250 ) ? 0.0
      : A*sin( ((double)(n)-(double)(N/2)) *WAVENUMBER ) *
      exp(-square(((double)(n)-(double)(N/4)))/2500.0) ;
  }  
}

// sets the velocity profile for a rounded bump of momentum hitting
// the surface
void   MakeDrop( double *v , int N  ) {
  for ( int n = 1 ; n <= N ; n ++ ) {
    v[n] = ( n < N/2 - 50 || n > N/2+50 ) ? 0.0
      : -exp(-square(((double)(n)-(double)(N/2)))/50.0) ;
  }  
}

// sets the velocity profile for a sudden square bump
// of momentum hitting the surface
void   MakeDropS( double *v , int N  ) {
  for ( int n = 1 ; n <= N ; n ++ ) {
    v[n] = ( n < N/2 - 5 || n > N/2 ) ? 0.0 : -1.0 ;
  }  
}

void   MakeSinusoids( double **M , int K , int N , double *k ) {
  // Make the sinusoids' matrix
  double scale = sqrt( 2.0/(double)(N) ) ; 
  for ( int m = 1 ; m <= K ; m ++ ) {
    double (*f)(double x) ; 
    if ( m%2 ) {
      f = cos ; 
    } else {
      f = sin ; 
    }
    int nperiods = ((m+1)/2) ;
    double thisk =  2.0 * PI * (double) nperiods / (double) N  ;
    for ( int n = 1 ; n <= N ; n ++ ) {
      M[m][n] = scale * f( thisk * (double) n )  ;
      // k * N = nperiods * 2pi
    } // if M = sin,   sum M**2 = N / 2.0 ; so scale**2 = 2.0/N
    // store the physical value of k, which is
    // 2 pi n / (N*SPACING_IN_METRES)
    k[m] = thisk / SPACING_IN_METRES ;
  }
}

void   CheckSinusoids( double **M , int K , int N ) {
  cout << "assert the following K quantities are all 1\n" ; 
  for ( int m = 1 ; m <= K ; m ++ ) {
    cout << m<<" "<<innerProduct( M[m], M[m], 1, N ) << endl; 
  }
  cout << "assert the following K-1 quantities are all 0\n" ; 
  for ( int m = 2 ; m <= K ; m ++ ) {
    cout << m<<" "<<innerProduct( M[1], M[m], 1, N ) << endl; 
  }
  cout << "assert the following K-1 quantities are all 0\n" ; 
  for ( int m = 1 ; m <= K-1 ; m ++ ) {
    cout << m<<" "<<innerProduct( M[K], M[m], 1, N ) << endl; 
  }
}

int main(int argc, char* argv[])
{
  // matrix dimensions N (length of ocean), K (number of stationary waves): 
  int N=4000, K=600  ;
  double t, T=3000, dt=60.0 ; // good time for tsunami: dt=60.0
  if(argc>1)   {
    sscanf( argv[1], "%d", &N ) ; // put the first command-line argument in N
  }
  if(argc>2) {
    sscanf( argv[2], "%d", &K ) ; // put the 2nd argument in seed
  }
  if(argc>3) {
    sscanf( argv[3], "%lf", &T ) ; // 
  }
  if(argc>4) {
    sscanf( argv[4], "%lf", &dt ) ; //
  }
  double *y , *yt, *yk ; 
  double *v , *vt, *vk ; 
  double *k , *omega ; 
  double **M ;
  int verbose = 0 ; 

  // memory allocation
  omega = dvector( 1 , K ) ; // frequencies
  k = dvector( 1 , K ) ; // wavevectors
  y = dvector( 1 , N ) ; // initial condition
  v = dvector( 1 , N ) ; 
  yt = dvector( 1 , N ) ; // state at some possibly later time 
  vt = dvector( 1 , N ) ; 
  vk = dvector( 1 , K ) ; // projection
  yk = dvector( 1 , K ) ; // projection

  M = dmatrix( 1, K, 1, N ) ; 

  MakeSinusoids( M , K , N , k ) ;
  int type = 3 ; // 1: non-dispersive   2: deep water   3: reasonably deep water
  DispersionRelation( k , omega , K , type ) ;
  //  printMatrix( M , 1,K , 1,N ) ;
  if( verbose ) printMatrixT( M , 1,K , 1,N ) ;
  if( verbose ) CheckSinusoids( M , K , N  ) ;

  int tsunami=1;
  if(tsunami) {
    // make initial conditions v,y 
    MakeDrop( v , N  ) ;
    constantdvector( y , 1 , N , 0.0 ) ;
  } else {
    MakePacket(y,v,N,type);
  }
  // turn initial conditions v,y into initial conditions for the
  // fourier components
  for ( int m = 1 ; m <= K ; m ++ ) {
    yk[m] = innerProduct( y , M[m] , 1 , N );
    vk[m] = innerProduct( v , M[m] , 1 , N );
  }

  // try to write output file
  char fileName[]="/tmp/y";
  ofstream fout1 ;  
  fout1.open(fileName);
  if(fout1.good()==false){
    cerr << "can't write to file " << fileName << endl;
    exit(0);
  }

  for ( t = 0.0; t <= T ; t += dt ) {
    Superpose( t , yt , vt , M , omega , N , K , yk , vk );
    fout1 << endl;
    fout1 << endl;
    fprintVector( yt, 1, N , fout1, 0 ) ;
    fout1.flush();
  }

  // free the memory
  freedvector( omega , 1 ) ;
  freedvector( k , 1 ) ;
  freedvector( v , 1 ) ;
  freedvector( y , 1 ) ;
  freedvector( vt , 1 ) ;
  freedvector( yt , 1 ) ;
  freedvector( vk , 1 ) ;
  freedvector( yk , 1 ) ;
  freedmatrix( M , 1, K, 1, N ) ;

  return 0;
}