/*	Extrap.c

	Wavefield extrapolation program for monochromatic waves.
	Provides a number of x-z timeslices that may be displayed as a movie.
	From Claerbout (1985) p. 105.

	Usage:		extrap >outfile

	Compile:	cc extrap.c ctris.c -fswitch -lm -o extrap
									*/
/* The usual includes */
#include <stdio.h>
#include <math.h>

/* Define complex multiplication and conjugate for declared struct complex */
#define  rmul(x,y)      (x.re * y.re - x.im * y.im)
#define  imul(x,y)      (x.im * y.re + x.re * y.im)
#define rcmul(x,y)      (x.re * y.re + x.im * y.im)
#define icmul(x,y)      (x.im * y.re - x.re * y.im)

/* Define complex division -- x/y for declared struct complex */
#define  rdiv(x,y) ((x.re * y.re + x.im * y.im)/(y.re * y.re + y.im * y.im))
#define  idiv(x,y) ((x.im * y.re - x.re * y.im)/(y.re * y.re + y.im * y.im))

/* Declare the structure to hold complex numbers */
struct complex 
	{
	float re; 
	float im;
	};

main()
	{
	int ix, nx, iz, nz, iw, nw, it, nt;
	float p[12][48][96], phase, pi2, dx, dz, v;
	float z0, x0, dt, dw, lambda, w, wov, x;
	struct complex cd[48], ce[48], cf[48], q[48], aa, a, b, c, cshift;
	struct complex num, den, num2;

	/* Set dimensions */
	nt = 12;
	nx = 48;
	nz = 96;
	dx = 2;
	dz = 1;
	nw = 2;
	/* Initialize constants */
	pi2 = 2.0*3.141592;
	v = 1;
	lambda = nz*dz/4;
	dw = v*pi2/lambda;
	dt = pi2/(nt*dw);

	/* Zero out data volume */
	for (it=0; it<nt; it++)
		for (ix=0; ix<nx; ix++)
			for (iz=0; iz<nz; iz++)
				p[it][ix][iz] = 0;

	/* Superimpose nw frequencies */
	for (iw=0; iw<nw; iw++)
		{
		w = (iw+1)*dw;
		/* Frequency / velocity */
		wov = w/v;
		x0 = nx*dx/3;
		z0 = nz*dz/3;

		/* Initial conditions for a collapsing spherical wave */
		for (ix=0; ix<nx; ix++)
			{
			x = (ix+1)*dx - x0;
			phase = -wov*(float)(sqrt((double)(z0*z0 + x*x)));
			/* q[ix] = e^(i phase) */
			q[ix].re = (float)(cos((double)(phase)));
			q[ix].im = (float)(sin((double)(phase)));
			}

		/* Tridiagonal matrix coefficients */
		num.re = dz;
		num.im = 0;
		den.re = 0;
		den.im = -4*wov*dx*dx;
		/* aa = dz/(-4 i wov dx^2) */
		aa.re = rdiv(num,den);
		aa.im = idiv(num,den);
		/* Diagonals of the matrix */
		a.re = -aa.re;
		a.im = -aa.im;
		b.re = 1 + 2*aa.re;
		b.im = 2*aa.im;
		c.re = -aa.re;
		c.im = -aa.im;

		/* Extrapolation in depth */
		for (iz=0; iz<nz; iz++)
			{
			/* Diffraction term */
			for (ix=1; ix<nx-1; ix++)
				{
				/* Apply finite difference star */
				/* cd[ix]=aa*q[ix+1]+(1-2aa)q[ix]+aa*q[ix-1] */
				num.re = 1 - 2*aa.re;
				num.im = -2*aa.im;
				/* Can only put one complex multiplication
				   macro on each line */
				num2.re = rmul(aa,q[ix-1]);
				num2.im = imul(aa,q[ix-1]);
				num2.re = rmul(num, q[ix]) + num2.re;
				num2.im = imul(num, q[ix]) + num2.im;
				cd[ix].re = rmul(aa,q[ix+1]) + num2.re;
				cd[ix].im = imul(aa,q[ix+1]) + num2.im;
				}
			/* Set side conditions */
			cd[0].re = 0;
			cd[0].im = 0;
			cd[nx-1].re = 0;
			cd[nx-1].im = 0;

			/* Set the left and right ends of the diagonal */
			/* Put endl into a complex buffer */
			num.re = -a.re;
			num.im = -a.im;
			/* Put endr into a complex buffer */
			num2.re = -c.re;
			num2.im = -c.im;
			/* Solve the complex tridiagonal system */
			ctris(nx, &num, &a, &b, &c, &num2, cd, q, ce, cf);

			/* cshift = e^(i wov dz) */
			cshift.re = (float)(cos((double)(wov*dz)));
			cshift.im = (float)(sin((double)(wov*dz)));

			/* Shifting term */
			for (ix=0; ix<nx; ix++)
				{
				num.re = rmul(q[ix],cshift);
				num.im = imul(q[ix],cshift);
				q[ix].re = num.re;
				q[ix].im = num.im;
				}
			
			/* Evolution in time */
			for (it=0; it<nt; it++)
				{
				/* cshift = e^(-i w it dt) */
				cshift.re = (float)(cos((double)(-w*it*dt)));
				cshift.im = (float)(sin((double)(-w*it*dt)));
				/* Pressure wavefield is real part of complex */
				for (ix=0; ix<nx; ix++)
					p[it][ix][iz] += rmul(q[ix],cshift);
				}
			}
		}

	/* Write data volume to standard output */
	write(1, p, nt*nx*nz*sizeof(float));

	/* All done */
	exit(0);
	}