/*
    KdV JEDNACINA
    implicitna - spektralna metoda

    http://bojan.info/diplomski/num/kdv.c

    Napomene:
    - Poziva gnuplot, koji treba biti instaliran na sistemu
        (po potrebi promijeniti "wgnuplot" u "gnuplot"
        na kraju funkcije cgplot.)
    - Izbor gnuplot terminala (GIF ili LateX)
        se podesava sa #define GIF
    - Za izlaz napraviti folder 'output/gif'
        odn. 'output/latex'.
    - Pocetni uvjet zadati sa #define START(x)
*/

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

/*
    PRETPROCESORSKE VARIJABLE
*/
#define PI      (4 * atan(1))

/* prostorna resetka */
#define N       1024 /* broj prostornih celija */
#define L       50.0 /* duzina prostornog intervala */
#define X(j)    (j*L/N) /* j-ti cvor */
#define A       (L/(2*PI))

/* vremenska diskretizacija */
#define DT      .005 /* vremenski korak */
#define T       11.0 /* max. vrijeme */
#define GPSTEP  99 /* plot svakih GPSTEP+1 koraka */

/* koeficijenti */
#define DISP    1.0 /* koeficijent disperzije */

/* pocetni uvjet */
#define START(x) 12/A*(.64/(cosh(.8*(x-5))*cosh(.8*(x-5)))\
 + .25/(cosh(.5*(x-5))*cosh(.5*(x-5))))
/* 2 solitona: 12/A*(.64/(cosh(.8*(x-5))*cosh(.8*(x-5)))\
 + .25/(cosh(.5*(x-15))*cosh(.5*(x-15)))) */
/* gausijan: ( exp(-2*(x-.25*L)*(x-.25*L)) / A ) */

/* gnuplot */
#define GIF     0 /* 0/1 = LaTeX/gif terminal */
#define XMIN    0.0*L /* min. apscise */
#define XMAX    1.0*L /* max. apscise */
#define YMIN    -.1 /* min. ordinate */
#define YMAX    8.0 /* max. ordinate */

/* NR */
#define NR_END  1
#define FREE_ARG    char*
#define SWAP(a,b)   tempr=(a);(a)=(b);(b)=tempr

/*
    PROTOTIPOVI FUNKCIJA
*/
void init(__complex__ float u[]);
void cgplot(__complex__ float u[], int n);

/* funkcije adaptirane iz NR */
__complex__ float *cfvector(long nl, long nh);
void free_cfvector(__complex__ float *v, long nl, long nh);
void cfour1(__complex__ float data[], unsigned long nn, int isign);

/* originalne NR funkcije */
void nrerror(char error_text[]);
float *vector(long nl, long nh);
int *ivector(long nl, long nh);
void free_vector(float *v, long nl, long nh);
void free_ivector(int *v, long nl, long nh);


int main()
{
    __complex__ float *u; /* valna funkcija */
    u = cfvector(0,N-1);
    init(u); /* zadavanje pocetnih uvjeta */

    __complex__ float *v, *w, *temp; /* pomocni vektori */
    v = cfvector(0,N-1);
    w = cfvector(0,N-1);
    temp = cfvector(0,N-1);

    __complex__ float *uop, *bop; /* dif. oper. u k-prostoru */
    uop = cfvector(0,N-1); /* operator U */
    bop = cfvector(0,N-1); /* operator B */

    int *kk; /* vektor sa valnim brojevima */
    kk = ivector(0,N-1);

    int j; /* brojac za koracanje kroz x-prostor */
    int k; /* brojac za koracanje kroz k-prostor */
    int i; /* brojac za iteraciju rjesenja implicitne jedn.*/

    float m = DISP * DT / pow(A,3); /* pomocni faktor */
    __complex__ naz; /* nazivnik za uop i bop */
    for(k=0; k<N; k++) {
        /* inicijalizacija kk */
        if(k<=N/2)
            kk[k] = k;
        else
            kk[k] = k - N;
        /* inicijalizacija uop i bop */
        naz = 1.0 + 0.5i * m * pow(kk[k],3);
        uop[k] = (1.0 - 0.5i * m * pow(kk[k],3)) / naz;
        bop[k] = (0.25i * DT * kk[k]) / naz;
    }

    int ti = 0; /* brojac vremenskih koraka */
    int tm = (int) T/DT; /* broj vremenskih koraka */
    int gpi = 0; /* brojac za iscrtavanje */
    for(ti=0; ti<tm; ti++) { /* koraci kroz vrijeme */
        printf("procesira se korak %d\n", ti);

        if(!gpi--) { /* ako je brojac gpi = 0 */
            printf("iscrtava se korak %d\n", ti);
            for(j=0; j<N; j++)
                temp[j] = A * u[j];
            cgplot(temp,ti); /* pozovi plot funkciju */
            gpi = GPSTEP; /* resetuj brojac */
        }

        for(j=0; j<N; j++)
            temp[j] = u[j]*u[j]; /* temp = u^2 */
        cfour1(u,N,1); /* u = FFT(u) */
        cfour1(temp,N,1); /* temp = FFT(u^2) */

        for(k=0; k<N; k++){
            /* v = v(k), prethodni vremenski korak */
            v[k] = uop[k]*u[k] + bop[k]*temp[k];
            /* w = w(k) = u_1, prvi korak iteracije */
            w[k] = uop[k]*u[k] + 2*bop[k]*temp[k];
        }
        cfour1(v,N,-1); /* v = v(x) = F^-1(v(k)) */
        cfour1(w,N,-1); /* w = w(x) = F^-1(w(k)) */
        for(j=0; j<N; j++) {
            v[j] = __real__ v[j]; /* v = Re(v) */
            w[j] = __real__ w[j]; /* w = Re(w) */
        }

        /* Iterativno rjesenje implicitne jednacine */
        for(i=2; i<5; i++) { /* (w: i=1); ide se do i=4 */
            for(j=0; j<N; j++)
                temp[j] = w[j]*w[j]; /* temp = w^2 */
            cfour1(temp,N,1); /* temp = FFT(w^2) */
            for(k=0; k<N; k++)
                temp[k] *= bop[k]; /* temp(k) = bop*FFT(w^2) */
            cfour1(temp,N,-1); /* temp(x) */
            for(j=0; j<N; j++)
                w[j] = v[j] + __real__ temp[j]; /* w = u_i */
        }

        for(j=0; j<N; j++)
            u[j] = w[j]; /* u(t+dt) = w = u_4 */
    }

    free_cfvector(uop,0,N-1);
    free_cfvector(bop,0,N-1);
    free_cfvector(u,0,N-1);
    free_cfvector(v,0,N-1);
    free_cfvector(w,0,N-1);
    free_cfvector(temp,0,N-1);

    return 0;
}

void init(__complex__ float u[])
/* inicijalizacija - pocetni uslov */
{
    int j;
    float xj;
    for(j=0; j<N; j++) {
        xj = X(j); /* j-ti cvor */
        u[j] = START(xj); /* pocetna vrijednost */
    }
}

void cgplot(__complex__ float y[], int n)
/*
    Iscrtava realni dio kompleksnog vektora y[0,1,...,N-1].
    n je samo indeks koraka za oznacavanje grafikona.
*/
{
    /* Ispis vrijednosti u data.dat */
    FILE *data; /* datoteka sa vrijednostima za iscrtavanje */
    data = fopen("data.dat", "w");
    int j;
    for(j=0; j<N; j++) /* zapisivanje vrijednosti u datoteku */
        fprintf(data, "%f\t%f\n", X(j), __real__ y[j]);
    fclose(data);

    /* Ispis gnuplot komandi u cmd.gp */
    FILE *gp; /* datoteka sa gnuplot komandama */
    gp = fopen("cmd.gp", "w");
    if(GIF) {
        fprintf(gp, "set terminal gif\n");
        fprintf(gp, "set output 'output\\gif\\%5.5d.gif'\n", n);
    } else {
        fprintf(gp, "set terminal latex\n");
        fprintf(gp, "set output 'output\\latex\\%5.5d.tex'\n", n);
    }
    fprintf(gp, "plot [%.0f:%.0f] ", XMIN, XMAX);
    fprintf(gp, "[%.1f:%.1f] 'data.dat' ", YMIN, YMAX);
    fprintf(gp, "title \"u(x, t=%.1f)\" ", n*DT);
    fprintf(gp, "smooth cspline\n");
    fclose(gp);

    /* Poziv gnuplota */
    system("wgnuplot cmd.gp");
}

/*
    FUNKCIJE ADAPTIRANE IZ NR
*/

__complex__ float *cfvector(long nl, long nh)
/* preradjena NR funkcija vector za __complex__ float */
{
    __complex__ float *v;

    v=(__complex__ float *)malloc((size_t)
     ((nh-nl+1+NR_END)*sizeof(__complex__ float)));
    if (!v) nrerror("allocation failure in vector()");
    return v-nl+NR_END;
}

void free_cfvector(__complex__ float *v, long nl, long nh)
/* preradjena NR funkcija free_vector za __complex__ float */
{
    free((FREE_ARG) (v+nl-NR_END));
}

void cfour1(__complex__ float cdata[], unsigned long nn, int isign)
/*
    Preradjena NR funkcija four1.
    Interna implementacija nije mijenjana, samo ulaz/izlaz.
    Ulaz: umjesto float data[1..2*nn],
        prima __complex__ float cdata[0..nn-1].
    Izlaz: za inverznu transformaciju,
        vraca normiran transformat (pomnozen faktorom 1/nn).
*/
{
    /* Procedure za konverziju ulaza/izlaza */
    void c2r(__complex__ float c[], float r[])
    /* pretvara kompleksni vektor c[0..nn-1]
        u realni vektor r[1..2*nn] */
    {
        int i;
        for(i=0; i<nn; i++) {
            r[2*i+1] = __real__ c[i]; /* (1,3,...) su realni */
            r[2*i+2] = __imag__ c[i]; /* (2,4,...) su imag. */
        }
    }
    void r2c(float r[], __complex__ float c[])
    /* pretvara realni vektor r[1..2*nn]
        u kompleksni vektor c[0..nn-1] */
    {
        int i;
        for(i=0; i<nn; i++)
            c[i] = r[2*i+1] + 1.0i * r[2*i+2];
    }

    /* Prerada ulaza */
    float *data;
    data = vector(1,2*nn);
    c2r(cdata,data);

    /* Originalni dio (four1) */
    unsigned long n,mmax,m,j,istep,i;
    double wtemp,wr,wpr,wpi,wi,theta;
    float tempr,tempi;

    n=nn << 1;
    j=1;
    for (i=1;i<n;i+=2) {
        if (j > i) {
            SWAP(data[j],data[i]);
            SWAP(data[j+1],data[i+1]);
        }
        m=nn;
        while (m >= 2 && j > m) {
            j -= m;
            m >>= 1;
        }
        j += m;
    }

    mmax=2;
    while (n > mmax) {
        istep=mmax << 1;
        theta=isign*(6.28318530717959/mmax);
        wtemp=sin(0.5*theta);
        wpr = -2.0*wtemp*wtemp;
        wpi=sin(theta);
        wr=1.0;
        wi=0.0;
        for (m=1;m<mmax;m+=2) {
            for (i=m;i<=n;i+=istep) {
                j=i+mmax;
                tempr=wr*data[j]-wi*data[j+1];
                tempi=wr*data[j+1]+wi*data[j];
                data[j]=data[i]-tempr;
                data[j+1]=data[i+1]-tempi;
                data[i] += tempr;
                data[i+1] += tempi;
            }
            wr=(wtemp=wr)*wpr-wi*wpi+wr;
            wi=wi*wpr+wtemp*wpi+wi;
        }
        mmax=istep;
    }
    /* Kraj originalnog dijela (four1) */

    /* Dio za preradu izlaza */
    r2c(data,cdata);
    free_vector(data,1,2*nn);
    /* Za inverznu transformaciju dodaje se faktor 1/nn */
    if(isign==-1){
        for(j=0; j<nn; j++) {
            cdata[j] *= 1.0/nn;
        }
    }
}


/*
    ORIGINALNE NR FUNKCIJE
*/

void nrerror(char error_text[])
/* Numerical Recipes standard error handler */
{
    fprintf(stderr,"Numerical Recipes run-time error...\n");
    fprintf(stderr,"%s\n",error_text);
    fprintf(stderr,"...now exiting to system...\n");
    exit(1);
}

float *vector(long nl, long nh)
/* allocate a float vector
    with subscript range v[nl..nh] */
{
    float *v;

    v=(float *)malloc((size_t)
     ((nh-nl+1+NR_END)*sizeof(float)));
    if (!v) nrerror("allocation failure in vector()");
    return v-nl+NR_END;
}

int *ivector(long nl, long nh)
/* allocate an int vector with subscript range v[nl..nh] */
{
    int *v;

    v=(int *)malloc((size_t) ((nh-nl+1+NR_END)*sizeof(int)));
    if (!v) nrerror("allocation failure in ivector()");
    return v-nl+NR_END;
}

void free_vector(float *v, long nl, long nh)
/* free a float vector allocated with vector() */
{
    free((FREE_ARG) (v+nl-NR_END));
}

void free_ivector(int *v, long nl, long nh)
/* free an int vector allocated with ivector() */
{
    free((FREE_ARG) (v+nl-NR_END));
}