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

/* ---------------- PARAMETERS ---------------- */
#define N 10000
#define RMAX 100.0
#define M 1.0
#define CFL 0.25
#define TMAX 20.0

#define OUTPUT_EVERY 10

#define SLICING_GEODESIC 0


/* ---------------- VARIABLES ---------------- */
enum {
    ALPHA, 
    A_BAR, B_BAR,
    D_A_BAR, D_B_BAR,
    K_A, K_B,
    NVARS
};

typedef struct {
    double v[NVARS][N];
} State;

/* ---------------- GRID ---------------- */
double r[N];
double dr;
double dt;

/* background fields */
double psi[N];
double E_arr[N];
double dE_arr[N];

/* horizon tracking */
double prev_theta[N];

/* symmetry */
double origin_sign[NVARS] = {+1,+1,+1,-1,-1,+1,+1};

/* asymptotic */
double u_inf[NVARS] = {1,1,1,0,0,0,0};

/* ------------------------------------------------ */

static inline double deriv(double *u,int i)
{
    return (u[i+1]-u[i-1])/(2.0*dr);
}

static inline double deriv_log(double *u,int i)
{
    double Dl = deriv(u,i-1)/u[i-1];
    double Dr = deriv(u,i+1)/u[i+1];
    return (Dr-Dl)/(2.0*dr);
}

/* ------------------------------------------------ */

void setup_grid()
{
    dr = RMAX/N;

    for(int i=0;i<N;i++)
    {
        r[i]     = (i-0.5)*dr;
        psi[i]   = 1.0 + M/(2.0*r[i]);
        E_arr[i] = -4.0*M/(2.0*r[i]*r[i] + M*r[i]);
        dE_arr[i] = 4.0*M*(M+4*r[i]) / (r[i]*r[i]*pow(M+2*r[i],2));

        prev_theta[i] = 1e100;
    }

    dt = CFL*dr;
}

/* ------------------------------------------------ */

static inline double dereg(double u,int i)
{
    return u * pow(psi[i],4);
}

static inline double deregD(double u,int i)
{
    return u + E_arr[i];
}

/* ------------------------------------------------ */

void enforce_boundaries(State *U)
{
    for(int v=0;v<NVARS;v++)
    {
        U->v[v][0]   = origin_sign[v]*U->v[v][1];
        U->v[v][N-1] = (U->v[v][N-2]*r[N-2] + u_inf[v]*dr) / r[N-1];
    }
}

/* ------------------------------------------------ */

void rhs(State *U, State *dU)
{
    enforce_boundaries(U);

    for(int i=1;i<N-1;i++)
    {
        double alpha = U->v[ALPHA][i];

#if (SLICING_GEODESIC)
        alpha = 1.0;
#endif

        double Abar    = U->v[A_BAR][i];
        double Bbar    = U->v[B_BAR][i];

        double KA      = U->v[K_A][i];
        double KB      = U->v[K_B][i];

        double A       = dereg(Abar,i);
        double B       = dereg(Bbar,i);
        double DA      = deregD(U->v[D_A_BAR][i],i);
        double DB      = deregD(U->v[D_B_BAR][i],i);

        double Dalpha  = deriv(U->v[ALPHA],i) / U->v[ALPHA][i];
        double dDalpha = deriv_log(U->v[ALPHA],i);

        double Ktrace  = KA + 2.0*KB;

#if (SLICING_GEODESIC)
        dU->v[ALPHA][i] = 0.0;
#else
        dU->v[ALPHA][i] = - 2.0*alpha*Ktrace;
#endif

        dU->v[A_BAR][i] = -2.0*alpha*Abar*KA;
        dU->v[B_BAR][i] = -2.0*alpha*Bbar*KB;

        dU->v[D_A_BAR][i] = -2.0*alpha*(KA*Dalpha + deriv(U->v[K_A],i));
        dU->v[D_B_BAR][i] = -2.0*alpha*(KB*Dalpha + deriv(U->v[K_B],i));

        double termA = deriv(U->v[D_B_BAR],i) + dDalpha + dE_arr[i]+ Dalpha*Dalpha -0.5*Dalpha*DA + 0.5*DB*DB -0.5*DA*DB -(1.0/r[i])*(DA-2.0*DB);
        double termB = deriv(U->v[D_B_BAR],i) + dE_arr[i] + Dalpha*DB + DB*DB -0.5*DA*DB - (1.0/r[i])*(DA-2.0*Dalpha-4.0*DB) - 2.0*(A-B)/(r[i]*r[i]*B);

        dU->v[K_A][i] = -alpha/A * termA + alpha*KA*Ktrace;
        dU->v[K_B][i] = -alpha/(2.0*A) * termB + alpha*KB*Ktrace;
    }
}

/* ------------------------------------------------ */

void rk4(State *U)
{
    static State k1,k2,k3,k4,tmp;

    rhs(U,&k1);
    for(int v=0;v<NVARS;v++)
    for(int i=0;i<N;i++)
        tmp.v[v][i] = U->v[v][i] + 0.5*dt*k1.v[v][i];

    rhs(&tmp,&k2);
    for(int v=0;v<NVARS;v++)
    for(int i=0;i<N;i++)
        tmp.v[v][i] = U->v[v][i] + 0.5*dt*k2.v[v][i];

    rhs(&tmp,&k3);
    for(int v=0;v<NVARS;v++)
    for(int i=0;i<N;i++)
        tmp.v[v][i] = U->v[v][i] +  dt*k3.v[v][i];

    rhs(&tmp,&k4);
    for(int v=0;v<NVARS;v++)
    for(int i=0;i<N;i++) 
        U->v[v][i] += dt*(k1.v[v][i]/6.0 + k2.v[v][i]/3.0 + k3.v[v][i]/3.0 + k4.v[v][i]/6.0
        );
}

/* ------------------------------------------------ */

double compute_Kmax(State *U)
{
    double Kmax = 0;

    for(int i=0;i<N;i++)
    {
        double K = U->v[K_A][i] + 2.0*U->v[K_B][i];
        double a = fabs(K);
        if(a > Kmax)
            Kmax = a;
    }

    return Kmax;
}

/* ------------------------------------------------ */

double find_horizon(State *U, double *area)
{
    double theta_prev = 0.0;
    int first = 1;

    for(int i=2;i<N-2;i++)
    {
        double A = dereg(U->v[A_BAR][i],i);
        double B = dereg(U->v[B_BAR][i],i);

        double dB = deriv(U->v[B_BAR],i)/U->v[B_BAR][i] + E_arr[i];
        double theta = (1.0/sqrt(A))*(2.0/r[i] + dB) - 2.0*U->v[K_B][i];

        if(!first)
        {
            if(theta_prev * theta < 0.0)
            {
                /* root interpolation */
                double w =
                    fabs(theta_prev) /
                    (fabs(theta_prev)+fabs(theta));

                double rH = r[i-1]*(1-w) + r[i]*w;

                /* interpolate B */
                double B1 = dereg(U->v[B_BAR][i-1],i-1);
                double B2 = dereg(U->v[B_BAR][i],i);

                double BH = B1*(1-w) + B2*w;

                /* horizon area */
                double R_areal = rH * sqrt(BH);
                *area = 4.0 * M_PI * R_areal * R_areal;

                return rH;
            }
        }

        theta_prev = theta;
        first = 0;
    }

    *area = -1.0;
    return -1.0;
}

/* ------------------------------------------------ */

void output_all(State *U,double t)
{
    static FILE *fa,*fA,*fB,*fKA,*fKB;

    if(t==0)
    {
        fa  = fopen("alpha.dat","w");
        fA  = fopen("Abar.dat","w");
        fB  = fopen("Bbar.dat","w");
        fKA = fopen("KA.dat","w");
        fKB = fopen("KB.dat","w");
    }

    for(int i=0;i<N;i++)
    {
        fprintf(fa ,"%e %e %e\n",t,r[i],U->v[ALPHA][i]);
        fprintf(fA ,"%e %e %e\n",t,r[i],U->v[A_BAR][i]);
        fprintf(fB ,"%e %e %e\n",t,r[i],U->v[B_BAR][i]);
        fprintf(fKA,"%e %e %e\n",t,r[i],U->v[K_A][i]);
        fprintf(fKB,"%e %e %e\n",t,r[i],U->v[K_B][i]);
    }

    fprintf(fa,"\n");
    fprintf(fA,"\n");
    fprintf(fB,"\n");
    fprintf(fKA,"\n");
    fprintf(fKB,"\n");
}

/* ------------------------------------------------ */

int check_crash(State *U)
{
    for(int v=0;v<NVARS;v++)
    for(int i=0;i<N;i++)
        if(!isfinite(U->v[v][i]))
            return 1;
    return 0;
}

/* ------------------------------------------------ */

int main()
{
    setup_grid();

    State U;

    for(int i=0;i<N;i++)
    {
        for(int v=0;v<NVARS;v++)
            U.v[v][i]=0.0;

        U.v[ALPHA][i] = 1.0;
        U.v[A_BAR][i] = 1.0;
        U.v[B_BAR][i] = 1.0;
    }

    double t=0;
    int step=0;

    while(t<TMAX)
    {
        rk4(&U);

        if(check_crash(&U))
        {
            printf("Crash at t=%f\n",t);
            break;
        }

        double Kmax = compute_Kmax(&U);
        double area;
        double rH = find_horizon(&U, &area);
        double Mirr = sqrt(area/(16.0*M_PI));
        printf("t=%f  Kmax=%e      ",t,Kmax);

        if(rH>0) 
          printf("Horizon: radius = %f Area = %.10f Mirr = %.10f  \n", rH, area, Mirr);

        if(step%OUTPUT_EVERY==0)
            output_all(&U,t);

        t+=dt;
        step++;
    }

    return 0;
}