/*
*			BlackScholesImp.c:
*
*		Sample code for AMS-691 homework.
*
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define         sqr(x)          ((x)*(x))

static void expiry_func(double,double*,double*,int);

/* functions for different schemes */
static void TridiagonalSolver(const int,const double*,const double*,
			      const double*, const double*,double*);

main(int argc, char **argv)
{ 
	double *x,*u_old;
	double time,dx,dt,dt_cfl;
	int i,n,step;
	char xg_name[200];
	FILE *xg_file;
	double E = 1.0;
	double gamma = 0.08;
	double sigma = 0.2;
	double CFL = 0.7;
	double *a,*b,*c,*d,*w;
	double print_time;
	int is_print_time;
        char fname[100];

	/* computational domain and mesh parameters */
	double L = 0.0;
	double U = 10.0;
	int  mesh_size = 401;

	x = (double*)malloc(mesh_size*sizeof(double));
        u_old = (double*)malloc(mesh_size*sizeof(double));
	a = (double*)malloc((mesh_size-2)*sizeof(double));
	b = (double*)malloc((mesh_size-2)*sizeof(double));
	c = (double*)malloc((mesh_size-2)*sizeof(double));
	d = (double*)malloc((mesh_size-2)*sizeof(double));
	w = (double*)malloc((mesh_size-2)*sizeof(double));
	n = mesh_size - 2;

	/* Initialization of states */

	dx = (U - L)/(mesh_size - 1);

	dt_cfl = dx/gamma/10.0;
	dt_cfl *= CFL;

	for (i = 0; i < mesh_size; i++)
	    x[i] = L + i*dx;

	/* Set the initial condition */

	expiry_func(E,x,u_old,mesh_size);

	/* Time loop */

	step = 0;   time = 0.0;    

        sprintf(fname,"bs_data-%3.1f",time);
        xg_file = fopen(fname,"w");
        fprintf(xg_file,"\"Option price vs. stock value\"\n");
        for (i = 0; i < mesh_size/4; i++)
            fprintf(xg_file,"%f  %f\n",x[i],u_old[i]);
        fclose(xg_file);

	print_time = 1.0; 
	while (time <= 4.0)
	{
	    dt = dt_cfl;
            is_print_time = 0;
            if (time < print_time && time+dt >= print_time)
            {
                dt = print_time - time;
                print_time += 1.0;
                is_print_time = 1;
            }
	
	    /* Setting coefficient matrix and RHS */
	    for (i = 1; i < mesh_size-1; ++i)
	    {
		a[i-1] =  dt/2.0/dx*gamma*x[i] - dt/sqr(dx)/2.0*sqr(sigma*x[i]);
		c[i-1] = -dt/2.0/dx*gamma*x[i] - dt/sqr(dx)/2.0*sqr(sigma*x[i]);
		b[i-1] = 1.0 + dt*sqr(sigma*x[i]/dx) + dt*gamma;
		d[i-1] = u_old[i];
	    }
	    a[0] = c[n-1] = 0.0;
	    d[0] += -(dt/dx*gamma*x[1]/2.0 - 
		dt/2.0*sqr(sigma*x[1]/dx))*u_old[0]; 
	    d[n-1] += -(-dt/dx*gamma*x[mesh_size-2]/2.0 - 
		dt/2.0*sqr(sigma*x[mesh_size-2]/dx))*u_old[mesh_size-1]; 

	    TridiagonalSolver(n,a,b,c,d,w);

	    step++;
	    time += dt;

	    for (i = 1; i < mesh_size-1; i++)
		u_old[i] = w[i-1];

	    /* Set boundary condition */
	    u_old[0] = 0.0;
	    u_old[mesh_size-1] = u_old[mesh_size-2] + dx;

	    /* Output date */
            if (is_print_time == 1)
            {
                sprintf(fname,"bs_data-%3.1f",time);
                xg_file = fopen(fname,"w");
                fprintf(xg_file,"\"Option price vs. stock value\"\n");
                for (i = 0; i < mesh_size/4; i++)
                    fprintf(xg_file,"%f  %f\n",x[i],u_old[i]);
                fclose(xg_file);
            }
	}
	printf("step = %d\n",step);
}

static void expiry_func(
	double E,
	double *x,
	double *u,
	int mesh_size)
{
	int i;
	double arg;

	for (i = 0; i < mesh_size; ++i)
	{
	    arg = x[i];
	    if (arg > E) u[i] = arg - E;
	    else u[i] = 0.0;
	}
}

static void TridiagonalSolver(
	const int n,
	const double *a,
	const double *b,
	const double *c,
	const double *d,
	double *x)
{
	int i;
	static double *A,*B;

	if (A == NULL)
	{
	    A = (double*)malloc(n*sizeof(double));
	    B = (double*)malloc(n*sizeof(double));
	}

	A[n-2] = -a[n-1]/b[n-1];
	B[n-2] = d[n-1]/b[n-1];

	for (i=n-3; i != -1; --i)
	{
		A[i] = -a[i+1]/(b[i+1]+c[i+1]*A[i+1]);
		B[i] = (d[i+1] - c[i+1]*B[i+1])/(b[i+1]+c[i+1]*A[i+1]);
	}

	x[0] = (d[0]-c[0]*B[0])/(b[0]+c[0]*A[0]);
	for (i=1; i < n; ++i)
		x[i] = A[i-1]*x[i-1]+B[i-1];
}