#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <dirent.h>
#include "hdf5.h"
#include <math.h>
#include <time.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_sf_bessel.h>
#include "mclib_3d.h"
//#include "mclib.h"
#include <omp.h>

#define R_DIM 1260
#define THETA_DIM 280
#define PHI_DIM 280

void read_hydro(char hydro_prefix[200], int frame, double r_inj, double **x, double **y,  double **z, double **szx, double **szy, double **r, double **theta, double **phi,\
  double **velx, double **vely, double **velz, double **dens, double **pres, double **gamma, double **dens_lab, double **temp, int *number,  int ph_inj, double min_r, double max_r, double fps, FILE *fPtr)
{
    FILE *hydroPtr=NULL;
    char hydrofile[200]="", file_num[200]="", full_file[200]=""  ;
    char buf[10]="";
    int i=0, j=0, k=0, elem=0;
    float buffer=0;
    float *dens_unprc=malloc(sizeof(float)*R_DIM*THETA_DIM*PHI_DIM);
    float *vel_r_unprc=malloc(sizeof(float)*R_DIM*THETA_DIM*PHI_DIM);
    float *vel_theta_unprc=malloc(sizeof(float)*R_DIM*THETA_DIM*PHI_DIM);
    float *vel_phi_unprc=malloc(sizeof(float)*R_DIM*THETA_DIM*PHI_DIM);
    float *pres_unprc=malloc(sizeof(float)*R_DIM*THETA_DIM*PHI_DIM);
    double ph_rmin=0, ph_rmax=0;
    double r_in=1e10, r_ref=2e13;
    double *r_edge=malloc(sizeof(double)*(R_DIM+1));
    double *dr=malloc(sizeof(double)*(R_DIM));
    double *r_unprc=malloc(sizeof(double)*R_DIM);
    double *theta_unprc=malloc(sizeof(double)*THETA_DIM);
    double *phi_unprc=malloc(sizeof(double)*PHI_DIM);

    if (ph_inj==0)
    {
        ph_rmin=min_r;
        ph_rmax=max_r;
    }

    //density
    snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"u0", 1,"-" );
    modifyFlashName(file_num, hydrofile, frame, 1);
    snprintf(full_file, sizeof(full_file), "%s%s", file_num, ".data");

    fprintf(fPtr,"Reading Density: %s\n", full_file);
    fflush(fPtr);

    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    fread(dens_unprc, sizeof(float)*R_DIM*THETA_DIM*PHI_DIM,R_DIM*THETA_DIM*PHI_DIM, hydroPtr); //data
    fclose(hydroPtr);

    for (i=0;i<R_DIM*THETA_DIM*PHI_DIM;i++)
    {
        if ((i>98784000-5) || (i<5))
        {
            fprintf(fPtr,"Density %d: %0.7e\n", i, *(dens_unprc+i));
            fflush(fPtr);
        }
    }

    //velocities divided by c
    //v_r
    snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"u0", 2,"-" );
    modifyFlashName(file_num, hydrofile, frame, 1);
    snprintf(full_file, sizeof(full_file), "%s%s", file_num, ".data");

    fprintf(fPtr,"Reading v_r: %s\n", full_file);
    fflush(fPtr);

    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    fread(vel_r_unprc, sizeof(float)*R_DIM*THETA_DIM*PHI_DIM,R_DIM*THETA_DIM*PHI_DIM, hydroPtr);
    fclose(hydroPtr);

    for (i=0;i<5;i++)
    {
        fprintf(fPtr,"V_r %d: %e\n", i, *(vel_r_unprc+i));
        fflush(fPtr);
    }

    //v_theta
    snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"u0", 3,"-" );
    modifyFlashName(file_num, hydrofile, frame, 1);
    snprintf(full_file, sizeof(full_file), "%s%s", file_num, ".data");

    fprintf(fPtr,"Reading v_theta: %s\n", full_file);
    fflush(fPtr);

    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    fread(vel_theta_unprc, sizeof(float)*R_DIM*THETA_DIM*PHI_DIM,R_DIM*THETA_DIM*PHI_DIM, hydroPtr);
    fclose(hydroPtr);

    for (i=0;i<5;i++)
    {
        fprintf(fPtr,"V_theta %d: %e\n", i, *(vel_theta_unprc+i));
        fflush(fPtr);
    }

    //v_phi
    snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"u0", 4,"-" );
    modifyFlashName(file_num, hydrofile, frame, 1);
    snprintf(full_file, sizeof(full_file), "%s%s", file_num, ".data");

    fprintf(fPtr,"Reading v_phi: %s\n", full_file);
    fflush(fPtr);

    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    fread(vel_phi_unprc, sizeof(float)*R_DIM*THETA_DIM*PHI_DIM,R_DIM*THETA_DIM*PHI_DIM, hydroPtr);
    fclose(hydroPtr);

    for (i=0;i<5;i++)
    {
        fprintf(fPtr,"V_phi %d: %e\n", i, *(vel_phi_unprc+i));
        fflush(fPtr);
    }

    //pressure (divided by c^2)
    snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"u0", 8,"-" );
    modifyFlashName(file_num, hydrofile, frame, 1);
    snprintf(full_file, sizeof(full_file), "%s%s", file_num, ".data");

    fprintf(fPtr,"Reading pres: %s\n", full_file);
    fflush(fPtr);

    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    fread(pres_unprc, sizeof(float)*R_DIM*THETA_DIM*PHI_DIM,R_DIM*THETA_DIM*PHI_DIM, hydroPtr);
    fclose(hydroPtr);

    for (i=PHI_DIM-1;i<PHI_DIM;i++)
    {
        for (j=THETA_DIM-1;j<THETA_DIM;j++)
        {
            for (k=R_DIM-5;k<R_DIM;k++)
            {

            fprintf(fPtr,"Pres %d: %e\n", (i*R_DIM*THETA_DIM + j*R_DIM + k  ), *(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )));
            fflush(fPtr);

            }
        }
    }

    // see how many elements there are to test if reading correctly
    /*
    hydroPtr=fopen(full_file, "rb");
    fread(&buffer, sizeof(float), 1,hydroPtr); //random stuff about the file from fortran
    while (1 == fread(&buffer,sizeof(float),1,hydroPtr))
    {
        elem++;
    }
    //fread(pres_unprc, sizeof(double)*HYDRO_DIM,HYDRO_DIM, hydroPtr);
    fclose(hydroPtr);

    fprintf(fPtr,"Elem %d\n", elem);
    */

    //R
    if (frame<=1300)
    {
        snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"grid0", 0,"-x1.data" );
    }
    else if (frame<=2000)
    {
        snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"grid0", 1,"-x1.data" );
    }
    else
    {
        snprintf(hydrofile,sizeof(hydrofile),"%s%s%d%s",hydro_prefix,"grid0", 2,"-x1.data" );
    }

    fprintf(fPtr,"Reading Radius: %s\n", hydrofile);
    fflush(fPtr);

    hydroPtr=fopen(hydrofile, "r");

    i=0;
    while (i<R_DIM)
    {
        fscanf(hydroPtr, "%lf", (r_unprc+i));  //read value
        fgets(buf, 3,hydroPtr); //read comma

        if (i<5)
        {
            fprintf(fPtr,"R %d: %e\n", i, *(r_unprc+i));
            fflush(fPtr);
        }

        i++;
    }

    fclose(hydroPtr);

    //calculate radial grid edges
    *(r_edge+0)=r_in;
    i=0;
    for (i=1;i<R_DIM;i++)
    {
        *(r_edge+i)=(*(r_edge+i-1))+((*(r_edge+i-1))*(M_PI/560)/(1+((*(r_edge+i-1))/r_ref))); //r_i = r_(i-1) + Dq r_(i-1) [1 + r_(i-1)/r0]-1
        *(dr+i-1)=(*(r_edge+i))-(*(r_edge+i-1));

        if (i<5)
        {
            fprintf(fPtr,"R Edge: %d: %e Dr: %e\n", i, *(r_edge+i), *(dr+i-1));
            fflush(fPtr);
        }

    }
    free(r_edge);

    //Theta
    snprintf(hydrofile,sizeof(hydrofile),"%s%s",hydro_prefix,"grid-x2.data" );
    fprintf(fPtr,"Reading Theta: %s\n", hydrofile);
    fflush(fPtr);

    hydroPtr=fopen(hydrofile, "r");

    i=0;
    while (i<THETA_DIM)
    {
        fscanf(hydroPtr, "%lf", (theta_unprc+i));  //read value
        fgets(buf, 3,hydroPtr); //read comma

        if (i<5)
        {
            fprintf(fPtr,"R %d: %e\n", i, *(theta_unprc+i));
            fflush(fPtr);
        }

        i++;
    }

    fclose(hydroPtr);

    //Phi
    snprintf(hydrofile,sizeof(hydrofile),"%s%s",hydro_prefix,"grid-x3.data" );
    fprintf(fPtr,"Reading Phi: %s\n", hydrofile);
    fflush(fPtr);

    hydroPtr=fopen(hydrofile, "r");

    i=0;
    while (i<PHI_DIM)
    {
        fscanf(hydroPtr, "%lf", (phi_unprc+i));  //read value
        fgets(buf, 3,hydroPtr); //read comma

        if (i<5)
        {
            fprintf(fPtr,"R %d: %e\n", i, *(phi_unprc+i));
            fflush(fPtr);
        }

        i++;
    }

    fclose(hydroPtr);

    //limit number of array elements
    elem=0;
    for (i=0;i<PHI_DIM;i++)
    {
        for (j=0;j<THETA_DIM;j++)
        {
            for (k=0;k<R_DIM;k++)
            {
                //if I have photons do selection differently than if injecting photons
                if (ph_inj==0)
                {
                    //if calling this function when propagating photons, choose blocks based on where the photons are
                    if (((ph_rmin - C_LIGHT/fps)<(*(r_unprc+k))) && (*(r_unprc+k)  < (ph_rmax + C_LIGHT/fps) ))
                    {
                        // *(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )
                        elem++;
                    }
                }
                else
                {
                    //if calling this function to inject photons choose blocks based on injection parameters, r_inj, which is sufficient
                    if (((r_inj - C_LIGHT/fps)<(*(r_unprc+k))) && (*(r_unprc+k)  < (r_inj + C_LIGHT/fps) ))
                    {
                        // *(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )
                        elem++;
                    }

                }

            }
        }
    }

    fprintf(fPtr,"Number of post restricted Elems: %d %e\n", elem, r_inj);

    //allocate space for new set of data
    (*pres)=malloc (elem * sizeof (double ));
    (*velx)=malloc (elem * sizeof (double ));
    (*vely)=malloc (elem * sizeof (double ));
    (*velz)=malloc (elem * sizeof (double ));
    (*dens)=malloc (elem * sizeof (double ));
    (*x)=malloc (elem * sizeof (double ));
    (*y)=malloc (elem * sizeof (double ));
    (*z)=malloc (elem * sizeof (double ));
    (*r)=malloc (elem * sizeof (double ));
    (*theta)=malloc (elem * sizeof (double ));
    (*phi)=malloc (elem * sizeof (double ));
    (*gamma)=malloc (elem * sizeof (double ));
    (*dens_lab)=malloc (elem * sizeof (double ));
    (*szx)=malloc (elem * sizeof (double )); //theta and phi resolution
    (*szy)=malloc (elem * sizeof (double )); //r resolution
    (*temp)=malloc (elem * sizeof (double ));

    //limit number of array elements
    elem=0;
    for (i=0;i<PHI_DIM;i++)
    {
        for (j=0;j<THETA_DIM;j++)
        {
            for (k=0;k<R_DIM;k++)
            {
                //if I have photons do selection differently than if injecting photons
                if (ph_inj==0)
                {
                    //if calling this function when propagating photons, choose blocks based on where the photons are
                    if (((ph_rmin - C_LIGHT/fps)<(*(r_unprc+k))) && (*(r_unprc+k)  < (ph_rmax + C_LIGHT/fps) ))
                    {
                        (*pres)[elem] = *(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  ));
                        (*dens)[elem] = *(dens_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  ));
                        (*temp)[elem] =  pow(3*(*(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )))*pow(C_LIGHT,2.0)/(A_RAD) ,1.0/4.0);
                        //fprintf(fPtr, "Selected Values: pres, %e, temp, %e\n",*(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )),pow(3*(*(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )))*pow(C_LIGHT,2.0)/(A_RAD) ,1.0/4.0) );

                        (*gamma)[elem] = pow(pow(1.0-(pow(*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+ pow(*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+pow(*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)),0.5),-1);
                        (*dens_lab)[elem] = (*(dens_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )))*pow(pow(1.0-(pow(*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+ pow(*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+pow(*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)),0.5),-1);
                        (*r)[elem] = *(r_unprc+k);
                        (*theta)[elem] = *(theta_unprc+j);
                        (*phi)[elem] = *(phi_unprc+i);
                        (*x)[elem] = (*(r_unprc+k))*sin(*(theta_unprc+j))*cos(*(phi_unprc+i));
                        (*y)[elem] = (*(r_unprc+k))*sin(*(theta_unprc+j))*sin(*(phi_unprc+i));
                        (*z)[elem] = (*(r_unprc+k))*cos(*(theta_unprc+j));
                        (*szx)[elem] = M_PI/560;
                        (*szy)[elem] =  *(dr+k);
                        (*velx)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j))*cos(*(phi_unprc+i))) + ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))*cos(*(phi_unprc+i))) - ((*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(phi_unprc+i)));
                        (*vely)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j))*sin(*(phi_unprc+i))) + ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))*sin(*(phi_unprc+i))) + ((*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(phi_unprc+i)));
                        (*velz)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))) - ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j)));

                        elem++;

                    }
                }
                else
                {
                    //if calling this function to inject photons choose blocks based on injection parameters, r_inj, which is sufficient
                    if (((r_inj - C_LIGHT/fps)<(*(r_unprc+k))) && (*(r_unprc+k)  < (r_inj + C_LIGHT/fps) ))
                    {
                        (*pres)[elem] = *(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  ));
                        (*dens)[elem] = *(dens_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  ));
                        (*temp)[elem] =  pow(3*(*(pres_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )))*pow(C_LIGHT,2.0)/(A_RAD) ,1.0/4.0);
                        (*gamma)[elem] = pow(pow(1.0-(pow(*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+ pow(*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+pow(*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)),0.5),-1);
                        (*dens_lab)[elem] = (*(dens_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k  )))*pow(pow(1.0-(pow(*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+ pow(*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)+pow(*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )),2)),0.5),-1);
                        (*r)[elem] = *(r_unprc+k);
                        (*theta)[elem] = *(theta_unprc+j);
                        (*phi)[elem] = *(phi_unprc+i);
                        (*x)[elem] = (*(r_unprc+k))*sin(*(theta_unprc+j))*cos(*(phi_unprc+i));
                        (*y)[elem] = (*(r_unprc+k))*sin(*(theta_unprc+j))*sin(*(phi_unprc+i));
                        (*z)[elem] = (*(r_unprc+k))*cos(*(theta_unprc+j));
                        (*szx)[elem] = M_PI/560;
                        (*szy)[elem] =  *(dr+k);
                        (*velx)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j))*cos(*(phi_unprc+i))) + ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))*cos(*(phi_unprc+i))) - ((*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(phi_unprc+i)));
                        (*vely)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j))*sin(*(phi_unprc+i))) + ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))*sin(*(phi_unprc+i))) + ((*(vel_phi_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(phi_unprc+i)));
                        (*velz)[elem]=((*(vel_r_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*cos(*(theta_unprc+j))) - ((*(vel_theta_unprc+(i*R_DIM*THETA_DIM + j*R_DIM + k )))*sin(*(theta_unprc+j)));

                        elem++;
                    }

                }

            }
        }
    }

    *number=elem;

    free(pres_unprc); free(dens_unprc); free(r_unprc); free(theta_unprc); free(phi_unprc);free(dr);free(vel_r_unprc); free(vel_theta_unprc); free(vel_phi_unprc);

}

void photonInjection3D( struct photon **ph, int *ph_num, double r_inj, double ph_weight, int min_photons, int max_photons, char spect, int array_length, double fps, double theta_min, double theta_max,\
 double *x, double *y, double *z, double *szx, double *szy, double *r, double *theta, double *phi, double *temps, double *vx, double *vy, double *vz, gsl_rng * rand)
{
    int i=0, block_cnt=0, *ph_dens=NULL, ph_tot=0, j=0,k=0;
    double ph_dens_calc=0.0, fr_dum=0.0, y_dum=0.0, yfr_dum=0.0, fr_max=0, bb_norm=0, position_phi, ph_weight_adjusted, theta_prime=0;
    double com_v_phi, com_v_theta, *p_comv=NULL, *boost=NULL; //comoving phi, theta, comoving 4 momentum for a photon, and boost for photon(to go to lab frame)
    double *l_boost=NULL; //pointer to hold array of lorentz boost, to lab frame, values
    float num_dens_coeff;

    if (spect=='w') //from MCRAT paper, w for wien spectrum
    {
        num_dens_coeff=8.44;
        printf("in wien spectrum\n");
    }
    else
    {
        num_dens_coeff=20.29; //this is for black body spectrum
        printf("in BB spectrum");
    }

    //find how many blocks are near the injection radius within the angles defined in mc.par, get temperatures and calculate number of photons to allocate memory for
    //and then rcord which blocks have to have "x" amount of photons injected there
    printf("%e, %e\n",*(phi+i), theta_max);
    for(i=0;i<array_length;i++)
    {
        //look at all boxes in width delta r=c/fps and within angles we are interested in NEED TO modify for RIKEN data- dont need r anymore, just theta and phi? (didnt work), just look at pojection on x-z plane
        theta_prime=acos(*(y+i)/(*(r+i))); //jet axis here is the y axis
            if ( (theta_prime< theta_max) && (theta_prime > theta_min) ) //(*(r+i) > (r_inj - C_LIGHT/fps))  &&   (*(r+i)  < (r_inj + C_LIGHT/fps)  ) &&
            {
                //printf("%e\n", theta_prime );
                block_cnt++;
            }
    }
    printf("Blocks: %d\n", block_cnt);

    ph_dens=malloc(block_cnt * sizeof(int));

    //calculate the photon density for each block and save it to the array
    j=0;
    ph_tot=0;
    ph_weight_adjusted=ph_weight;
    //printf("%d %d\n", max_photons, min_photons);
    while ((ph_tot>max_photons) || (ph_tot<min_photons) )
    {
        j=0;
        ph_tot=0;
        //allocate memory to record density of photons for each block
        //ph_dens=malloc(block_cnt * sizeof(int));

        for (i=0;i<array_length;i++)
        {
            //printf("%d\n",i);
            //printf("%e, %e, %e, %e, %e, %e\n", *(r+i),(r_inj - C_LIGHT/fps), (r_inj + C_LIGHT/fps), *(theta+i) , theta_max, theta_min);
                //NEED TO modify for RIKEN data - modified
                theta_prime=acos(*(y+i)/(*(r+i)));
                if ( (theta_prime< theta_max) && (theta_prime > theta_min) )
                {
                    //NEED TO modify for RIKEN data - modified
                    ph_dens_calc=(num_dens_coeff*pow(*(temps+i),3.0)*pow(*(r+i),2)*sin(*(theta+i))* pow(*(szx+i),2.0)*(*(szy+i)) /(ph_weight_adjusted))*pow(pow(1.0-(pow(*(vx+i),2)+pow(*(vy+i),2)),0.5),-1) ; //a*T^3/(weight) dV, dV=2*PI*x*dx^2,

                     (*(ph_dens+j))=gsl_ran_poisson(rand,ph_dens_calc) ; //choose from poission distribution with mean of ph_dens_calc

                    //printf("%d, %lf \n",*(ph_dens+j), ph_dens_calc);

                     //sum up all the densities to get total number of photons
                     ph_tot+=(*(ph_dens+j));

                     j++;
                }
        }

        if (ph_tot>max_photons)
        {
            //if the number of photons is too big make ph_weight larger
            ph_weight_adjusted*=10;
            //free(ph_dens);
        }
        else if (ph_tot<min_photons)
        {
            ph_weight_adjusted*=0.5;
            //free(ph_dens);
        }
        //printf("dens: %d, photons: %d\n", *(ph_dens+(j-1)), ph_tot);

    }

    printf("%d\n", ph_tot);

    //allocate memory for that many photons and also allocate memory to hold comoving 4 momentum of each photon and the velocity of the fluid
    (*ph)=malloc (ph_tot * sizeof (struct photon ));

    p_comv=malloc(4*sizeof(double));
    boost=malloc(3*sizeof(double));
    l_boost=malloc(4*sizeof(double));

    //go through blocks and assign random energies/locations to proper number of photons
    ph_tot=0;
    k=0;
    for (i=0;i<array_length;i++)
    {
        theta_prime=acos(*(y+i)/(*(r+i)));
        if ( (theta_prime< theta_max) && (theta_prime > theta_min) )  //NEED TO modify for RIKEN data - modified
        {

            //*(temps+i)=0.76*(*(temps+i));
            for(j=0;j<( *(ph_dens+k) ); j++ )
            {
                    //have to get random frequency for the photon comoving frequency
                    y_dum=1; //initalize loop
                    yfr_dum=0;
                    while (y_dum>yfr_dum)
                    {
                        fr_dum=gsl_rng_uniform_pos(rand)*6.3e11*(*(temps+i)); //in Hz
                        //printf("%lf, %lf ",gsl_rng_uniform_pos(rand), (*(temps+i)));
                        y_dum=gsl_rng_uniform_pos(rand);
                        //printf("%lf ",fr_dum);

                        if (spect=='w')
                        {
                            yfr_dum=(1.0/(1.29e31))*pow((fr_dum/(*(temps+i))),3.0)/(exp((PL_CONST*fr_dum)/(K_B*(*(temps+i)) ))-1); //curve is normalized to maximum
                        }
                        else
                        {
                            fr_max=(5.88e10)*(*(temps+i));//(C_LIGHT*(*(temps+i)))/(0.29); //max frequency of bb
                            bb_norm=(PL_CONST*fr_max * pow((fr_max/C_LIGHT),2.0))/(exp(PL_CONST*fr_max/(K_B*(*(temps+i))))-1); //find value of bb at fr_max
                            yfr_dum=((1.0/bb_norm)*PL_CONST*fr_dum * pow((fr_dum/C_LIGHT),2.0))/(exp(PL_CONST*fr_dum/(K_B*(*(temps+i))))-1); //curve is normalized to vaue of bb @ max frequency
                        }
                        //printf("%lf, %lf,%lf,%e \n",(*(temps+i)),fr_dum, y_dum, yfr_dum);

                    }
                   //printf("%lf\n ",fr_dum);
                   //position_phi= gsl_rng_uniform(rand)*2*M_PI; //NEED TO modify for RIKEN data-modified, dont need anymore
                   com_v_phi=gsl_rng_uniform(rand)*2*M_PI;
                   com_v_theta=acos((gsl_rng_uniform(rand)*2)-1);
                   //printf("%lf, %lf, %lf\n", position_phi, com_v_phi, com_v_theta);

                   //populate 4 momentum comoving array
                   *(p_comv+0)=PL_CONST*fr_dum/C_LIGHT;
                   *(p_comv+1)=(PL_CONST*fr_dum/C_LIGHT)*sin(com_v_theta)*cos(com_v_phi);
                   *(p_comv+2)=(PL_CONST*fr_dum/C_LIGHT)*sin(com_v_theta)*sin(com_v_phi);
                   *(p_comv+3)=(PL_CONST*fr_dum/C_LIGHT)*cos(com_v_theta);

                    //populate boost matrix, not sure why multiplying by -1, seems to give correct answer in old python code...
                    //NEED TO modify for RIKEN data - modified
                    *(boost+0)=-1*(*(vx+i));
                    *(boost+1)=-1*(*(vy+i));
                    *(boost+2)=-1*(*(vz+i));
                    //printf("%lf, %lf, %lf\n", *(boost+0), *(boost+1), *(boost+2));

                    //boost to lab frame
                    lorentzBoost(boost, p_comv, l_boost, 'p');
                    //printf("Assignemnt: %e, %e, %e, %e\n", *(l_boost+0), *(l_boost+1), *(l_boost+2),*(l_boost+3));

                (*ph)[ph_tot].p0=(*(l_boost+0));
                (*ph)[ph_tot].p1=(*(l_boost+1));
                (*ph)[ph_tot].p2=(*(l_boost+2));
                (*ph)[ph_tot].p3=(*(l_boost+3));
                //NEED TO modify for RIKEN data-modified
                (*ph)[ph_tot].r0= (*(x+i)); //put photons @ center of box that they are supposed to be in with random phi
                (*ph)[ph_tot].r1=(*(y+i)) ;
                (*ph)[ph_tot].r2=(*(z+i)); //y coordinate in flash becomes z coordinate in MCRaT
                (*ph)[ph_tot].num_scatt=0;
                (*ph)[ph_tot].weight=ph_weight_adjusted;
                //printf("%d\n",ph_tot);
                ph_tot++;
            }
            k++;
        }
    }

    *ph_num=ph_tot; //save number of photons
    //printf(" %d: %d\n", *(ph_dens+(k-1)), *ph_num);
    free(ph_dens); free(p_comv);free(boost); free(l_boost);

}

 void phMinMax(struct photon *ph, int ph_num, double *min, double *max)
 {
      double temp_r_max=0, temp_r_min=-1;
        int i=0;
      double ph_r=0;

      for (i=0;i<ph_num;i++)
    {
        ph_r=pow(pow( ((ph+i)->r0), 2.0) + pow(((ph+i)->r1),2.0 ) + pow(((ph+i)->r2) , 2.0),0.5);
        if (ph_r > temp_r_max )
        {
            temp_r_max=ph_r;
            printf("The new max is: %e\n", temp_r_max);
        }

        if ((i==0) || (ph_r<temp_r_min))
        {
            temp_r_min=ph_r;
            printf("The new min is: %e\n", temp_r_min);
        }
    }

    *max=temp_r_max;
    *min=temp_r_min;

 }


#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <stdlib.h>

/* Representation */
gsl_complex *alloc_gsl_complex(void)
{
  gsl_complex *c = malloc(sizeof(gsl_complex));
  return c;
}

void free_gsl_complex(gsl_complex *c)
{
  free(c);
}

gsl_complex_float *alloc_gsl_complex_float(void)
{
  gsl_complex_float *c = malloc(sizeof(gsl_complex_float));
  return c;
}

void free_gsl_complex_float(gsl_complex_float *c)
{
  free(c);
}

void mgsl_complex_rect(double x, double y, gsl_complex *res)
{
  gsl_complex ret = gsl_complex_rect(x, y);
  *res = ret;
}

void mgsl_complex_float_rect(float x, float y, gsl_complex_float *res)
{
  GSL_SET_COMPLEX(res, x, y);
}

void mgsl_complex_polar(double r, double theta, gsl_complex *res)
{
  gsl_complex ret = gsl_complex_polar(r, theta);
  *res = ret;
}

void mgsl_complex_float_polar(float r, float theta, gsl_complex_float *res)
{
  gsl_complex ret = gsl_complex_polar((double)r, (double)theta);
  GSL_SET_COMPLEX(res, ret.dat[0], ret.dat[1]);
}

/* Properties */
double mgsl_complex_arg(gsl_complex *c)
{
  return gsl_complex_arg(*c);
}

double mgsl_complex_abs(gsl_complex *c)
{
  return gsl_complex_abs(*c);
}

double mgsl_complex_abs2(gsl_complex *c)
{
  return gsl_complex_abs2(*c);
}

double mgsl_complex_logabs(gsl_complex *c)
{
  return gsl_complex_logabs(*c);
}

/* Arithmetic operators */

void mgsl_complex_add(gsl_complex *a, gsl_complex *b, gsl_complex *res)
{
  *res = gsl_complex_add(*a, *b);
}

void mgsl_complex_sub(gsl_complex *a, gsl_complex *b, gsl_complex *res)
{
  *res = gsl_complex_sub(*a, *b);
}

void mgsl_complex_mul(gsl_complex *a, gsl_complex *b, gsl_complex *res)
{
  *res = gsl_complex_mul(*a, *b);
}

void mgsl_complex_div(gsl_complex *a, gsl_complex *b, gsl_complex *res)
{
  *res = gsl_complex_div(*a, *b);
}

void mgsl_complex_add_real(gsl_complex *a, double x, gsl_complex *res)
{
  *res = gsl_complex_add_real(*a, x);
}

void mgsl_complex_sub_real(gsl_complex *a, double x, gsl_complex *res)
{
  *res = gsl_complex_sub_real(*a, x);
}

void mgsl_complex_mul_real(gsl_complex *a, double x, gsl_complex *res)
{
  *res = gsl_complex_mul_real(*a, x);
}

void mgsl_complex_div_real(gsl_complex *a, double x, gsl_complex *res)
{
  *res = gsl_complex_div_real(*a, x);
}

void mgsl_complex_add_imag(gsl_complex *a, double y, gsl_complex *res)
{
  *res = gsl_complex_add_imag(*a, y);
}

void mgsl_complex_sub_imag(gsl_complex *a, double y, gsl_complex *res)
{
  *res = gsl_complex_sub_imag(*a, y);
}

void mgsl_complex_mul_imag(gsl_complex *a, double y, gsl_complex *res)
{
  *res = gsl_complex_mul_imag(*a, y);
}

void mgsl_complex_div_imag(gsl_complex *a, double y, gsl_complex *res)
{
  *res = gsl_complex_div_imag(*a, y);
}

void mgsl_complex_conjugate(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_conjugate(*z);
}

void mgsl_complex_inverse(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_inverse(*z);
}

void mgsl_complex_negative(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_negative(*z);
}

/* Elementary Complex Functions */
void mgsl_complex_sqrt(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_sqrt(*z);
}

void mgsl_complex_sqrt_real(double x, gsl_complex *res)
{
  *res = gsl_complex_sqrt_real(x);
}

void mgsl_complex_pow(gsl_complex *z, gsl_complex *a, gsl_complex *res)
{
  *res = gsl_complex_pow(*z, *a);
}

void mgsl_complex_pow_real(gsl_complex *z, double x, gsl_complex *res)
{
  *res = gsl_complex_pow_real(*z, x);
}

void mgsl_complex_exp(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_exp(*z);
}

void mgsl_complex_log(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_log(*z);
}

void mgsl_complex_log10(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_log10(*z);
}

void mgsl_complex_log_b(gsl_complex *z, gsl_complex *b, gsl_complex *res)
{
  *res = gsl_complex_log_b(*z, *b);
}

/* Complex Trigonometric Functions */
void mgsl_complex_sin(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_sin(*z);
}

void mgsl_complex_cos(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_cos(*z);
}

void mgsl_complex_tan(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_tan(*z);
}

void mgsl_complex_sec(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_sec(*z);
}

void mgsl_complex_csc(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_csc(*z);
}

void mgsl_complex_cot(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_cot(*z);
}

/* Inverse Complex Trigonometric Functions */
void mgsl_complex_arcsin(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arcsin(*z);
}

void mgsl_complex_arcsin_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arcsin_real(z);
}

void mgsl_complex_arccos(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccos(*z);
}

void mgsl_complex_arccos_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arccos_real(z);
}

void mgsl_complex_arctan(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arctan(*z);
}

void mgsl_complex_arcsec(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arcsec(*z);
}

void mgsl_complex_arcsec_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arcsec_real(z);
}

void mgsl_complex_arccsc(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccsc(*z);
}

void mgsl_complex_arccsc_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arccsc_real(z);
}

void mgsl_complex_arccot(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccot(*z);
}

/* Complex Hyperbolic Functions */
void mgsl_complex_sinh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_sinh(*z);
}

void mgsl_complex_cosh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_cosh(*z);
}

void mgsl_complex_tanh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_tanh(*z);
}

void mgsl_complex_sech(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_sech(*z);
}

void mgsl_complex_csch(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_csch(*z);
}

void mgsl_complex_coth(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_coth(*z);
}

/* Inverse Complex Hyperbolic Functions */
void mgsl_complex_arcsinh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arcsinh(*z);
}

void mgsl_complex_arccosh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccosh(*z);
}

void mgsl_complex_arccosh_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arccosh_real(z);
}

void mgsl_complex_arctanh(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arctanh(*z);
}

void mgsl_complex_arctanh_real(double z, gsl_complex *res)
{
  *res = gsl_complex_arctanh_real(z);
}

void mgsl_complex_arcsech(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arcsech(*z);
}

void mgsl_complex_arccsch(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccsch(*z);
}

void mgsl_complex_arccoth(gsl_complex *z, gsl_complex *res)
{
  *res = gsl_complex_arccoth(*z);
}


/*
ODE: a program to get optime Runge-Kutta and multi-steps methods.

Copyright 2011-2019, Javier Burguete Tolosa.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

	1. Redistributions of source code must retain the above copyright notice,
		this list of conditions and the following disclaimer.

	2. Redistributions in binary form must reproduce the above copyright notice,
		this list of conditions and the following disclaimer in the
		documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY Javier Burguete Tolosa ``AS IS'' AND ANY EXPRESS OR
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
SHALL Javier Burguete Tolosa OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/**
 * \file rk_6_4.c
 * \brief Source file to optimize Runge-Kutta 6 steps 4th order methods.
 * \author Javier Burguete Tolosa.
 * \copyright Copyright 2011-2019.
 */
#define _GNU_SOURCE
#include <string.h>
#include <math.h>
#include <libxml/parser.h>
#include <glib.h>
#include <libintl.h>
#include <gsl/gsl_rng.h>
#include "config.h"
#include "utils.h"
#include "optimize.h"
#include "rk.h"
#include "rk_6_4.h"

#define DEBUG_RK_6_4 0          ///< macro to debug.

/**
 * Function to obtain the coefficients of a 6 steps 4th order Runge-Kutta
 * method.
 */
int
rk_tb_6_4 (Optimize * optimize) ///< Optimize struct.
{
  long double A[4], B[4], C[4], D[4], E[4];
  long double *tb, *r;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  b21 (tb) = r[2];
  t3 (tb) = r[3];
  b31 (tb) = r[4];
  b32 (tb) = r[5];
  t4 (tb) = r[6];
  b41 (tb) = r[7];
  b42 (tb) = r[8];
  b43 (tb) = r[9];
  t5 (tb) = r[10];
  b54 (tb) = r[11];
  b65 (tb) = r[12];
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = t4 (tb);
  E[0] = 0.5L - b65 (tb) * t5 (tb);
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = D[0] * t4 (tb);
  E[1] = 1.L / 3.L - b65 (tb) * sqr (t5 (tb));
  A[2] = A[1] * t1 (tb);
  B[2] = B[1] * t2 (tb);
  C[2] = C[1] * t3 (tb);
  D[2] = D[1] * t4 (tb);
  E[2] = 0.25L - b65 (tb) * sqr (t5 (tb)) * t5 (tb);
  A[3] = 0.L;
  B[3] = b21 (tb) * t1 (tb) * (t2 (tb) - t5 (tb));
  C[3] = (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb)) * (t3 (tb) - t5 (tb));
  D[3] = (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb) + b43 (tb) * t3 (tb))
    * (t4 (tb) - t5 (tb));
  E[3] = 0.125L - 1.L / 6.L * t5 (tb);
  solve_4 (A, B, C, D, E);
  if (isnan (E[0]) || isnan (E[1]) || isnan (E[2]) || isnan (E[3]))
    return 0;
  b64 (tb) = E[3];
  b63 (tb) = E[2];
  b62 (tb) = E[1];
  b61 (tb) = E[0];
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = (1.L / 6.L - b62 (tb) * b21 (tb) * t1 (tb)
          - b63 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))
          - b64 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                        + b43 (tb) * t3 (tb))) / b65 (tb) - b54 (tb) * t4 (tb);
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = (1.L / 12.L - b62 (tb) * b21 (tb) * sqr (t1 (tb))
          - b63 (tb) * (b31 (tb) * sqr (t1 (tb)) + b32 (tb) * sqr (t2 (tb)))
          - b64 (tb) * (b41 (tb) * sqr (t1 (tb)) + b42 (tb) * sqr (t2 (tb))
                        + b43 (tb) * sqr (t3 (tb)))) / b65 (tb)
    - b54 (tb) * sqr (t4 (tb));
  A[2] = 0.L;
  B[2] = b21 (tb) * t1 (tb);
  C[2] = b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb);
  D[2] = (1.L / 24.L - b63 (tb) * b32 (tb) * b21 (tb) * t1 (tb)
          - b64 (tb) * (b42 (tb) * b21 (tb) * t1 (tb)
                        + b43 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))))
    / b65 (tb) - b54 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                             + b43 (tb) * t3 (tb));
  solve_3 (A, B, C, D);
  if (isnan (D[0]) || isnan (D[1]) || isnan (D[2]))
    return 0;
  b53 (tb) = D[2];
  b52 (tb) = D[1];
  b51 (tb) = D[0];
  rk_b_6 (tb);
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 4th order, 5th order in
 * equations depending only on time, Runge-Kutta method.
 */
int
rk_tb_6_4t (Optimize * optimize)        ///< Optimize struct.
{
  long double A[5], B[5], C[5], D[5], E[5], F[5];
  long double *tb, *r;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4t: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  b21 (tb) = r[2];
  t3 (tb) = r[3];
  b31 (tb) = r[4];
  b32 (tb) = r[5];
  t4 (tb) = r[6];
  b41 (tb) = r[7];
  b42 (tb) = r[8];
  b43 (tb) = r[9];
  t5 (tb) = r[10];
  b54 (tb) = r[11];
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = t4 (tb);
  E[0] = t5 (tb);
  F[0] = 0.5L;
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = D[0] * t4 (tb);
  E[1] = E[0] * t5 (tb);
  F[1] = 1.L / 3.L;
  A[2] = A[1] * t1 (tb);
  B[2] = B[1] * t2 (tb);
  C[2] = C[1] * t3 (tb);
  D[2] = D[1] * t4 (tb);
  E[2] = E[1] * t5 (tb);
  F[2] = 0.25L;
  A[3] = A[2] * t1 (tb);
  B[3] = B[2] * t2 (tb);
  C[3] = C[2] * t3 (tb);
  D[3] = D[2] * t4 (tb);
  E[3] = E[2] * t5 (tb);
  F[3] = 0.2L;
  A[4] = A[3] * t1 (tb);
  B[4] = B[3] * t2 (tb);
  C[4] = C[3] * t3 (tb);
  D[4] = D[3] * t4 (tb);
  E[4] = E[3] * t5 (tb);
  F[4] = 1.L / 6.L;
  solve_5 (A, B, C, D, E, F);
  if (isnan (F[0]) || isnan (F[1]) || isnan (F[2]) || isnan (F[3])
      || isnan (F[4]))
    return 0;
  b65 (tb) = F[4];
  b64 (tb) = F[3];
  b63 (tb) = F[2];
  b62 (tb) = F[1];
  b61 (tb) = F[0];
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = (1.L / 6.L - b62 (tb) * b21 (tb) * t1 (tb)
          - b63 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))
          - b64 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                        + b43 (tb) * t3 (tb))) / b65 (tb) - b54 (tb) * t4 (tb);
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = (1.L / 12.L - b62 (tb) * b21 (tb) * sqr (t1 (tb))
          - b63 (tb) * (b31 (tb) * sqr (t1 (tb)) + b32 (tb) * sqr (t2 (tb)))
          - b64 (tb) * (b41 (tb) * sqr (t1 (tb)) + b42 (tb) * sqr (t2 (tb))
                        + b43 (tb) * sqr (t3 (tb)))) / b65 (tb)
    - b54 (tb) * sqr (t4 (tb));
  A[2] = 0.L;
  B[2] = b21 (tb) * t1 (tb);
  C[2] = b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb);
  D[2] = (1.L / 24.L - b63 (tb) * b32 (tb) * b21 (tb) * t1 (tb)
          - b64 (tb) * (b42 (tb) * b21 (tb) * t1 (tb)
                        + b43 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))))
    / b65 (tb) - b54 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                             + b43 (tb) * t3 (tb));
  solve_3 (A, B, C, D);
  if (isnan (D[0]) || isnan (D[1]) || isnan (D[2]))
    return 0;
  b53 (tb) = D[2];
  b52 (tb) = D[1];
  b51 (tb) = D[0];
  rk_b_6 (tb);
#if DEBUG_RK_6_4
  rk_print_tb (optimize, "rk_tb_6_4t", stderr);
  fprintf (stderr, "rk_tb_6_4t: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 3th-4th order Runge-Kutta
 * pair.
 */
int
rk_tb_6_4p (Optimize * optimize)        ///< Optimize struct.
{
  long double A[4], B[4], C[4], D[4], E[4], AA[4], BB[4], CC[4], DD[4], EE[4];
  long double *tb, *r;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4p: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  b21 (tb) = r[2];
  t3 (tb) = r[3];
  b31 (tb) = r[4];
  b32 (tb) = r[5];
  t4 (tb) = r[6];
  b41 (tb) = r[7];
  b42 (tb) = r[8];
  b43 (tb) = r[9];
  t5 (tb) = r[10];
  b54 (tb) = r[11];
  b65 (tb) = r[12];
  A[0] = AA[0] = t1 (tb);
  B[0] = BB[0] = t2 (tb);
  C[0] = CC[0] = t3 (tb);
  D[0] = DD[0] = t4 (tb);
  EE[0] = 0.5L;
  E[0] = 0.5L - b65 (tb) * t5 (tb);
  A[1] = AA[1] = A[0] * t1 (tb);
  B[1] = BB[1] = B[0] * t2 (tb);
  C[1] = CC[1] = C[0] * t3 (tb);
  D[1] = DD[1] = D[0] * t4 (tb);
  EE[1] = 1.L / 3.L;
  E[1] = 1.L / 3.L - b65 (tb) * sqr (t5 (tb));
  A[2] = AA[2] = A[1] * t1 (tb);
  B[2] = BB[2] = B[1] * t2 (tb);
  C[2] = CC[2] = C[1] * t3 (tb);
  D[2] = DD[2] = D[1] * t4 (tb);
  EE[2] = 0.25L;
  E[2] = 0.25L - b65 (tb) * sqr (t5 (tb)) * t5 (tb);
  A[3] = AA[3] = 0.L;
  BB[3] = b21 (tb) * t1 (tb);
  B[3] = BB[3] * (t2 (tb) - t5 (tb));
  CC[3] = (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb));
  C[3] = CC[3] * (t3 (tb) - t5 (tb));
  DD[3] = (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb) + b43 (tb) * t3 (tb));
  D[3] = DD[3] * (t4 (tb) - t5 (tb));
  EE[3] = 1.L / 6.L;
  E[3] = 0.125L - 1.L / 6.L * t5 (tb);
  solve_4 (A, B, C, D, E);
  if (isnan (E[0]) || isnan (E[1]) || isnan (E[2]) || isnan (E[3]))
    return 0;
  b64 (tb) = E[3];
  b63 (tb) = E[2];
  b62 (tb) = E[1];
  b61 (tb) = E[0];
  solve_4 (AA, BB, CC, DD, EE);
  if (isnan (EE[0]) || isnan (EE[1]) || isnan (EE[2]) || isnan (EE[3]))
    return 0;
  e64 (tb) = EE[3];
  e63 (tb) = EE[2];
  e62 (tb) = EE[1];
  e61 (tb) = EE[0];
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = (1.L / 6.L - b62 (tb) * b21 (tb) * t1 (tb)
          - b63 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))
          - b64 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                        + b43 (tb) * t3 (tb))) / b65 (tb) - b54 (tb) * t4 (tb);
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = (1.L / 12.L - b62 (tb) * b21 (tb) * sqr (t1 (tb))
          - b63 (tb) * (b31 (tb) * sqr (t1 (tb)) + b32 (tb) * sqr (t2 (tb)))
          - b64 (tb) * (b41 (tb) * sqr (t1 (tb)) + b42 (tb) * sqr (t2 (tb))
                        + b43 (tb) * sqr (t3 (tb)))) / b65 (tb)
    - b54 (tb) * sqr (t4 (tb));
  A[2] = 0.L;
  B[2] = b21 (tb) * t1 (tb);
  C[2] = b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb);
  D[2] = (1.L / 24.L - b63 (tb) * b32 (tb) * b21 (tb) * t1 (tb)
          - b64 (tb) * (b42 (tb) * b21 (tb) * t1 (tb)
                        + b43 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))))
    / b65 (tb) - b54 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                             + b43 (tb) * t3 (tb));
  solve_3 (A, B, C, D);
  if (isnan (D[0]) || isnan (D[1]) || isnan (D[2]))
    return 0;
  b53 (tb) = D[2];
  b52 (tb) = D[1];
  b51 (tb) = D[0];
  rk_b_6 (tb);
  rk_e_6 (tb);
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4p: end\n");
#endif
  return 1;
}

/**
 * Function to obtain the coefficients of a 6 steps 3rd-4th order, 4th-5th order
 * in equations depending only on time, Runge-Kutta method.
 */
int
rk_tb_6_4tp (Optimize * optimize)       ///< Optimize struct.
{
  long double A[5], B[5], C[5], D[5], E[5], F[5], AA[4], BB[4], CC[4], DD[4],
    EE[4];
  long double *tb, *r;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_tb_6_4tp: start\n");
#endif
  tb = optimize->coefficient;
  r = optimize->random_data;
  t6 (tb) = 1.L;
  t1 (tb) = r[0];
  t2 (tb) = r[1];
  t3 (tb) = r[2];
  b31 (tb) = r[3];
  b32 (tb) = r[4];
  t4 (tb) = r[5];
  b41 (tb) = r[6];
  b42 (tb) = r[7];
  b43 (tb) = r[8];
  t5 (tb) = r[9];
  b54 (tb) = r[10];
  b65 (tb) = r[11];
  A[0] = AA[0] = t1 (tb);
  B[0] = BB[0] = t2 (tb);
  C[0] = CC[0] = t3 (tb);
  D[0] = DD[0] = t4 (tb);
  E[0] = t5 (tb);
  F[0] = EE[0] = 0.5L;
  A[1] = AA[1] = A[0] * t1 (tb);
  B[1] = BB[1] = B[0] * t2 (tb);
  C[1] = CC[1] = C[0] * t3 (tb);
  D[1] = DD[1] = D[0] * t4 (tb);
  E[1] = E[0] * t5 (tb);
  F[1] = EE[1] = 1.L / 3.L;
  A[2] = AA[2] = A[1] * t1 (tb);
  B[2] = BB[2] = B[1] * t2 (tb);
  C[2] = CC[2] = C[1] * t3 (tb);
  D[2] = DD[2] = D[1] * t4 (tb);
  E[2] = E[1] * t5 (tb);
  F[2] = EE[2] = 0.25L;
  A[3] = AA[3] = A[2] * t1 (tb);
  B[3] = BB[3] = B[2] * t2 (tb);
  C[3] = CC[3] = C[2] * t3 (tb);
  D[3] = DD[3] = D[2] * t4 (tb);
  E[3] = E[2] * t5 (tb);
  F[3] = EE[3] = 0.2L;
  A[4] = A[3] * t1 (tb);
  B[4] = B[3] * t2 (tb);
  C[4] = C[3] * t3 (tb);
  D[4] = D[3] * t4 (tb);
  E[4] = E[3] * t5 (tb);
  F[4] = 1.L / 6.L;
  solve_4 (AA, BB, CC, DD, EE);
  if (isnan (EE[0]) || isnan (EE[1]) || isnan (EE[2]) || isnan (EE[3]))
    return 0;
  e64 (tb) = EE[3];
  e63 (tb) = EE[2];
  e62 (tb) = EE[1];
  e61 (tb) = EE[0];
  solve_5 (A, B, C, D, E, F);
  if (isnan (F[0]) || isnan (F[1]) || isnan (F[2]) || isnan (F[3])
      || isnan (F[4]))
    return 0;
  b65 (tb) = F[4];
  b64 (tb) = F[3];
  b63 (tb) = F[2];
  b62 (tb) = F[1];
  b61 (tb) = F[0];
  b21 (tb) = (1.L / 6.L - e63 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))
              - e64 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                            + b43 (tb) * t3 (tb))) / (e62 (tb) * t1 (tb));
  if (isnan (b21 (tb)))
    return 0;
  A[0] = t1 (tb);
  B[0] = t2 (tb);
  C[0] = t3 (tb);
  D[0] = (1.L / 6.L - b62 (tb) * b21 (tb) * t1 (tb)
          - b63 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))
          - b64 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                        + b43 (tb) * t3 (tb))) / b65 (tb) - b54 (tb) * t4 (tb);
  A[1] = A[0] * t1 (tb);
  B[1] = B[0] * t2 (tb);
  C[1] = C[0] * t3 (tb);
  D[1] = (1.L / 12.L - b62 (tb) * b21 (tb) * sqr (t1 (tb))
          - b63 (tb) * (b31 (tb) * sqr (t1 (tb)) + b32 (tb) * sqr (t2 (tb)))
          - b64 (tb) * (b41 (tb) * sqr (t1 (tb)) + b42 (tb) * sqr (t2 (tb))
                        + b43 (tb) * sqr (t3 (tb)))) / b65 (tb)
    - b54 (tb) * sqr (t4 (tb));
  A[2] = 0.L;
  B[2] = b21 (tb) * t1 (tb);
  C[2] = b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb);
  D[2] = (1.L / 24.L - b63 (tb) * b32 (tb) * b21 (tb) * t1 (tb)
          - b64 (tb) * (b42 (tb) * b21 (tb) * t1 (tb)
                        + b43 (tb) * (b31 (tb) * t1 (tb) + b32 (tb) * t2 (tb))))
    / b65 (tb) - b54 (tb) * (b41 (tb) * t1 (tb) + b42 (tb) * t2 (tb)
                             + b43 (tb) * t3 (tb));
  solve_3 (A, B, C, D);
  if (isnan (D[0]) || isnan (D[1]) || isnan (D[2]))
    return 0;
  b53 (tb) = D[2];
  b52 (tb) = D[1];
  b51 (tb) = D[0];
  rk_b_6 (tb);
#if DEBUG_RK_6_4
  rk_print_tb (optimize, "rk_tb_6_4t", stderr);
  fprintf (stderr, "rk_tb_6_4t: end\n");
#endif
  return 1;
}

/**
 * Function to calculate the objective function of a 6 steps 4th order
 * Runge-Kutta method.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_4 (RK * rk)   ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4: start\n");
#endif
  tb = rk->tb->coefficient;
  o = fminl (0.L, b20 (tb));
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b51 (tb) < 0.L)
    o += b51 (tb);
  if (b52 (tb) < 0.L)
    o += b52 (tb);
  if (b53 (tb) < 0.L)
    o += b53 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b61 (tb) < 0.L)
    o += b61 (tb);
  if (b62 (tb) < 0.L)
    o += b62 (tb);
  if (b63 (tb) < 0.L)
    o += b63 (tb);
  if (b64 (tb) < 0.L)
    o += b64 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_4: end\n");
#endif
  return o;
}

/**
 * Function to calculate the objective function of a 6 steps 4th order, 5th
 * order in equations depending only on time, Runge-Kutta method.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_4t (RK * rk)  ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4t: start\n");
#endif
  tb = rk->tb->coefficient;
#if DEBUG_RK_6_4
  rk_print_tb (optimize, "rk_objective_tb_6_4t", stderr);
#endif
  o = fminl (0.L, b20 (tb));
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b51 (tb) < 0.L)
    o += b51 (tb);
  if (b52 (tb) < 0.L)
    o += b52 (tb);
  if (b53 (tb) < 0.L)
    o += b53 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b61 (tb) < 0.L)
    o += b61 (tb);
  if (b62 (tb) < 0.L)
    o += b62 (tb);
  if (b63 (tb) < 0.L)
    o += b63 (tb);
  if (b64 (tb) < 0.L)
    o += b64 (tb);
  if (b65 (tb) < 0.L)
    o += b65 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4t: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_4t: end\n");
#endif
  return o;
}

/**
 * Function to calculate the objective function of a 6 steps 3rd-4th order
 * Runge-Kutta pair.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_4p (RK * rk)  ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4p: start\n");
#endif
  tb = rk->tb->coefficient;
  o = fminl (0.L, b20 (tb));
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b51 (tb) < 0.L)
    o += b51 (tb);
  if (b52 (tb) < 0.L)
    o += b52 (tb);
  if (b53 (tb) < 0.L)
    o += b53 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b61 (tb) < 0.L)
    o += b61 (tb);
  if (b62 (tb) < 0.L)
    o += b62 (tb);
  if (b63 (tb) < 0.L)
    o += b63 (tb);
  if (b64 (tb) < 0.L)
    o += b64 (tb);
  if (e60 (tb) < 0.L)
    o += e60 (tb);
  if (e61 (tb) < 0.L)
    o += e61 (tb);
  if (e62 (tb) < 0.L)
    o += e62 (tb);
  if (e63 (tb) < 0.L)
    o += e63 (tb);
  if (e64 (tb) < 0.L)
    o += e64 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4p: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_4p: end\n");
#endif
  return o;
}

/**
 * Function to calculate the objective function of a 6 steps 3rd-4th order,
 * 4th-5th order in equations depending only in time, Runge-Kutta pair.
 *
 * \return objective function value.
 */
long double
rk_objective_tb_6_4tp (RK * rk) ///< RK struct.
{
  long double *tb;
  long double o;
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4tp: start\n");
#endif
  tb = rk->tb->coefficient;
#if DEBUG_RK_6_4
  rk_print_tb (optimize, "rk_objective_tb_6_4tp", stderr);
#endif
  o = fminl (0.L, b20 (tb));
  if (b21 (tb) < 0.L)
    o += b21 (tb);
  if (b30 (tb) < 0.L)
    o += b30 (tb);
  if (b40 (tb) < 0.L)
    o += b40 (tb);
  if (b50 (tb) < 0.L)
    o += b50 (tb);
  if (b51 (tb) < 0.L)
    o += b51 (tb);
  if (b52 (tb) < 0.L)
    o += b52 (tb);
  if (b53 (tb) < 0.L)
    o += b53 (tb);
  if (b60 (tb) < 0.L)
    o += b60 (tb);
  if (b61 (tb) < 0.L)
    o += b61 (tb);
  if (b62 (tb) < 0.L)
    o += b62 (tb);
  if (b63 (tb) < 0.L)
    o += b63 (tb);
  if (b64 (tb) < 0.L)
    o += b64 (tb);
  if (b65 (tb) < 0.L)
    o += b65 (tb);
  if (e60 (tb) < 0.L)
    o += e60 (tb);
  if (e61 (tb) < 0.L)
    o += e61 (tb);
  if (e62 (tb) < 0.L)
    o += e62 (tb);
  if (e63 (tb) < 0.L)
    o += e63 (tb);
  if (e64 (tb) < 0.L)
    o += e64 (tb);
  if (o < 0.L)
    {
      o = 40.L - o;
      goto end;
    }
  o = 30.L
    + fmaxl (1.L,
             fmaxl (t1 (tb),
                    fmaxl (t2 (tb),
                           fmaxl (t3 (tb), fmaxl (t4 (tb), t5 (tb))))));
  if (rk->strong)
    {
      rk_bucle_ac (rk);
      o = fminl (o, *rk->ac0->optimal);
    }
end:
#if DEBUG_RK_6_4
  fprintf (stderr, "rk_objective_tb_6_4tp: optimal=%Lg\n", o);
  fprintf (stderr, "rk_objective_tb_6_4tp: end\n");
#endif
  return o;
}


#ifndef __CNN_CONV_H__
#define __CNN_CONV_H__

#include <cblas.h>
#include <string.h>

#include "cnn_macro.h"
#include "cnn_types.h"

#ifdef CNN_WITH_CUDA
#include <cublas_v2.h>
#include <cuda_runtime.h>

#include "cnn_init.h"

void cnn_map_gpu(float* dst, float* src, int* map, int len);
void cnn_map_inv_gpu(float* dst, float* src, int* map, int len);
#endif

static inline void cnn_conv_unroll_2d_valid(int* indexMap, int dstHeight,
                                            int dstWidth, int kSize,
                                            int srcHeight, int srcWidth,
                                            int srcCh)
{
    int __kMemSize = kSize * kSize;
    int __srcImSize = srcHeight * srcWidth;
    int __indexMapCols = __kMemSize * srcCh;

    for (int __h = 0; __h < dstHeight; __h++)
    {
        int __dstRowShift = __h * dstWidth;

        for (int __w = 0; __w < dstWidth; __w++)
        {
            int __indexMapRow = __dstRowShift + __w;
            int __indexMemBase = __indexMapRow * __indexMapCols;

            for (int __ch = 0; __ch < srcCh; __ch++)
            {
                int __indexMemShiftBase = __indexMemBase + __kMemSize * __ch;
                int __srcChShift = __ch * __srcImSize;

                for (int __convH = 0; __convH < kSize; __convH++)
                {
                    int __indexMemShift = __indexMemShiftBase + __convH * kSize;
                    int __srcShift = (__h + __convH) * srcWidth + __srcChShift;

                    for (int __convW = 0; __convW < kSize; __convW++)
                    {
                        indexMap[__indexMemShift + __convW] =
                            __srcShift + (__w + __convW);
                    }
                }
            }
        }
    }
}

static inline void cnn_conv_unroll_2d_same(int* indexMap, int dstHeight,
                                           int dstWidth, int kSize,
                                           int srcHeight, int srcWidth,
                                           int srcCh)
{
    int __kMemSize = kSize * kSize;
    int __srcImSize = srcHeight * srcWidth;
    int __indexMapCols = __kMemSize * srcCh;

    int __convHBase = -kSize / 2;
    int __convWBase = -kSize / 2;

    for (int __h = 0; __h < dstHeight; __h++)
    {
        int __dstRowShift = __h * dstWidth;

        for (int __w = 0; __w < dstWidth; __w++)
        {
            int __indexMapRow = __dstRowShift + __w;
            int __indexMemBase = __indexMapRow * __indexMapCols;

            for (int __ch = 0; __ch < srcCh; __ch++)
            {
                int __indexMemShiftBase = __indexMemBase + __kMemSize * __ch;
                int __srcChShift = __ch * __srcImSize;

                for (int __convH = 0; __convH < kSize; __convH++)
                {
                    int __indexMemShift = __indexMemShiftBase + __convH * kSize;
                    int __convHIndex = __h + __convH + __convHBase;

                    if (__convHIndex >= 0 && __convHIndex < srcHeight)
                    {
                        int __srcShift = __convHIndex * srcWidth + __srcChShift;

                        for (int __convW = 0; __convW < kSize; __convW++)
                        {
                            int __convWIndex = __w + __convW + __convWBase;
                            if (__convWIndex >= 0 && __convWIndex < srcWidth)
                            {
                                int __tmpIndex = __srcShift + __convWIndex;

                                indexMap[__indexMemShift + __convW] =
                                    __tmpIndex;
                            }
                        }
                    }
                }
            }
        }
    }
}

static inline void cnn_conv_2d(float* dst, int dstHeight, int dstWidth,
                               float* kernel, int kSize, int chIn, int chOut,
                               float* src, int srcHeight, int srcWidth)
{
    int __kMemSize = kSize * kSize;
    int __filterSize = chIn * __kMemSize;
    int __dstImSize = dstHeight * dstWidth;
    int __srcImSize = srcHeight * srcWidth;

    for (int __chOut = 0; __chOut < chOut; __chOut++)
    {
        int __filterShift = __chOut * __filterSize;
        int __dstChShift = __chOut * __dstImSize;

        for (int __chIn = 0; __chIn < chIn; __chIn++)
        {
            int __kShiftBase = __chIn * __kMemSize + __filterShift;
            int __srcChShift = __chIn * __srcImSize;

            for (int __h = 0; __h < dstHeight; __h++)
            {
                int __dstShift = __h * dstWidth + __dstChShift;
                for (int __w = 0; __w < dstWidth; __w++)
                {
                    float __conv = 0;
                    for (int __convH = 0; __convH < kSize; __convH++)
                    {
                        int __kShift = __convH * kSize + __kShiftBase;
                        int __srcShift =
                            (__h + __convH) * srcWidth + __srcChShift;

                        for (int __convW = 0; __convW < kSize; __convW++)
                        {
                            __conv += kernel[__kShift + __convW] *
                                      src[__srcShift + (__w + __convW)];
                        }
                    }

                    dst[__dstShift + __w] += __conv;
                }
            }
        }
    }
}

static inline void cnn_conv_2d_grad(float* srcGrad, int srcHeight, int srcWidth,
                                    float* kernel, int kSize, int srcCh,
                                    int lCh, float* lGrad, int lHeight,
                                    int lWidth)
{
    int __kMemSize = kSize * kSize;
    int __filterSize = srcCh * __kMemSize;
    int __lImSize = lHeight * lWidth;
    int __srcImSize = srcHeight * srcWidth;

    for (int __lCh = 0; __lCh < lCh; __lCh++)
    {
        int __filterShift = __lCh * __filterSize;
        int __lChShift = __lCh * __lImSize;

        for (int __srcCh = 0; __srcCh < srcCh; __srcCh++)
        {
            int __srcChShift = __srcCh * __srcImSize;
            int __kShiftBase = __srcCh * __kMemSize + __filterShift;

            for (int __h = 0; __h < lHeight; __h++)
            {
                int __lShift = __h * lHeight + __lChShift;
                for (int __w = 0; __w < lWidth; __w++)
                {
                    for (int __convH = 0; __convH < kSize; __convH++)
                    {
                        int __kShift = __convH * kSize + __kShiftBase;
                        int __srcShift =
                            (__h + __convH) * srcWidth + __srcChShift;

                        for (int __convW = 0; __convW < kSize; __convW++)
                        {
                            srcGrad[__srcShift + (__w + __convW)] +=
                                lGrad[__lShift + __w] *
                                kernel[__kShift + __convW];
                        }
                    }
                }
            }
        }
    }
}

static inline void cnn_conv_2d_kernel_grad(float* lGrad, int lHeight,
                                           int lWidth, float* kGrad, int kSize,
                                           int lCh, int srcCh, float* src,
                                           int srcHeight, int srcWidth)
{
    int __kMemSize = kSize * kSize;
    int __filterSize = srcCh * __kMemSize;
    int __lImSize = lHeight * lWidth;
    int __srcImSize = srcHeight * srcWidth;

    for (int __lCh = 0; __lCh < lCh; __lCh++)
    {
        int __filterShift = __lCh * __filterSize;
        int __lChShift = __lCh * __lImSize;

        for (int __srcCh = 0; __srcCh < srcCh; __srcCh++)
        {
            int __srcChShift = __srcCh * __srcImSize;
            int __kShiftBase = __srcCh * __kMemSize + __filterShift;

            for (int __h = 0; __h < lHeight; __h++)
            {
                int __lShift = __h * lWidth + __lChShift;
                for (int __w = 0; __w < lWidth; __w++)
                {
                    for (int __convH = 0; __convH < kSize; __convH++)
                    {
                        int __kShift = __convH * kSize + __kShiftBase;
                        int __srcShift =
                            (__h + __convH) * srcWidth + __srcChShift;

                        for (int __convW = 0; __convW < kSize; __convW++)
                        {
                            kGrad[__kShift + __convW] +=
                                lGrad[__lShift + __w] *
                                src[__srcShift + (__w + __convW)];
                        }
                    }
                }
            }
        }
    }
}

static inline void cnn_forward_conv(union CNN_LAYER* layerRef,
                                    struct CNN_CONFIG* cfgRef, int layerIndex)
{
#ifdef CNN_WITH_CUDA
    float alpha = 1.0;
    float beta = 0.0;

    struct CNN_LAYER_CONV* layerPtr = &layerRef[layerIndex].conv;

    struct CNN_MAT* outData = &layerPtr->outMat.data;
    struct CNN_MAT* preOutData = &layerRef[layerIndex - 1].outMat.data;

    cnn_assert_cudnn(cudnnConvolutionForward(
        cnnInit.cudnnHandle,                        //
        &alpha,                                     //
        layerPtr->srcTen, preOutData->mat,          //
        layerPtr->kernelTen, layerPtr->kernel.mat,  //
        layerPtr->convDesc, layerPtr->convAlgoFW, cnnInit.wsData,
        cnnInit.wsSize,  //
        &beta,           //
        layerPtr->dstTen, outData->mat));

#if defined(CNN_CONV_BIAS_FILTER)
    beta = 1.0;
    cnn_assert_cudnn(cudnnAddTensor(cnnInit.cudnnHandle,                    //
                                    &alpha,                                 //
                                    layerPtr->biasTen, layerPtr->bias.mat,  //
                                    &beta,                                  //
                                    layerPtr->dstTen, outData->mat));
#endif

#else
    // Cache
    int mapRows =
        layerRef[layerIndex].outMat.width * layerRef[layerIndex].outMat.height;
    int mapCols = layerRef[layerIndex - 1].outMat.channel *
                  cfgRef->layerCfg[layerIndex].conv.size *
                  cfgRef->layerCfg[layerIndex].conv.size;
    int mapSize = mapRows * mapCols;
    int* indexMap = layerRef[layerIndex].conv.indexMap;

    int chOut = layerRef[layerIndex].outMat.channel;
    float* kernel = layerRef[layerIndex].conv.kernel.mat;

    // Clear outputs
    // memset(layerRef[layerIndex].outMat.data.mat, 0,
    //       sizeof(float) * layerRef[layerIndex].outMat.data.rows *
    //           layerRef[layerIndex].outMat.data.cols);

    for (int j = 0; j < cfgRef->batch; j++)
    {
        int srcShift = j * layerRef[layerIndex - 1].outMat.data.cols;
        int dstShift = j * layerRef[layerIndex].outMat.data.cols;
        int mapShift = j * mapSize;

        float* srcPtr = layerRef[layerIndex - 1].outMat.data.mat + srcShift;
        float* dstPtr = layerRef[layerIndex].outMat.data.mat + dstShift;
        float* mapPtr = layerRef[layerIndex].conv.unroll.mat + mapShift;

        for (int k = 0; k < mapSize; k++)
        {
            int tmpIndex = indexMap[k];
            if (tmpIndex >= 0)
            {
                mapPtr[k] = srcPtr[tmpIndex];
            }
        }

        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, chOut, mapRows,
                    mapCols, 1.0, kernel, mapCols, mapPtr, mapCols, 0.0, dstPtr,
                    mapRows);

        // Add bias
#if defined(CNN_CONV_BIAS_FILTER)
#ifdef DEBUG
#pragma message("cnn_forward_conv(): Enable convolution filter bias")
#endif
        for (int ch = 0; ch < chOut; ch++)
        {
            cblas_saxpy(
                mapRows, 1.0, &layerRef[layerIndex].conv.bias.mat[ch], 0,
                &layerRef[layerIndex].outMat.data.mat[dstShift + ch * mapRows],
                1);
        }
#elif defined(CNN_CONV_BIAS_LAYER)
#ifdef DEBUG
#pragma message("cnn_forward_conv(): Enable convolution layer bias")
#endif
        cblas_saxpy(layerRef[layerIndex].conv.bias.cols, 1.0,
                    layerRef[layerIndex].conv.bias.mat, 1,
                    &layerRef[layerIndex].outMat.data.mat[dstShift], 1);
#endif
    }
#endif
}

static inline void cnn_backward_conv(union CNN_LAYER* layerRef,
                                     struct CNN_CONFIG* cfgRef, int layerIndex)
{
#ifdef CNN_WITH_CUDA
    float alpha = 1.0;
    float beta = 1.0;

    struct CNN_LAYER_CONV* layerPtr = &layerRef[layerIndex].conv;

    struct CNN_MAT* outData = &layerPtr->outMat.data;
    struct CNN_MAT* preOutData = &layerRef[layerIndex - 1].outMat.data;

    cnn_assert_cudnn(cudnnConvolutionBackwardFilter(
        cnnInit.cudnnHandle,                //
        &alpha,                             //
        layerPtr->srcTen, preOutData->mat,  //
        layerPtr->dstTen, outData->grad,    //
        layerPtr->convDesc, layerPtr->convAlgoBWFilter, cnnInit.wsData,
        cnnInit.wsSize,  //
        &beta,           //
        layerPtr->kernelTen, layerPtr->kernel.grad));

    cnn_assert_cudnn(
        cudnnConvolutionBackwardBias(cnnInit.cudnnHandle,              //
                                     &alpha,                           //
                                     layerPtr->dstTen, outData->grad,  //
                                     &beta,                            //
                                     layerPtr->biasTen, layerPtr->bias.grad));

#else
    // Cache
    int mapRows =
        layerRef[layerIndex].outMat.width * layerRef[layerIndex].outMat.height;
    int mapCols = layerRef[layerIndex - 1].outMat.channel *
                  cfgRef->layerCfg[layerIndex].conv.size *
                  cfgRef->layerCfg[layerIndex].conv.size;
    int mapSize = mapRows * mapCols;

    int* indexMap = layerRef[layerIndex].conv.indexMap;

    int chOut = layerRef[layerIndex].outMat.channel;
    float* kernel = layerRef[layerIndex].conv.kernel.mat;
    float* kGrad = layerRef[layerIndex].conv.kernel.grad;

    // Sum gradient
    for (int j = 0; j < cfgRef->batch; j++)
    {
        int gradShift = j * layerRef[layerIndex].outMat.data.cols;
        int mapShift = j * mapSize;

        float* gradPtr = &layerRef[layerIndex].outMat.data.grad[gradShift];
        float* mapPtr = &layerRef[layerIndex].conv.unroll.mat[mapShift];

        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, chOut, mapCols,
                    mapRows, 1.0, gradPtr, mapRows, mapPtr, mapCols, 1.0, kGrad,
                    mapCols);

        // Sum bias gradient matrix
#if defined(CNN_CONV_BIAS_FILTER)
#ifdef DEBUG
#pragma message("cnn_forward_conv(): Enable convolution filter bias")
#endif
        for (int ch = 0; ch < chOut; ch++)
        {
            cblas_saxpy(mapRows, 1.0, &gradPtr[ch * mapRows], 1,
                        &layerRef[layerIndex].conv.bias.grad[ch], 0);
        }
#elif defined(CNN_CONV_BIAS_LAYER)
#ifdef DEBUG
#pragma message("cnn_forward_conv(): Enable convolution layer bias")
#endif
        cblas_saxpy(layerRef[layerIndex].conv.bias.cols, 1.0, gradPtr, 1,
                    layerRef[layerIndex].conv.bias.grad, 1);
#endif
    }
#endif

    // Find layer gradient
    if (layerIndex > 1)
    {
#ifdef CNN_WITH_CUDA
        cudaMemset(preOutData->grad, 0,
                   sizeof(float) * preOutData->rows * preOutData->cols);
        beta = 0.0;
        cnn_assert_cudnn(cudnnConvolutionBackwardData(
            cnnInit.cudnnHandle,                        //
            &alpha,                                     //
            layerPtr->kernelTen, layerPtr->kernel.mat,  //
            layerPtr->dstTen, outData->grad,            //
            layerPtr->convDesc, layerPtr->convAlgoBWGrad, cnnInit.wsData,
            cnnInit.wsSize,  //
            &beta,           //
            layerPtr->srcTen, preOutData->grad));
#else
        memset(layerRef[layerIndex - 1].outMat.data.grad, 0,
               sizeof(float) * layerRef[layerIndex - 1].outMat.data.rows *
                   layerRef[layerIndex - 1].outMat.data.cols);
        for (int j = 0; j < cfgRef->batch; j++)
        {
            int gradShift = j * layerRef[layerIndex].outMat.data.cols;
            int preGradShift = j * layerRef[layerIndex - 1].outMat.data.cols;
            int mapShift = j * mapSize;

            float* gradPtr = layerRef[layerIndex].outMat.data.grad + gradShift;
            float* preGradPtr =
                layerRef[layerIndex - 1].outMat.data.grad + preGradShift;
            float* mapPtr = layerRef[layerIndex].conv.unroll.grad + mapShift;

            cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, mapRows,
                        mapCols, chOut, 1.0, gradPtr, mapRows, kernel, mapCols,
                        0.0, mapPtr, mapCols);

            for (int i = 0; i < mapSize; i++)
            {
                int tmpIndex = indexMap[i];
                if (tmpIndex >= 0)
                {
                    preGradPtr[tmpIndex] += mapPtr[i];
                }
            }
        }
#endif
    }
}

#endif


/* multifit_nlinear/fdfvv.c
 *
 * Copyright (C) 2015 Patrick Alken
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_multifit_nlinear.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>

/*
fdfvv()
  Compute approximate second directional derivative using
finite differences.

See Eq. 19 of:

M. K. Transtrum, J. P. Sethna, Improvements to the Levenberg
Marquardt algorithm for nonlinear least-squares minimization,
arXiv:1201.5885, 2012.

Inputs: h     - step size for finite difference
        x     - parameter vector, size p
        v     - geodesic velocity, size p
        f     - vector of function values f_i(x), size n
        J     - Jacobian matrix J(x), n-by-p
        swts  - data weights
        fdf   - fdf struct
        fvv   - (output) approximate second directional derivative
                vector D_v^2 f(x)
        work  - workspace, size p

Return: success or error
*/

static int
fdfvv(const double h, const gsl_vector *x, const gsl_vector *v,
      const gsl_vector *f, const gsl_matrix *J, const gsl_vector *swts,
      gsl_multifit_nlinear_fdf *fdf, gsl_vector *fvv, gsl_vector *work)
{
  int status;
  const size_t n = fdf->n;
  const size_t p = fdf->p;
  const double hinv = 1.0 / h;
  size_t i;

  /* compute work = x + h*v */
  for (i = 0; i < p; ++i)
    {
      double xi = gsl_vector_get(x, i);
      double vi = gsl_vector_get(v, i);

      gsl_vector_set(work, i, xi + h * vi);
    }

  /* compute f(x + h*v) */
  status = gsl_multifit_nlinear_eval_f (fdf, work, swts, fvv);
  if (status)
    return status;

  for (i = 0; i < n; ++i)
    {
      double fi = gsl_vector_get(f, i);    /* f_i(x) */
      double fip = gsl_vector_get(fvv, i); /* f_i(x + h*v) */
      gsl_vector_const_view row = gsl_matrix_const_row(J, i);
      double u, fvvi;

      /* compute u = sum_{ij} J_{ij} D v_j */
      gsl_blas_ddot(&row.vector, v, &u);

      fvvi = (2.0 * hinv) * ((fip - fi) * hinv - u);

      gsl_vector_set(fvv, i, fvvi);
    }

  return status;
}

/*
gsl_multifit_nlinear_fdfvv()
  Compute approximate second directional derivative
using finite differences

Inputs: h    - step size for finite difference
        x    - parameter vector, size p
        v    - geodesic velocity, size p
        f    - function values f_i(x), size n
        J    - Jacobian matrix J(x), n-by-p
        swts - sqrt data weights (set to NULL if not needed)
        fdf  - fdf
        fvv  - (output) approximate (weighted) second directional derivative
               vector, size n, sqrt(W) fvv
        work - workspace, size p

Return: success or error
*/

int
gsl_multifit_nlinear_fdfvv(const double h, const gsl_vector *x, const gsl_vector *v,
                           const gsl_vector *f, const gsl_matrix *J,
                           const gsl_vector *swts, gsl_multifit_nlinear_fdf *fdf,
                           gsl_vector *fvv, gsl_vector *work)
{
  return fdfvv(h, x, v, f, J, swts, fdf, fvv, work);
}


/******************************************************************************
 *                                                                            *
 * ESTIMATE_THETAE.C                                                          *
 *                                                                            *
 * ESTIMATE THETAE AT END OF TIMESTEP DUE TO RADIATION SOURCES                *
 *                                                                            *
 ******************************************************************************/

#include "decs.h"

#if RADIATION
#if ESTIMATE_THETAE

#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_roots.h>

#define NTHREADS (24)

// ESTIMATE THETAE BEFORE PUSH TO AVOID ISSUES WITH MPI COMMUNICATION

double Thetae_est[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG];
double Thetae_old[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG];

double Ucon[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NDIM];
double Ucov[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NDIM];
double Bcov[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NDIM];
double Ne[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG];
double Bmag[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG];

int        Nsph_zone[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
int        Nsph_cntd[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
double *   w[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
double *   nu[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
double *   dlam[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
double *   theta[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NTHREADS];
static int type        = 0; // TODO: remove me when type loops are in place.
static int interaction = 0; // TODO: remove me when type loops are in place.

double get_Thetae_est(int i, int j, int k) { return Thetae_est[i][j][k]; }

// Rootfinding parameters and function
struct of_params {
  int    i, j, k;
  double rho, Ne, Bmag, Thetaei, Ucon0, dt;
  double extra[EOS_NUM_EXTRA];
};

double dEdt(double Thetae, void *params) {
  struct of_params *p     = (struct of_params *)params;
  int               i     = p->i;
  int               j     = p->j;
  int               k     = p->k;
  double            rho   = p->rho;
  double            dt    = p->dt;
  double *          extra = p->extra;

  // TODO:  encapsulate electrons in EOS framework.
  double uei, uef;
#if ELECTRONS && EOS == EOS_TYPE_GAMMA
  uei = Ne[i][j][k] * Thetae_old[i][j][k] * ME * CL * CL / (game - 1.);
  uef = Ne[i][j][k] * Thetae * ME * CL * CL / (game - 1.);
#else
  uei = EOS_u_N_Theta(rho, Ne[i][j][k], Thetae_old[i][j][k], extra);
  uef = EOS_u_N_Theta(rho, Ne[i][j][k], Thetae, extra);
#endif

  struct of_microphysics micro;
  micro.Thetae = Thetae;
  micro.Ne     = Ne[i][j][k];
  micro.B      = Bmag[i][j][k];

  double J   = get_J(&m);
  double vol = ggeom[i][j][CENT].g * dx[1] * dx[2] * dx[3] * dt;

  // Loop over superphotons
  double udotG_abs   = 0.;
  double udotG_scatt = 0.;
  for (int n = 0; n < nthreads; n++) {
    for (int m = 0; m < Nsph_zone[i][j][k][n]; m++) {
      double udotk = HPL * nu[i][j][k][n][m] / (ME * CL * CL);

      // Absorption
      double alpha_inv_a =
          alpha_inv_abs(nu[i][j][k][n][m], type, &micro, theta[i][j][k][n][m]);
      double dtau_a =
          alpha_inv_a * L_unit * HPL / (ME * CL * CL) * dlam[i][j][k][n][m];
      double dw_a = w[i][j][k][n][m] * (1. - exp(-dtau_a));
      udotG_abs += kphys_to_num * dw_a * udotk / vol;

      // dEtau_abs += HPL*nu[i][j][k][n][m]*w[i][j][k][n][m]*(1. -
      // exp(-dtau_a));

      // Scattering (assuming h \nu << k_B T_e)
      double alpha_inv_s =
          alpha_inv_scatt(nu[i][j][k][n][m], type, interaction, &micro);
      double dtau_s =
          alpha_inv_s * L_unit * HPL / (ME * CL * CL) * dlam[i][j][k][n][m];
      double dw_s = w[i][j][k][n][m] * (1. - exp(-dtau_s));
      double amp =
          1. + 4. * Thetae - 2. * pow(Thetae, 3. / 2.) + 16. * pow(Thetae, 2.);
      udotG_scatt += kphys_to_num * (1. - amp) * dw_s * udotk / vol;
      // dEdtau_scatt -= HPL*nu[i][j][k][n][m]*;
    }
  }
  // dEdtau_abs *= Ucon[i][j][k][0]/(/* d3xi! */dt*T_unit);
  // dEdtau_scatt *= Ucon[i][j][k][0]/(dt*T_unit);
  udotG_abs *= U_unit / T_unit;
  udotG_scatt *= U_unit / T_unit;

  // if (udotG_abs != 0. || udotG_scatt != 0.) {
  /*if (i == 80 && j == 64) {
    printf("Thetae Ne Bmag = %e %e %e\n", Thetae, Ne[i][j][k], Bmag[i][j][k]);
    printf("%e %e %e %e\n", Ucon[i][j][k][0]*(uef - uei)/dt,
      J, udotG_abs, udotG_scatt);
  }*/

  // Solve entropy equation in cgs:
  //   d u_e / d \tau = -emission + absorption - upscattering
  double resid = Ucon[i][j][k][0] * (uef - uei) / (dt * T_unit);
  resid += J;
  resid += udotG_abs;
  resid += udotG_scatt;
  return resid;
}

void estimate_Thetae(
    grid_prim_type P, grid_eosvar_type extra, double t, double dt) {
  if (NTHREADS != nthreads) {
    fprintf(stderr, "NTHREADS = %i nthreads = %i! Exiting...\n", NTHREADS,
        nthreads);
    exit(-1);
  }
  struct of_microphysics micro;

// Count superphotons in each zone
#pragma omp parallel
  {
    int               n  = omp_get_thread_num();
    struct of_photon *ph = photon_lists[n];
    while (ph != NULL) {
      int    i, j, k;
      double X[NDIM], Kcov[NDIM], Kcon[NDIM];
      get_X_K_interp(ph, t, P, X, Kcov, Kcon);
      Xtoijk(X, &i, &j, &k);
      Nsph_zone[i][j][k][n]++;
      ph = ph->next;
    }
  } // omp parallel

// malloc required memory and store ucon for convenience
#pragma omp parallel for collapse(3)
  ZLOOP {
    double Bcon[NDIM];
    get_fluid_zone(i, j, k, P, extra, &micro, Ucon[i][j][k], Ucov[i][j][k],
        Bcon, Bcov[i][j][k]);
    Ne[i][j][k]         = micro.Ne;
    Thetae_old[i][j][k] = micro.Thetae;
    Bmag[i][j][k]       = micro.B;
    for (int n = 0; n < nthreads; n++) {
      w[i][j][k][n]     = safe_malloc(Nsph_zone[i][j][k][n] * sizeof(double));
      nu[i][j][k][n]    = safe_malloc(Nsph_zone[i][j][k][n] * sizeof(double));
      dlam[i][j][k][n]  = safe_malloc(Nsph_zone[i][j][k][n] * sizeof(double));
      theta[i][j][k][n] = safe_malloc(Nsph_zone[i][j][k][n] * sizeof(double));
      Nsph_cntd[i][j][k][n] = 0;
    }
  } // omp parallel

// Create per-zone lists of w, nu, dlam
#pragma omp parallel
  {
    int               n  = omp_get_thread_num();
    struct of_photon *ph = photon_lists[n];
    while (ph != NULL) {
      int    i, j, k;
      double X[NDIM], Kcov[NDIM], Kcon[NDIM];
      get_X_K_interp(ph, t, P, X, Kcov, Kcon);
      Xtoijk(X, &i, &j, &k);

      if (i < NG || i > NG + N1 - 1 || j < NG || j > NG + N2 - 1 || k < NG ||
          k > NG + N3 - 1) {
        printf("BAD PH????\n");
        printf("[%i] %i %i %i X[] = %e %e %e %e\n", n, i, j, k, X[0], X[1],
            X[2], X[3]);
        for (int mu = 0; mu < 3; mu++) {
          printf("X[%i][] = %e %e %e %e\n", mu, ph->X[mu][0], ph->X[mu][1],
              ph->X[mu][2], ph->X[mu][3]);
          printf("Kcov[%i][] = %e %e %e %e\n", mu, ph->Kcov[mu][0],
              ph->Kcov[mu][1], ph->Kcov[mu][2], ph->Kcov[mu][3]);
          printf("Kcon[%i][] = %e %e %e %e\n", mu, ph->Kcon[mu][0],
              ph->Kcon[mu][1], ph->Kcon[mu][2], ph->Kcon[mu][3]);
        }
        printf("origin %i %i %i %i\n", ph->origin[0], ph->origin[1],
            ph->origin[2], ph->origin[3]);
        printf("nscatt = %i\n", ph->nscatt);
        exit(-1);
      }

      double freq = 0.;
      for (int mu = 0; mu < NDIM; mu++) {
        freq -= Ucon[i][j][k][mu] * Kcov[mu];
      }
      freq *= ME * CL * CL / HPL;

      w[i][j][k][n][Nsph_cntd[i][j][k][n]]    = ph->w;
      nu[i][j][k][n][Nsph_cntd[i][j][k][n]]   = freq;
      dlam[i][j][k][n][Nsph_cntd[i][j][k][n]] = dt / Kcon[0];
      theta[i][j][k][n][Nsph_cntd[i][j][k][n]] =
          get_bk_angle(X, Kcon, Ucov[i][j][k], Bcov[i][j][k], Bmag[i][j][k]);

      Nsph_cntd[i][j][k][n]++;
      ph = ph->next;
    }
  } // omp parallel

// In each zone, rootfind to Thetae that matches sources. If failure, simply
// return current Thetae
#pragma omp parallel
  {
    const gsl_root_fsolver_type *T;
    gsl_root_fsolver *           s;
    gsl_function                 F;
    F.function = &dEdt;
    T          = gsl_root_fsolver_brent;
    s          = gsl_root_fsolver_alloc(T);

    static double extra[EOS_NUM_EXTRA];
#pragma omp threadprivate(extra)
#if EOS == EOS_TYPE_GAMMA
    EOS_ELOOP { extra[e] = 0.0; }
#endif

#pragma omp for collapse(3)
    ZLOOP {
#if EOS == EOS_TYPE_TABLE
      extra[EOS_YE] = P[i][j][k][YE];
#endif

      // fill parameters
      struct of_params params;
      params.i       = i;
      params.j       = j;
      params.k       = k;
      params.rho     = P[i][j][k][RHO];
      params.Ne      = Ne[i][j][k];
      params.Bmag    = Bmag[i][j][k];
      params.Thetaei = Thetae_old[i][j][k];
      params.Ucon0   = Ucon[i][j][k][0];
      params.dt      = dt;
      // hopefully this doesn't break openmp
      EOS_ELOOP { params.extra[e] = extra[e]; }

      F.params         = &params;
      double r         = 0.;
      double Thetae_lo = 0.5 * Thetae_old[i][j][k];
      double Thetae_hi = 2. * Thetae_old[i][j][k];

      // Test interval for sanity
      double rmin = dEdt(Thetae_lo, &params);
      double rmax = dEdt(Thetae_hi, &params);
      if (rmin * rmax > 0.) {
        fprintf(stderr, "[%i %i %i] Root not bracketed!\n", i, j, k);
        Thetae_est[i][j][k] = Thetae_old[i][j][k];
        continue;
      }

      gsl_root_fsolver_set(s, &F, Thetae_lo, Thetae_hi);

      int iter = 0, max_iter = 100;
      int status;
      do {
        status    = gsl_root_fsolver_iterate(s);
        r         = gsl_root_fsolver_root(s);
        Thetae_lo = gsl_root_fsolver_x_lower(s);
        Thetae_hi = gsl_root_fsolver_x_upper(s);
        status    = gsl_root_test_interval(Thetae_lo, Thetae_hi, 0, 0.001);
      } while (status == GSL_CONTINUE && iter < max_iter);

      if (status != GSL_SUCCESS) {
        Thetae_est[i][j][k] = Thetae_old[i][j][k];
      } else {
        Thetae_est[i][j][k] = r;
      }
    }

    gsl_root_fsolver_free(s);
  } // omp parallel

// Clean up mallocs and reset counters
#pragma omp parallel for collapse(3)
  ZLOOP {
    for (int n = 0; n < nthreads; n++) {
      free(w[i][j][k][n]);
      free(nu[i][j][k][n]);
      free(dlam[i][j][k][n]);
      free(theta[i][j][k][n]);
      Nsph_zone[i][j][k][n] = 0.;
    }
  } // omp parallel
}

#endif // ESTIMATE_THETAE
#endif // RADIATION