#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <cmath>
#include <cstdlib>
#include <TMath.h>
#include <TF1.h>
#include <TRandom3.h>
#include <Math/ProbFuncMathCore.h>
#include <Math/QuantFuncMathCore.h>
#include "Prof.h"
#include "fitPar.h"
#include <Math/Random.h>
#include <Math/GSLRndmEngines.h>

using namespace std;

Prof::Prof (int N, int M) {

  // set m_N and m_npar, create minuit, everything else initialized to zero.

  m_N = N;
  m_M = M;
  m_npar = 2*m_N + m_M;

  m_yVec.clear();
  m_sigma_y_Vec.clear();
  m_rVec.clear();
  m_uVec.clear();
  m_sVec.clear();
  for (int i=0; i<N; i++){
    m_yVec.push_back(0.);
    m_sigma_y_Vec.push_back(0.);
    m_sVec.push_back(0.);
    m_uVec.push_back(0.);
    m_rVec.push_back(0.);
  }

  m_parVec.clear();
  m_sigParVec.clear();
  for (int i=0; i<m_npar; i++){
    m_parVec.push_back(0.);
    m_sigParVec.push_back(0.);
  }

  m_minuit = new TMinuit(m_npar);
  m_ran = 0;
  m_rgsl = 0;

}


Prof::Prof (string inputFile) {

  // Constructor assumes use wants an average (M=1)

  ifstream f;
  f.open(inputFile.c_str());
  if ( f.fail() ){
    cout << "Sorry, couldn't open file" << endl;
    exit(1);
  }

  istringstream instream;       // Declare an input string stream
  string line;
  int lineNum = 0;
  while ( getline(f, line) ) {
    bool useLine = !( line.substr(0,1) == "#" || line.substr(0,1) == "!");
    if ( useLine ) {
      instream.clear();          // Reset from possible previous errors
      instream.str(line);        // Use line as source of input
      lineNum++;
      double y, sigma_y, s, r, u;
      instream >> y >> sigma_y >> s >> r;
      u = 0;         // nominal values zero
      m_yVec.push_back(y);  
      m_sigma_y_Vec.push_back(sigma_y);
      m_sVec.push_back(s);
      m_rVec.push_back(r);
      m_uVec.push_back(u);
    }
  }
  f.close();

  m_N = m_yVec.size();
  m_M = 1;                // for average, only one param. of interest
  m_npar = 2*m_N + m_M;
  m_minuit = new TMinuit(m_npar);

  m_parVec.clear();
  m_sigParVec.clear();
  for (int i=0; i<m_npar; i++){
    m_parVec.push_back(0.);
    m_sigParVec.push_back(0.);
  }


}

vector<double> Prof::mu(vector<double> theta){

  vector<double> muVec(m_N);
  int N = this->N();
  int M = theta.size();      // could differ from nominal m_M
  if ( M == 0 || M == 1 ) {
    for (int i=0; i<N; i++){
      muVec[i] = theta[0];
    }
  }
  else if (M == m_N) {
    muVec = theta;          // all components fitted
  }
  return muVec;

}

double Prof::lnL(vector<double> parVec, int M){

  // y_i ~ Gauss (mu_i + b_i, sigma_y_i)
  // u_i ~ Gauss (b_i, sigma_u_i)
  // v_i ~ gamma (alpha_i, beta_i), 
  // alpha_i = 1./4r_i^2, beta_i = alpha_i / sigma_u_i^2
  // parVec[i] = mu_i, i = 0,...,M-1  parameters of interest
  // parVec[i+M] = b_i, i  0,...,M-1  bias parameters
  // parVec[N+i+1] = eta_i, where eta_i = 1/sigma_u_i^2

  int N = m_N;
  vector<double> theta(M);
  for (int i=0; i<M; i++){
    theta[i] = parVec[i];
  }
  vector<double> muVec = this->mu(theta);
  vector<double> bVec(N);
  vector<double> sigma_u_Vec(N);
  for (int i=0; i<N; i++){
    bVec[i] = parVec[M+i];
    double lambda = parVec[M+N+i];
    sigma_u_Vec[i] = exp(lambda);
  }

  const double epsilon = 1.e-9;
  double chi2 = 0.;
  for (int i=0; i<N; i++){
    double alpha;
    if ( m_rVec[i] <= 0 ){
      alpha = epsilon;
    }
    else {
      alpha = 1./ (4.*m_rVec[i]*m_rVec[i]);
    }
    double beta = alpha / (sigma_u_Vec[i]*sigma_u_Vec[i]);
    double q = m_yVec[i] - muVec[i] - bVec[i];
    double v = m_uVec[i] - bVec[i];  // don't confuse with su2
    double w = m_sigma_y_Vec[i]*m_sigma_y_Vec[i];  
    double z = sigma_u_Vec[i]*sigma_u_Vec[i];
    double s = m_sVec[i];         // "measurement" for sigma_u
    double su2 = s*s;                // "measurement" of sigma_u2
    // note alpha is indep of parameters but beta depends on sigma_u
    double g = alpha*log(beta) + (alpha-1)*log(su2) - beta*su2;
    chi2 += q*q/w  +  v*v/z  + 2.*log(sigma_u_Vec[i]) - 2.*g;
  }

  double logL = -0.5*chi2;
  return logL;

}

int Prof::fit(int M){
  
  int N = this->N();
  this->setM(M);                   // sets m_npar
  vector<bool> freePar(m_npar);
  for (int i=0; i<m_npar; i++){
    freePar[i] = true;
  }

  int status = fitPar(this, freePar);
  return status;

}

void Prof::getStartPar(){

  int N = this->N();
  int M = this->M();
  int npar = this->npar();
  m_startPar.resize(npar);
  m_stepSize.resize(npar);

  // use ordinary LS for start value of mu  (for M = N use y_i)

  double wSum = 0.;
  vector<double> w(N);
  for (int i=0; i<N; i++){
    double sy = this->sigma_y(i);
    double su = this->s(i);
    w[i] = 1./(sy*sy + su*su);
    wSum += w[i];
  }
  double sigmaMuHatLS = 1./sqrt(wSum);

  double muHatLS = 0.;
  for (int i=0; i<N; i++){
    w[i] /= wSum;
    muHatLS += w[i] * this->y(i);
  }

  for (int i=0; i<M; i++){
    if ( M != N ) {
      m_startPar[i] = muHatLS;
    }
    else {
      m_startPar[i] = this->y(i);
    }
  }

  // use 0 for b_i, s_i for sigma_ui
  for (int i=0; i<N; i++){
    m_startPar[M+i] = 0.;
    m_startPar[M+N+i] = log( this->s(i) );
  }

  m_stepSize[0] = 0.1 * sigmaMuHatLS;
  for (int i=0; i<N; i++){
    m_stepSize[M+i] = 0.1 * this->s(i);      // for b_i
    m_stepSize[M+N+i] = 0.1;                 // for log sigma_u_i
  }

}

double Prof::q(){
  int N = this->N();
  int M = this->M();
  int status = this->fit(M);            // fills m_parVec
  double chi2_numer = -2.*this->lnL(m_parVec,M);

  // for chi2_denom, equivalent to fitting all mui
  double chi2_denom = 0.;
  for (int i=0; i<N; i++){
    double r2 = m_rVec[i]*m_rVec[i];
    double alpha = 1./(4.*r2);
    double sigma_u = exp(m_parVec[M+N+i]);          
    double beta = alpha / (sigma_u*sigma_u);
    double v = m_sVec[i]*m_sVec[i];
    chi2_denom += 2.*log(sigma_u) - 2.*alpha*log(beta) 
                - 2.*(alpha-1)*log(v) + 2.*beta*v;
  }
  return chi2_numer - chi2_denom;
}

double Prof::q(int M, int N){

  // likelihood ratio statistic with M fitted parameters in
  // numerator and N in denominator

  Prof modelM = *this;       // copies data
  modelM.setM(M);
  modelM.getStartPar();
  modelM.fit(M);
  vector<double> fitParM = modelM.parVec();

  int npm = modelM.npar();
  //   for (int i=0; i<npm; i++){
  //   cout << "i, fitParM = " << i << "  " << fitParM[i] << endl;
  // }
  // for (int i=0; i<N; i++){
  //   cout << "i, y, u, r = " << i << "  " <<
  //     modelM.y(i) << "  " << modelM.u(i) << "  " << modelM.r(i) << endl;
  // }

  double lnLM = modelM.lnL(fitParM, M);
  // cout << "lnLM = " << lnLM << endl;

  Prof modelN = *this;
  modelN.setM(N);
  modelN.getStartPar();
  modelN.fit(N);
  vector<double> fitParN = modelN.parVec();

  int npn = modelN.npar();
  // for (int i=0; i<npn; i++){
  //   cout << "i, fitParN = " << i << "  " << fitParN[i] << endl;
  // }
  // for (int i=0; i<N; i++){
  //   cout << "i, y, u, r = " << i << "  " <<
  //     modelN.y(i) << "  " << modelN.u(i) << "  " << modelN.r(i) << endl;
  // }

  double lnLN = modelN.lnL(fitParN, N);
  // cout << "lnLN = " << lnLN << endl;

  double qVal = -2. * (lnLM - lnLN);
  return qVal;

}

void Prof::generateData(int N, int M, vector<double> parVec,
     TRandom3* ran, ROOT::Math::Random<ROOT::Math::GSLRngMT>* rgsl) {

  m_yVec.clear();
  m_uVec.clear();
  m_sVec.clear();
  m_N = N;
  m_M = M;
  vector<double> theta(M);
  for (int i=0; i<M; i++){
    theta[i] = parVec[i];
  }
  vector<double> muVec = this->mu(theta);
  for (int i=0; i<N; i++){
    double sigma_y = m_sigma_y_Vec[i];
    double mu = muVec[i];
    double b = parVec[M+i];
    double lambda = parVec[M+N+i];
    double sigma_u = exp(lambda);
    // cout << "i, mu, b, sigma_u, r: " << i << "  " 
    //	 << mu << "  " << b << "  " << sigma_u << "  " << m_rVec[i] << endl;
    double r2 = m_rVec[i]*m_rVec[i];
    double alpha = 1./r2;
    double beta = alpha / sigma_u;
    double y = ran->Gaus(mu + b, sigma_y);
    double u = ran->Gaus(b, sigma_u);
    // cout << "alpha, beta = " << alpha << "  " << beta << endl;
    double s = rgsl->Gamma(alpha, 1./beta);  // gsl's convention for beta
    // cout << "i, y, u, s: " << i << "  " << y << "  " << u << "  " << s << endl;
    m_yVec.push_back(y);
    m_uVec.push_back(u);
    m_sVec.push_back(s);
  }

}