//
//      This software generates a UPF pseudopotential file to simulate
//      a spherical jellium particle using the Quantum Espresso package.
//      Compilation: 
//              gcc -o jelly-maker.x jelly-maker.c -lm
//
//      Usage:
//              jelly-maker.x <rs> <ne> <filename>
//
//      where rs is the Seitz radius of jellium, ne is the number of
//      electrons, and filename is the output filename.
//
//      Please cite this code as follows:
//      V. I. Artyukhov, M. Liu, E. S. Penev, and B. I. Yakobson, 
//      "A jellium model of a catalyst particle in carbon nanotube growth",
//     `J. Chem. Phys. 146, 244701 (2017)
//
//      URL: http://dx.doi.org/10.1063/1.4986949
//

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

#define logrmin -7      // log of lower grid limit
#define logrmax 4.6     // log of upper grid limit
#define npoints 1451    // number of radial grid points
#define rabfac 125      // ratio of R to RAB grids
#define rhostep 0.1     // electron density decay width, unimportant
#define Bohr 0.529177   // must be god to tweak!

double  sie(double, double);                    // self-interaction energy
int     writeheader(double, double, FILE *);
int     writemesh(double *, FILE *);
double  gpot(double, double, double);           // compute the potential
int     writerho(double, double, double *, FILE *);
int     writepot(double, double, double *, FILE *);

int main(int argc, char *argv[])

//      check if the arguments are enough
{

if (argc != 4) {
        printf("Please provide the Seitz radius, number of electons, and output filename\n");
        return 1;
}

//      variable declaration
double  rs, ne, energy, radius;
double  mesh[npoints];
int     status, i;
char    *name;
FILE    *of;

//      set variables
rs = atof(argv[1]);
ne = atof(argv[2]);

//      compute radius, energy, and prepare the grid
radius = rs * pow(ne, 1./3.);
energy = sie(rs, ne);
for(i=1; i<=npoints; i++) {
        mesh[i] = exp( logrmin + ( i - 1 ) * ( logrmax - logrmin ) / ( npoints - 1 ) );
}

//      output radius & energy
printf("%f Ha | %f Angstrom\n", energy, radius*Bohr);

//      write the output file
of = fopen(argv[3], "w");
status = writeheader(rs, ne, of);
status = writemesh(mesh, of);
status = writerho(rs, ne, mesh, of);
status = writepot(rs, ne, mesh, of);
fclose(of);

//      done
return 0;

}



double sie(double rs, double ne) {
        // self-interaction energy of jellium globe
        return ( 3. / 5.) * ne * ne / rs / pow(ne, 1./3.);
}



int writeheader(double rs, double ne, FILE *of) {
        // write the UPF file header
        fprintf(of, "<PP_INFO>\n\
        Spherical jellium.\n\
        Seitz radius (a.u.):            %f\n\
        Number of electrons:            %f\n\
        Self-interaction energy (Ha):   %f\n\
        Radius (Angstrom):              %f\n\
</PP_INFO>\n\
\n\
<PP_HEADER>\n\
0                       Version Number\n\
H                       Element\n\
NC                      Norm - conserving\n\
F                       No Nonlinear Core Correction\n\
SLA PZ NOGX NOGC        LDA Perdew - Zunger\n\
%f              Z valence\n\
0                       Total energy\n\
0 0                     Suggested cutoff for wfc and rho\n\
0                       Max angular momentum component\n\
%d                      Number of points in mesh\n\
0 0                     Number of Wavefunctions, Number of Projectors\n\
</PP_HEADER>\n\
", rs, ne, sie(rs, ne), rs * pow(ne, 1./3.) * Bohr, ne, npoints);

        return 0;

}



int writemesh (double *mesh,  FILE *of) {
        // write the R and RAB grids
        int i;
        fprintf(of, "\n<PP_MESH>\n  <PP_R>\n");
        for (i=1; i<=npoints; i++) {
                fprintf(of, "  %.11E", mesh[i]);
                if(i % 4 == 0) {
                        fprintf(of, "\n");
                }
        }
        if(npoints % 4 != 0) {
                                fprintf(of, "\n");
                }
        fprintf(of, "  </PP_R>\n  <PP_RAB>\n");
        for (i=1; i<=npoints; i++) {
                fprintf(of, "  %.11E", mesh[i] / rabfac);
                if(i % 4 == 0) {
                        fprintf(of, "\n");
                }
        }
        if(npoints % 4 != 0) {
                                fprintf(of, "\n");
                }

        fprintf(of, "  </PP_RAB>\n</PP_MESH>\n");
        return 0;
}



int writerho(double rs, double ne, double *mesh, FILE *of) {
        // write the pseudoatomic density
        int i;
        double r0, rho; // real rho times 4pi*x^2
        r0 = rs * pow(ne, 1./3.);
        fprintf(of, "\n<PP_RHOATOM>\n");
        for (i=1; i<=npoints; i++) {
                rho = mesh[i] * mesh[i] * 3 / rs / rs / rs
                         / ( exp( ( mesh[i] - r0 ) / rhostep ) + 1 );
                fprintf(of, "  %.11E", rho);
                if(i % 4 == 0) {
                        fprintf(of, "\n");
                }
        }
        if(npoints % 4 != 0) {
                fprintf(of, "\n");
        }
        fprintf(of,"</PP_RHOATOM>\n");
        return 0;
}



double gpot(double rs, double ne, double r) {
        // computes the electrostatic potential of jellium globe
        double r0;
        r0 = rs * pow(ne, 1./3.);
        if (r >= r0) {
                return -2 * ne / r;
        }
        else {
                return ( -2 * ne / ( 2 * r0 ) ) * ( 3 - r * r / r0 / r0 );
        }
        return 0;
}



int writepot (double rs, double ne, double *mesh, FILE *of) {
        // write the pseudopotential
        int i;
        fprintf(of, "\n<PP_LOCAL>\n");
        for (i=1; i<=npoints; i++) {
                fprintf(of, "  %.11E", gpot( rs, ne, mesh[i]));
                if(i % 4 == 0) {
                        fprintf(of, "\n");
                }
        }
        if(npoints % 4 != 0) {
                                fprintf(of, "\n");
                }
        fprintf(of, "</PP_LOCAL>\n");
        return 0;
}