#include "def.h"
#include <math.h>
#include <stdio.h>

//f(tij|ui)
float fuij(float fpar[NFPAR],Pool pl, float ui){
      int i;
      float prod=1, ita;
      for(i=0; i<pl.nsub; i++){
        ita = exp(fpar[0]+fpar[1]*pl.xijk[i]+ui);
        prod = prod * (1/(1+ita));      
      }

      if(pl.tij == 1){
          return SE+(1-SE-SP)*prod;      
      }
      else{
          return 1 - SE - (1-SE-SP)*prod;
      }      
}

//replace random effect
void Replace(float fpar[NFPAR], Data*data, float rparm, Strf*rf, long*idum)
{
     float rnorm(long *idum);
     float runif(long *idum);
     int i,j;
     float prop;
     for(i=0;i<(*rf).nst;i++){
         //generate proposal random sample from proposal density
         prop=rnorm(idum)*sqrt(rparm);       
         
       
         //accept function
         float accept=1;
         for(j=0;j<(*data).st[i].npl;j++){               
            accept = accept * fuij(fpar,(*data).st[i].pls[j],prop)/fuij(fpar,(*data).st[i].pls[j],(*rf).uh[i]);          
         }     
          
          //move?
          if(accept > runif(idum)){
             (*rf).uh[i]=prop;              
          }
     }
     
}

//metroplis step:generate random draws from condistonal density
void Metroplis(float fpar[NFPAR], float rpar,Data*data, int nburn, int nskip, Reff*rf, Strf*u0, long*idum)
//rf:outside buffer to store random effect 
//uo:starting values for random effect
//note:
{
     void Replace(float fpar[NFPAR], Data*data, float rparm, Strf*rf, long*idum);
     
     int k,h,i,j;
    
     //burn in phase
     for(k=0;k<nburn;k++){
        Replace(fpar, data, rpar, u0, idum);     
     }
     
     //generation phase
     for(h=0;h<(*rf).M;h++){
         //skip?
         for(i=0;i<nskip;i++){
            Replace(fpar, data, rpar, u0, idum);     
         }     

        /* for(k=0;k<(*u0).nst;k++){
            printf("u0_%d=%f %f\n", k, (*u0).uh[k][0],(*u0).uh[k][1]);
         }*/
         //stor random draws
         (*rf).u[h].nst = (*u0).nst;
         for(j=0;j<(*rf).u[h].nst;j++){
             (*rf).u[h].uh[j]=(*u0).uh[j];
         }        
         
     }          
}

//first and second order derivative of log(fuij)
//w.r.t fixed effect parameter 
void score_betaij(float fpar[NFPAR],Pool*pl, float ui, float dfij[NFPAR],float ddfij[NFPAR][NFPAR])
{
     int k;
     float ita,prod=1,sij, weight, x;
     float dsij[NFPAR]={0,0};
     float sum1[NFPAR]={0,0};
     float sum2[NFPAR][NFPAR]={0,0,0,0};
     float ddsij[NFPAR][NFPAR]={0,0,0,0};
     
     //prod=prod(1+exp(*))
     for(k=0;k<(*pl).nsub;k++){
        x = (*pl).xijk[k];
        ita = exp(fpar[0]+fpar[1]*x+ui);
        prod = prod*(1+ita);
        
        sum1[0] += ita/(1+ita);
        sum1[1] += (ita/(1+ita))*x;                
        
        //for ddsij,only the upper triangular is necessary
        weight  = ita/((1+ita)*(1+ita));
        sum2[0][0] += weight;
        sum2[0][1] += weight*x;
        sum2[1][1] += weight*(x*x);
     }
     
     //sij
     sij = SE + (1 - SE - SP)/prod;
     //dsij
     dsij[0] = -(1 - SE - SP)/prod*sum1[0];
     dsij[1] = -(1 - SE - SP)/prod*sum1[1];
     
     //ddsij
     ddsij[0][0] = - (dsij[0]*sum1[0] + (sij - SE)*sum2[0][0]);
     ddsij[0][1] = - (dsij[0]*sum1[1] + (sij - SE)*sum2[0][1]);
     //ddsij[1][0] = - (dsij[1]*sum1[0] + (sij - SE)*sum2[0][1]);
     ddsij[1][1] = - (dsij[1]*sum1[1] + (sij - SE)*sum2[1][1]); 
     
     if((*pl).tij == 1){
         dfij[0] = dsij[0]/sij;
         dfij[1] = dsij[1]/sij;     
         
         //second order derivative
         weight  = sij*sij;
         ddfij[0][0] = (sij*ddsij[0][0] - dsij[0]*dsij[0])/weight;
         ddfij[0][1] = (sij*ddsij[0][1] - dsij[0]*dsij[1])/weight;
         ddfij[1][0] = ddfij[0][1];
         ddfij[1][1] = (sij*ddsij[1][1] - dsij[1]*dsij[1])/weight;
     }
     else{
          //first order derivative          
         dfij[0] = -dsij[0]/(1-sij);
         dfij[1] = -dsij[1]/(1-sij);
         
         //second order derivative
         weight  = (1-sij)*(1-sij);
         ddfij[0][0] = ((sij-1)*ddsij[0][0] - dsij[0]*dsij[0])/weight;
         ddfij[0][1] = ((sij-1)*ddsij[0][1] - dsij[0]*dsij[1])/weight;
         ddfij[1][0] = ddfij[0][1];
         ddfij[1][1] = ((sij-1)*ddsij[1][1] - dsij[1]*dsij[1])/weight;
     }
}

//update parameters
//one step Newton-Raphson for fixed effect parameters
//closed form of MLE for random effect parameters
void parm_update(Reff*rf,
                 Data*data,
                 float oldfpar[NFPAR],
                 float delta_fpar[NFPAR],
                 float *newrpar){
      float square(float x);
      float score[NFPAR]={0,0}, tempscore[NFPAR]={0,0};
      float hess[NFPAR][NFPAR]={0,0,0,0},temphess[NFPAR][NFPAR]={0,0,0,0};
      //parameter for random effects
      float var=0;
      
      int h,i,j;
      
      //use monte-carlo method to estimate score and hessian 
      for(h = 0;h<(*rf).M;h++){
            for(i = 0; i<(*data).nst; i++){
               for(j = 0; j<(*data).st[i].npl;j++){
                     score_betaij(oldfpar,
                                  &((*data).st[i].pls[j]),                                        
                                  (*rf).u[h].uh[i],
                                  tempscore,
                                  temphess);
                     score[0] += tempscore[0];
                     score[1] += tempscore[1];
                     hess[0][0]+=temphess[0][0];
                     hess[0][1]+=temphess[0][1];
                     hess[1][0]+=temphess[1][0];
                     hess[1][1]+=temphess[1][1];               
               }            
               
            var  += square((*rf).u[h].uh[i]);
            }            
      }
      
      //approximated score and hessian
      score[0] /= (*rf).M;
      score[1] /= (*rf).M;
      hess[0][0] /= -(*rf).M;
      hess[0][1] /= -(*rf).M;
      hess[1][0] /= -(*rf).M;
      hess[1][1] /= -(*rf).M;
      
      Invert22(hess);
      
      //updates for beta
      delta_fpar[0] = hess[0][0]*score[0] + hess[0][1]*score[1];
      delta_fpar[1] = hess[1][0]*score[0] + hess[1][1]*score[1];
      
      //new variance component
       int weight = (*rf).M*((*data).nst);
       *newrpar = var/weight;
}