/* hydrogen.c: graphs electron orbitals.
version 1.1
 arguments:
	-res=<pixels per division>
	-dx -dy -dz=<divisions in graph>
	-sqr  <show square of psi.  If this isn't selected,
					graph psi with positive values in red and negatives in blue.>
	-sat=<saturation 1.0 nominal.>
 -s -m -lm= <state variables>
*/

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

#define HEADSIZE 54

unsigned char bmphead[]={
	'B','M',
	0,0,0,0, 		//offset 2 -> file size
	0,0,0,0, 		//reserved
	HEADSIZE,0,0,0, 	//offset 10 -> header size (including any palette data...)
	0x28,0,0,0, 	//offset 14 -> subheader size (not including palette...)
	0,0,0,0, 		//offset 18 -> width
	0,0,0,0, 		//offset 22 -> height
	1,0,				//offset 26 -> planes (usually 1)
	24,0,				//offset 28 -> bits
	0,0,0,0,		//offset 30 -> compression (usually 0)
	0,0,0,0,		//offset 34 -> size of image
	0,0,0,0,		//offset 38 -> X pixels per meter (or 0) 
	0,0,0,0,		//offset 42 -> Y pixels per meter (or 0)
	0,0,0,0,		//offset 46 -> Colors used (or 0 for nonpalette)
	0,0,0,0			//offset 50 -> number of important colors (or 0 for nonpalette)
}; //Total size 54

int WriteBMPHead(FILE *fp,int dx,int dy,int bits)
{
long t,bytes;

	if (dx<1)
		return 0;
	if (dy<1)
		return 0;

	if (bits!=24)
		return 0;
	bytes=3;

	t=(long)dx*(long)dy*bytes;
	/* file size */
	bmphead[2]=(t+HEADSIZE)&0xFF;
	bmphead[3]=((t+HEADSIZE)>>8)&0xFF;
	bmphead[4]=((t+HEADSIZE)>>16)&0xFF;
	bmphead[5]=((t+HEADSIZE)>>24)&0xFF;

	/* width */
	bmphead[18]=dx&0xFF;
	bmphead[19]=(dx>>8)&0xFF;
	bmphead[20]=(dx>>16)&0xFF;
	bmphead[21]=(dx>>24)&0xFF;

	/* height */
	bmphead[22]=dy&0xFF;
	bmphead[23]=(dy>>8)&0xFF;
	bmphead[24]=(dy>>16)&0xFF;
	bmphead[25]=(dy>>24)&0xFF;

	/* image size */
	0,0,0,0,		//offset 34 -> size of image
	bmphead[34]=t&0xFF;
	bmphead[35]=(t>>8)&0xFF;
	bmphead[36]=(t>>16)&0xFF;
	bmphead[37]=(t>>24)&0xFF;

	if (fwrite(bmphead,1,sizeof(bmphead),fp)!=sizeof(bmphead)) {
		puts("bad write");
		return 0;
	}
	return 1;
}

typedef double REAL;

long dx=3;
long dy=3;
long dz=3;
long res=4;
REAL sat=1.0;
int s=1;
int l=0;
int ml=0;
int aml;
int sqrd=0;


void writeheader(FILE *f)
{
	WriteBMPHead(f,2*dx*res,2*dy*res,24);
}

void calc(long x,long y,long z,int *cols);
void axes(long x,long y,long z,int *cols);
void setup(void);

FILE *parseargs(char *argv[])
{
int i;
char *arg;
FILE *fp=NULL;

    for (i=1;argv[i]!=NULL;i++) {
        arg=argv[i];
        if (arg[0]!='-') {
            if (fp==NULL) {
                fp=fopen(arg,"wb");
                if (fp==NULL) {
                    printf("couldn't create file %s\n",arg);
                    return NULL;
                }
            }
            else {
                fclose(fp);
                printf("more than one file specified: %s\n",arg);
                return NULL;
            }
        }

        if (strncmp("-dx=",arg,4)==0) {
            dx=atoi(arg+4);
            if (dx<1)
                dx=1;
        }
        if (strncmp("-dy=",arg,4)==0) {
            dy=atoi(arg+4);
            if (dy<1)
                dy=1;
        }
        if (strncmp("-dz=",arg,4)==0) {
            dz=atoi(arg+4);
            if (dz<1)
                dz=1;
        }
        if (strncmp("-dx=",arg,4)==0) {
            dx=atoi(arg+4);
            if (dx<1)
                dx=1;
        }
        if (strncmp("-res=",arg,5)==0) {
            res=atoi(arg+5);
            if (res<1)
                res=1;
        }
        if (strncmp("-sat=",arg,5)==0) {
            sat=atof(arg+5);
            if (sat<0.0)
                sat=0.0;
            if (sat>10.0)
                sat=10.0;
        }
        if (strncmp("-s=",arg,3)==0) {
            s=atoi(arg+3);
            if (s<1)
                s=1;
            if (s>10)
                s=10;
        }
        if (strncmp("-l=",arg,3)==0) {
            l=atoi(arg+3);
            if (l<0)
                l=0;
            if (l>=s)
                l=s-1;
        }
        if (strncmp("-ml=",arg,4)==0) {
            ml=atoi(arg+4);
            if (ml<-l)
                ml=-l;
            if (ml>l)
                ml=l;
        }
        if (strncmp("-sqr",arg,4)==0) {
            sqrd=1;
        }

    }
    return fp;
}

void main(int argc,char *argv[])
    {
    long x,y,z;
    int cols[3];
    FILE *fp;

    fp=parseargs(argv);
    if (!fp) {
        puts("incorrect output file specified");
        return;
    }

    printf("graphing s=%d l=%d ml=%d\n",s,l,ml);

    writeheader(fp);

    setup();

    for (y=-dy*res;y<dy*res;y++) {
        for (x=-dx*res;x<dx*res;x++) {
            cols[0]=255;
            cols[1]=255;
            cols[2]=255;
            for (z=-dz*res;z<dz*res;z+=2)  {
                calc(x+(z>>1),y+(z>>1),z,cols);
                axes(x+(z>>1),y+(z>>1),z,cols);
            }
            putc(cols[2],fp);
            putc(cols[1],fp);
            putc(cols[0],fp);
        }
        //printf("y=%f\n",(float)y/(float)res);
    }
    fclose(fp);
}


REAL thetap[11];
REAL rp[10];

void setup(void)
{
int i;
REAL t,n,b,max;

    if (ml<0)
        aml=-ml;
    else
        aml=ml;

    if (l&1)
        thetap[1]=1.0;
    else
        thetap[0]=1.0;
    rp[0]=1.0;
    for(i=0;i<9;i++) {
        thetap[i+2]=thetap[i]*(i*(i+1)-l*(l+1))/((i+2)*(i+1));
        thetap[i+2]*=(i+1)*(i+2);
        if (i>=aml)
            thetap[i+2]=thetap[i+2]/((i+1-aml)*(i+2-aml));

        rp[i+1]=rp[i]*(i+l+1-s)/((i+l+1)*(i+l+2)-(l+1)*l);
    }

    max=0.0;
    for (n=-1.0;n<1.0;n+=0.01) {
        t=0.0;
        b=1.0;
        for (i=aml;i<=l;i++) {
            t+=thetap[i]*b;
            b*=n;
        }
        if (t>max)
            max=t;
        if (-t>max)
            max=-t;
    }
    printf("max of thetap = %f\n",max);
    for (i=0;i<11;i++)
        thetap[i]/=max;

    max=0.0;
    for (n=0;n<dx;n+=1.0/res) {
        b=1.0;
        t=0.0;
        for (i=0;i<l;i++)
            b*=n;
        for (i=0;i<s-l;i++) {
            t+=rp[i]*b;
            b*=n;
        }
        t*=exp(-n/s);
        //printf("r=%f psi(r)=%f\n",n,t);
        if (t>max)
            max=t;
        if (-t>max)
            max=-t;

    }
    printf("max of rp = %f\n",max);
    for (i=0;i<10;i++)
        rp[i]/=max;
}

void calc(long x,long y,long z,int *cols)
{
int v;
REAL r,theta,phi,psi,f1,f2,a,b;
int i;

    r=sqrt(x*x+y*y+z*z);
    theta=acos(y/r);
    phi=atan2(x,z);
    r=r/res;


    a=cos(theta);
    b=1.0;
    f2=0.0;
    for (i=aml;i<=l;i++) {
        f2+=thetap[i]*b;
        b*=a;
    }


    a=sin(theta);
    for (i=0;i<aml;i++)
        f2*=a;

    if (ml<0) {
          f1=sin(aml*phi);
    }
    else if (ml>0) {
        f1=cos(aml*phi);
    }
    else
        f1=1.0;

    psi=0.0;

    a=1.0;
    for (i=0;i<l;i++)
        a*=r;
    for (i=0;i<s-l;i++) {
        psi+=rp[i]*a;
        a*=r;
    }
    psi*=exp(-r/s)*f1*f2;
    if (sqrd)
        psi=psi*psi;

    v=(int)(sat*psi*256);

    if (v>0) {
        if (v>256)
            v=256;
        cols[0]=(int)(((cols[0]*(256L-v))>>8)+v);
        if (cols[0]>255)
            cols[0]=255;
        cols[1]=(int)((cols[1]*(256L-v))>>8);
        cols[2]=(int)((cols[2]*(256L-v))>>8);
    }
    else {
        if (v<-256)
            v=-256;
        cols[0]=(int)((cols[0]*(256L+v))>>8);
        cols[1]=(int)((cols[1]*(256L+v))>>8);
        cols[2]=(int)(((cols[2]*(256L+v))>>8)-v);
        if (cols[2]>255)
            cols[2]=255;

    }
}

void axes(long x,long y,long z,int *cols)
{
    if ((x!=0)&&(y!=0))
        return;
    if ((x!=0)&&(z!=0))
        return;
    if ((y!=0)&&(z!=0))
        return;
    if ((x+y+z)%res==0) {
        cols[0]=0;
        cols[1]=0;
        cols[2]=0;
    }
    else {
        cols[0]=128;
        cols[1]=128;
        cols[2]=128;
    }
}