
DATA_SECTION
  int n
 !! n=50000;
  vector obs(1,n)
  vector index(1,n)
 !! obs.fill_seqadd(1,.001);
 !! index.fill_seqadd(2,.001);
 int gpu_flag;
 !! gpu_flag=0;
 !! cout << " enter 1 for GPU, 2 for CPU " << endl;
 !! cin >> gpu_flag;

PARAMETER_SECTION
  init_number log_mu
  init_number log_alpha
  objective_function_value f
PROCEDURE_SECTION
  dvar_vector mu=exp(log_mu)+index;
  dvar_vector tau=exp(log_alpha)+index;
  int nreps=5;
  if (gpu_flag==1)
  {
    for (int i=1;i<=nreps;i++)
      f-=sum(log_negbinomial_density_gpu(obs,mu,tau));
  }
  else
  {
    for (int i=1;i<=nreps;i++)
      f-=sum(my_log_negbinomial_density(obs,mu,tau));
  }
GLOBALS_SECTION

  #include <admodel.h>
  #include <stdio.h>
  
  #include <cutil_inline.h>
  #include <shrUtils.h>
  #include <shrQATest.h>
  
  #include <stdlib.h>
  
  #include <vector>
  #include <algorithm>
  

 extern "C"
 {
 
 double mygammln(double xx)
 {
   double x,tmp,ser;
   static double cof[6]={76.18009173,-86.50532033,24.01409822,
     -1.231739516,0.120858003e-2,-0.536382e-5};
   int j;
   x=xx-1.0;
   tmp=x+5.5;
   tmp -= (x+0.5)*log(tmp);
   ser=1.0;
   for (j=0;j<=5;j++) 
   {
     x += 1.0;
     ser += cof[j]/x;
   }
   return -tmp+log(2.50662827465*ser);
 }
 __device__ double mygammln_kernel(double xx)
 {
   double x,tmp,ser;
   const double cof[6]={76.18009173,-86.50532033,24.01409822,
     -1.231739516,0.120858003e-2,-0.536382e-5};
   int j;
   x=xx-1.0;
   tmp=x+5.5;
   tmp -= (x+0.5)*log(tmp);
   ser=1.0;
   for (j=0;j<=5;j++) 
   {
     x += 1.0;
     ser += cof[j]/x;
   }
   return -tmp+log(2.50662827465*ser);
 }
 double mylog_negbinomial_density(double x,double mu,double tau)
 {
   if (tau<=1.0)
   {
     printf("tau <=1 in log_negbinomial_density \n");
     exit(1);
   }
   double r=mu/(1.e-120+(tau-1.0));
   return mygammln(x+r)-gammln(r)-gammln(x+1)
     +r*log(r)+x*log(mu)-(r+x)*log(r+mu);
 }
 

  struct df11variable
  {
    double v[2]; // v[0] holds value, v[1] holds derivative
  };
  
  struct df11variable df11gammln(struct df11variable xx)
  {
    struct  df11variable x,tmp,ser,z;
    static double cof[6]={76.18009173,-86.50532033,24.01409822,
      -1.231739516,0.120858003e-2,-0.536382e-5};
  
    xx.v[1]=1.0;   // initialize 
    int j;
    x.v[0]=xx.v[0]-1.0;
    x.v[1]=xx.v[1];
  
    tmp.v[1]=x.v[1];
    tmp.v[0]=x.v[0]+5.5;
    
  
    tmp.v[1] -= ( x.v[1]*log(tmp.v[0]) + (x.v[0]+0.5)/tmp.v[0]*tmp.v[1] );
    tmp.v[0] -= (x.v[0]+0.5)*log(tmp.v[0]);
  
  
    ser.v[1]=0.0;
    ser.v[0]=1.0;
  
    for (j=0;j<=5;j++) 
    {
      x.v[0] += 1.0;
      ser.v[1] -= cof[j]/(x.v[0]*x.v[0])*x.v[1];
      ser.v[0] += cof[j]/x.v[0];
    }
   
  
    z.v[0]= -tmp.v[0] + log(2.50662827465*ser.v[0]);
    z.v[1]= -tmp.v[1]+ 1.0/ser.v[0]*ser.v[1];
  
    return z;
  }

  __device__ struct df11variable df11gammln_kernel(struct df11variable xx)
  {
    struct  df11variable x,tmp,ser,z;
    //static double cof[6]={76.18009173,-86.50532033,24.01409822,
    //  -1.231739516,0.120858003e-2,-0.536382e-5};
    const double cof[6]={76.18009173,-86.50532033,24.01409822,
      -1.231739516,0.120858003e-2,-0.536382e-5};
  
    xx.v[1]=1.0;   // initialize 
    int j;
    x.v[0]=xx.v[0]-1.0;
    x.v[1]=xx.v[1];
  
    tmp.v[1]=x.v[1];
    tmp.v[0]=x.v[0]+5.5;
    
  
    tmp.v[1] -= ( x.v[1]*log(tmp.v[0]) + (x.v[0]+0.5)/tmp.v[0]*tmp.v[1] );
    tmp.v[0] -= (x.v[0]+0.5)*log(tmp.v[0]);
  
  
    ser.v[1]=0.0;
    ser.v[0]=1.0;
  
    for (j=0;j<=5;j++) 
    {
      x.v[0] += 1.0;
      ser.v[1] -= cof[j]/(x.v[0]*x.v[0])*x.v[1];
      ser.v[0] += cof[j]/x.v[0];
    }
   
  
    z.v[0]= -tmp.v[0] + log(2.50662827465*ser.v[0]);
    z.v[1]= -tmp.v[1]+ 1.0/ser.v[0]*ser.v[1];
  
    return z;
  }
  struct df12variable
  {
    double v[3]; // v[0] holds value,
                //v[1] holds derivative wrt first variable
                //v[2] holds derivative wrt second variable
  };
  
  
  
  struct df12variable df12log_negbinomial_density(double x,struct df12variable mu,
    struct df12variable tau)
  {
    struct df12variable r;
    struct df12variable z;
    struct df11variable r1;
    struct df11variable r2;
    struct df11variable z1;
    struct df11variable z2;
    mu.v[1]=1.0;
    mu.v[2]=0.0;
    tau.v[1]=0.0;
    tau.v[2]=1.0;
    if (tau.v[0]<=1.0)
    {
      printf("tau <=1 in log_negbinomial_density \n");
      exit(1);
    }
    r.v[0]=mu.v[0]/(1.e-120+(tau.v[0]-1.0));
    r.v[1]=mu.v[1]/(1.e-120+(tau.v[0]-1.0)); //der wrt mu
    r.v[2]=-mu.v[0]/((1.e-120+(tau.v[0]-1.0))*(1.e-120+(tau.v[0]-1.0)))*
              tau.v[2]; //der wrt tau
    // return mygammln(x+r)-gammln(r)-gammln(x+1)
    //  +r*log(r)+x*log(mu)-(r+x)*log(r+mu);
    r1.v[0]=r.v[0]+x;
    z1=df11gammln(r1); 
    r2.v[0]=r.v[0];
    z2=df11gammln(r2); 
    z.v[0]=0.0;
    z.v[1]=0.0;
    z.v[2]=0.0;
    // do it step by step
    z.v[0]+=z1.v[0]; // this is z+=gammln(x+r)
    z.v[1]+=z1.v[1]*r.v[1]; // chain rule for der wrt mu
    z.v[2]+=z1.v[1]*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]-=z2.v[0]; // this is z-=gammln(r)
    z.v[1]-=z2.v[1]*r.v[1]; // chain rule for der wrt mu
    z.v[2]-=z2.v[1]*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]-=mygammln(x+1.0); // constant term so no ders
  
    z.v[0]+= r.v[0]*log(r.v[0]); // this is z+=r*log(r)
    z.v[1]+= (log(r.v[0])+1.0)*r.v[1]; // chain rule for der wrt mu
    z.v[2]+= (log(r.v[0])+1.0)*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]+= x*log(mu.v[0]);
    z.v[1]+= x/mu.v[0]*mu.v[1];; // chain rule for der wrt mu
                               // note that mu.v[2]=0.0
    z.v[0]-=(r.v[0]+x)*log(r.v[0]+mu.v[0]);
    z.v[1]-=log(r.v[0]+mu.v[0])*r.v[1];
    z.v[1]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*r.v[1];
    z.v[1]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*mu.v[1];
    z.v[2]-=log(r.v[0]+mu.v[0])*r.v[2];
    z.v[2]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*r.v[2];
    return z;
  }

  __device__ void df12log_negbinomial_density_kernel(double x,double _mu,
    double _tau,double *dfx,double * dfy,double *tmp)
  {
    struct df12variable mu;
    struct df12variable tau;
    struct df12variable r;
    struct df12variable z;
    struct df11variable r1;
    struct df11variable r2;
    struct df11variable z1;
    struct df11variable z2;
    mu.v[0]=_mu;
    mu.v[1]=1.0;
    mu.v[2]=0.0;
    tau.v[0]=_tau;
    tau.v[1]=0.0;
    tau.v[2]=1.0;
   /*
    if (tau.v[0]<=1.0)
    {
      printf("tau <=1 in log_negbinomial_density \n");
      exit(1);
    }
   */
    r.v[0]=mu.v[0]/(1.e-120+(tau.v[0]-1.0));
    r.v[1]=mu.v[1]/(1.e-120+(tau.v[0]-1.0)); //der wrt mu
    r.v[2]=-mu.v[0]/((1.e-120+(tau.v[0]-1.0))*(1.e-120+(tau.v[0]-1.0)))*
              tau.v[2]; //der wrt tau
    // return mygammln(x+r)-gammln(r)-gammln(x+1)
    //  +r*log(r)+x*log(mu)-(r+x)*log(r+mu);
    r1.v[0]=r.v[0]+x;
    z1=df11gammln_kernel(r1); 
    r2.v[0]=r.v[0];
    z2=df11gammln_kernel(r2); 
    z.v[0]=0.0;
    z.v[1]=0.0;
    z.v[2]=0.0;
    // do it step by step
    z.v[0]+=z1.v[0]; // this is z+=gammln(x+r)
    z.v[1]+=z1.v[1]*r.v[1]; // chain rule for der wrt mu
    z.v[2]+=z1.v[1]*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]-=z2.v[0]; // this is z-=gammln(r)
    z.v[1]-=z2.v[1]*r.v[1]; // chain rule for der wrt mu
    z.v[2]-=z2.v[1]*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]-=mygammln_kernel(x+1.0); // constant term so no ders
  
    z.v[0]+= r.v[0]*log(r.v[0]); // this is z+=r*log(r)
    z.v[1]+= (log(r.v[0])+1.0)*r.v[1]; // chain rule for der wrt mu
    z.v[2]+= (log(r.v[0])+1.0)*r.v[2]; // chain rule for der wrt tau
  
    z.v[0]+= x*log(mu.v[0]);
    z.v[1]+= x/mu.v[0]*mu.v[1];; // chain rule for der wrt mu
                               // note that mu.v[2]=0.0
    z.v[0]-=(r.v[0]+x)*log(r.v[0]+mu.v[0]);
    z.v[1]-=log(r.v[0]+mu.v[0])*r.v[1];
    z.v[1]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*r.v[1];
    z.v[1]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*mu.v[1];
    z.v[2]-=log(r.v[0]+mu.v[0])*r.v[2];
    z.v[2]-=(r.v[0]+x)/(r.v[0]+mu.v[0])*r.v[2];
    //return z;
    *tmp=z.v[0];
    *dfx=z.v[1];
    *dfy=z.v[2];
  }
  
  
  } // extern "C"


  void DF_dvdvnegbin(void);
  
  dvar_vector my_log_negbinomial_density(dvector& x,dvar_vector& mu,dvar_vector& tau){
    int mmin=x.indexmin();
    int mmax=x.indexmax();
    if (mmin !=mu.indexmin() || mmin != tau.indexmin() ||
        mmax !=tau.indexmax() || mmax != tau.indexmax())
    {
      cerr << "shape error in log_negbinomial_density" << endl;
      ad_exit(1);
    }
    dvar_vector tmp(mmin,mmax);
    dvector dfx(mmin,mmax);
    dvector dfy(mmin,mmax);
    for (int i=mmin;i<=mmax;i++)
    {
      struct df12variable z,xmu,xtau;
      xmu.v[0]=value(mu(i));
      xtau.v[0]=value(tau(i));
      z=df12log_negbinomial_density(x(i),xmu,xtau);
      value(tmp(i))=z.v[0];
      dfx(i)=z.v[1];
      dfy(i)=z.v[2];
    }
    save_identifier_string("tg");
    tmp.save_dvar_vector_position();
    mu.save_dvar_vector_position();
    tau.save_dvar_vector_position();
    save_identifier_string("ez");
    dfx.save_dvector_value();
    dfx.save_dvector_position();
    dfy.save_dvector_value();
    dfy.save_dvector_position();
    save_identifier_string("ey");
    gradient_structure::GRAD_STACK1->
              set_gradient_stack(DF_dvdvnegbin);
    return tmp; 
  }

  __global__ void log_negbinomial_density_kernel(double *x ,double * mu,
    double * tau, double * dfx, double * dfy, double * tmp,int n)
  {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < n)
    {
      df12log_negbinomial_density_kernel(x[i],mu[i],tau[i],&(dfx[i]),
        &(dfy[i]),&(tmp[i]));
    }
  }

      double * d_x=0;
      double * d_mu=0;
      double * d_tau=0;
      double * d_dfx=0;
      double * d_dfy=0;
      double * d_tmp=0;

   
    dvar_vector log_negbinomial_density_gpu(dvector& x,dvar_vector& mu,dvar_vector& tau){
      int mmin=x.indexmin();
      int mmax=x.indexmax();
      if (mmin !=mu.indexmin() || mmin != tau.indexmin() ||
          mmax !=tau.indexmax() || mmax != tau.indexmax())
      {
        cerr << "shape error in log_negbinomial_density" << endl;
        ad_exit(1);
      }
      dvar_vector tmp(mmin,mmax);
      dvector dfx(mmin,mmax);
      dvector dfy(mmin,mmax);
      dvector cmu=value(mu);
      dvector ctau=value(tau);
      int N = mmax-mmin+1;
      size_t size = N * sizeof(double);

      double * h_x=&(x(mmin));
      double * h_mu=&(cmu(mmin));
      double * h_tau=&(ctau(mmin));
      double * h_dfx=&(dfx(mmin));
      double * h_dfy=&(dfy(mmin));
      double * h_tmp=&(value(tmp(mmin)));

     /*
      double * d_x;
      double * d_mu;
      double * d_tau;
      double * d_dfx;
      double * d_dfy;
      double * d_tmp;
     */

    // Allocate vectors in device memory
    if (!d_x)
     cutilSafeCall( cudaMalloc((void**)&d_x, size) );
    if (!d_mu)
      cutilSafeCall( cudaMalloc((void**)&d_mu, size) );
    if (!d_tau)
      cutilSafeCall( cudaMalloc((void**)&d_tau, size) );
    if (!d_dfx)
      cutilSafeCall( cudaMalloc((void**)&d_dfx, size) );
    if (!d_dfy)
      cutilSafeCall( cudaMalloc((void**)&d_dfy, size) );
    if (!d_tmp)
      cutilSafeCall( cudaMalloc((void**)&d_tmp, size) );

   // Copy vectors from host memory to device memory
    cutilSafeCall( cudaMemcpy(d_x, h_x, size, cudaMemcpyHostToDevice) );
    cutilSafeCall( cudaMemcpy(d_mu, h_mu, size, cudaMemcpyHostToDevice) );
    cutilSafeCall( cudaMemcpy(d_tau, h_tau, size, cudaMemcpyHostToDevice) );


  // Invoke kernel
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    log_negbinomial_density_kernel<<<blocksPerGrid, threadsPerBlock>>>(d_x, d_mu, d_tau,d_dfx,d_dfy,d_tmp,N);
    cutilCheckMsg("kernel launch failure");
  #ifdef _DEBUG
    cutilSafeCall( cutilDeviceSynchronize() );
  #endif

    cutilSafeCall( cudaMemcpy(h_tmp, d_tmp, size, cudaMemcpyDeviceToHost) );
    cutilSafeCall( cudaMemcpy(h_dfx, d_dfx, size, cudaMemcpyDeviceToHost) );
    cutilSafeCall( cudaMemcpy(h_dfy, d_dfy, size, cudaMemcpyDeviceToHost) );

    save_identifier_string("tg");
    tmp.save_dvar_vector_position();
    mu.save_dvar_vector_position();
    tau.save_dvar_vector_position();
    save_identifier_string("ez");
    dfx.save_dvector_value();
    dfx.save_dvector_position();
    dfy.save_dvector_value();
    dfy.save_dvector_position();
    save_identifier_string("ey");
    gradient_structure::GRAD_STACK1->
              set_gradient_stack(DF_dvdvnegbin);
    return tmp; 
  }
  
  
  void DF_dvdvnegbin(void)
  {
    // int ierr=fsetpos(gradient_structure::get_fp(),&filepos);
    verify_identifier_string("ey");
    dvector_position dfy_pos=restore_dvector_position();
    dvector dfy=restore_dvector_value(dfy_pos);
    dvector_position dfx_pos=restore_dvector_position();
    dvector dfx=restore_dvector_value(dfx_pos);
    verify_identifier_string("ez");
  
    dvar_vector_position tau_pos=restore_dvar_vector_position();
    dvar_vector_position mu_pos=restore_dvar_vector_position();
    dvar_vector_position tmp_pos=restore_dvar_vector_position();
  
    verify_identifier_string("tg");
    dvector dfvtmp=restore_dvar_vector_derivatives(tmp_pos);
  
    dvector dfmu(dfvtmp.indexmin(),dfvtmp.indexmax());
    dvector dftau(dfvtmp.indexmin(),dfvtmp.indexmax());
  
    for (int i=dfvtmp.indexmin();i<=dfvtmp.indexmax();i++)
    {
      dfmu(i)=dfx(i)*dfvtmp(i);
      dftau(i)=dfy(i)*dfvtmp(i);
    }
    dfmu.save_dvector_derivatives(mu_pos);
    dftau.save_dvector_derivatives(tau_pos);
  
   
  }
  
