#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <cmath>
#include <cstdlib>

#include <TMath.h>
#include <TF1.h>
#include <TRandom3.h>
#include <TMatrixD.h>
#include <TVectorD.h>
#include <TMatrixDSym.h>
#include "TDecompSVD.h"
#include <Math/ProbFuncMathCore.h>
#include <Math/QuantFuncMathCore.h>
#include <Math/Random.h>
#include <Math/GSLRndmEngines.h>
#include <gsl/gsl_poly.h>
#include "Math/WrappedTF1.h"
#include "Math/BrentRootFinder.h"

#include "Eoe.h"
#include "fitPar.h"

using namespace std;

// Constructor

Eoe::Eoe (string inputFileName, TRandom3* ran, 
	  ROOT::Math::Random<ROOT::Math::GSLRngMT>* rgsl) {

  // read in steering file with data
  // line 1:  L -- number of measurements
  // lines 2:L+1 -- measured values y_i and stat. errors sigma_y_i
  // line L+2 -- N, number of nuisance params (=number of control meas)
  // lines L+3:L+N+2 -- control meas u_i and estimate of its std. dev. s_i
  // lines L+N+3:2L+N+2 -- matrix R_{ij), one row per line, defined
  // by E[y_i] = phi_i(mu) + \sum_{j} R_{ij} theta_j

  ifstream f;
  f.open(inputFileName.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;

      if ( lineNum == 1 ) {
        instream >> m_L;             // number of main measurements
        m_R.resize(m_L);             // set number of rows to m_L
        m_cov_bHat.resize(m_L);
        for (int i=0; i<m_L; i++){
          m_cov_bHat[i].resize(m_L);   // cov matrix of bias est is m_L by m_L
        }
      }

      if ( lineNum >= 2 && lineNum <= m_L+1 ){
        instream >> y >> sigma_y;
        m_yVec.push_back(y);  
        m_sigma_y_Vec.push_back(sigma_y);
      }

      if ( lineNum == m_L+2 ) {
        instream >> m_N;             // number nuisance par / control meas.
        for (int i = 0; i<m_L; i++){
          m_R[i].resize(m_N);        // set each row of m_R to have m_N columns
        }
      }

      if ( lineNum >= m_L+3 && lineNum <= m_L+m_N+2 ) {
        instream >> u >> s >> r;
        m_sVec.push_back(s);
        double v = s*s;
        m_vVec.push_back(v);
        m_rVec.push_back(r);
        m_uVec.push_back(u);
      }
  
      if ( lineNum >= m_L+m_N+3 && lineNum <= 2*m_L+m_N+2) {
        int i = lineNum - m_L - m_N - 3;    //   start at row 0
        for (int j=0; j<m_N; j++){
	  instream >> m_R[i][j];
	}
      }

    }
  }
  f.close();

  // initialize stuff

  m_M = 1;                 // for now support one parameter of interest
  m_npar = m_N + m_M;      // for default fit
  m_up = 1.;               // for MINOS

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

  // random generators (ran for uniform, rgsl for gamma)

  m_ran = ran;
  m_rgsl = rgsl;

  // initial (i.e. prefit) estimates of bias

  m_bHatVec.clear();
  for (int i=0; i<m_L; i++){
    double bHat = 0.;
    for (int j=0; j<m_N; j++){
      bHat += m_R[i][j]*m_uVec[j];
    }
    m_bHatVec.push_back(bHat);
    for (int j=0; j<m_L; j++){
      double cov_bHat = 0.;
      for (int k=0; k<m_N; k++){
        cov_bHat += m_R[i][k] * m_R[j][k] * m_vVec[k];
      }
      m_cov_bHat[i][j] = cov_bHat;
    }
  }

}

double Eoe::phi(int i, vector<double> muVec){    
  double mu = muVec[0];
  return mu;             //  phi_i(mu) = E[y_i] (nominal model, no bias)
}

double Eoe::phi(int i, double mu){
  vector<double> muVec(1);
  muVec[0] = mu;
  return this->phi(i, muVec);
}

double Eoe::lnL(vector<double> parVec){

  // y_i ~ Gauss (mu_i + (R theta)_i, sigma_y_i), i = 1,..., L
  // u_i ~ Gauss (theta_i, sigma_u_i), i = 1,..., N
  // v_i ~ gamma (alpha_i, beta_i), i = 1,..., N
  // alpha_i = 1./4r_i^2, beta_i = alpha_i / sigma_u_i^2
  // parVec[0:m_M-1] =  parameters of interest (e.g. parVec[0] = mu)
  // parVec[m_M:m_M+m_N-1] = theta_i, i = 0,...,N-1  nuisance parameters

  // for now implement m_M = 1, fit single value of mu

  int L = m_L;
  int M = m_M;
  int N = m_N;

  double mu = parVec[0];
  vector<double> thetaVec(N);
  for (int i=0; i<N; i++){
    thetaVec[i] = parVec[M+i];
  }

  double sum = 0.;
  for (int i=0; i<L; i++){
    double Rtheta_i = 0;
    for (int j=0; j<N; j++){
      Rtheta_i += m_R[i][j]*thetaVec[j];
    }
    double q = m_yVec[i] - this->phi(i,mu) - Rtheta_i;
    double w = m_sigma_y_Vec[i]*m_sigma_y_Vec[i];  
    sum += q*q/w;
  }

  for (int i=0; i<N; i++){
    double e = 2.*m_rVec[i]*m_rVec[i];
    double v = m_vVec[i];                // estimate of sigma_u2
    double p = m_uVec[i] - thetaVec[i];  
    sum += (1. + 1./e) * log(1. + e*p*p/v);
  }

  double logL = -0.5*sum;
  return logL;

}

// Profile likelihood lnL(muVec, thetaVec, rVec), profiled over sigma_u,
// terms dropped that do not depend on mu or theta (i.e.
// here r dependence retained, cf. EOE Eq. (77).

double lnLPrime(vector<double> muVec, vector<double> thetaVec, 
	 vector<double> rVec, vector<double> uVec, vector<double> vVec){

  double lnL = 0.;
  int N = thetaVec.size();
  for (int i=0; i<N; i++){
    double r = rVec[i];
    double nu = 1./(2.*r*r);
    double u = uVec[i];
    double theta = thetaVec[i];
    double v = vVec[i];
    double z2 = (u - theta)*(u - theta)/v;
    lnL  += -0.5*(1. + nu)*log(1. + 2.*r*r*z2)
      - 0.5*(1. + nu)*(1. - log(1. + 2.*r*r))
      + (1./(4.*r*r)) * log (1./(4.*r*r)) - TMath::LnGamma(nu/2.);
  }
  return lnL;

}

int Eoe::fit(){                   // use to fix/rel parameters in fit
  vector<bool> freePar(m_npar);
  for (int i=0; i<m_npar; i++){
    freePar[i] = true;
  }
  int status = fitPar(this, freePar);
  return status;
}

// nominal gof statistic (EOE Eq. (61))
double Eoe::q(){
  int status = this->fit();                  // fills m_parVec
  double qVal = -2.*this->lnL(m_parVec);
  return qVal;
}

// gof statistic q' (q with Bartlett correction based on MC)
double Eoe::qBCMC(){
  int L = this->L();  // number of measurements
  int M = this->M();  // number of parameters of interest
  int N = this->N();  // number of nuisance parameters & control meas.
  double ndof = static_cast<double>(L-M);
  double qVal = this->q();
  int status = this->fit();         // sets m_parVec;
  vector<double> thetaVec(N);
  for (int i=0; i<N; i++){
    thetaVec[i] = m_parVec[i+M];
  }
  vector<double> sigma_u_hh = sigma_u_HatHatVec(thetaVec);
  vector<double> genParVec(2*N+M);
  genParVec[0] = m_parVec[0];        // mu
  for (int i=0; i<N; i++){
    genParVec[i+M] = thetaVec[i];      // theta_i
    genParVec[N+i+M] = sigma_u_hh[i];
  }
  this->setEqMC(genParVec); 
  qVal *= ndof / m_EqMC;                       // apply Bartlett correction
  return qVal;
}

// mean values of LR statistics for Bartlett correction; also sets p-value
void Eoe::setEqMC(vector<double> parVec) {
  const int numEvents = 10000;              // adjust as needed
  double eq = 0.;
  double qObs = this->q();
  Eoe eoe = *this;                     // local object
  int m = 0;
  for (int i=0; i<numEvents; i++){
    eoe.generateData(parVec);
    double q = eoe.q();
    eq += q;
    if ( q >= qObs ) {  m++; }
  }
  m_p = static_cast<double>(m)/static_cast<double>(numEvents);   // set p-value
  eq /= static_cast<double>(numEvents);
  m_EqMC = eq;
  m_EqMCSet = true;
  m_pValSet = true;
}

// Generate MC data based on Eoe model and input parameter values
void Eoe::generateData(vector<double> parVec) {

  m_yVec.clear();
  m_uVec.clear();
  m_sVec.clear();
  m_vVec.clear();
  int N = m_N;

  double mu = parVec[0];       // for now supports only one param of int.
  for (int i=0; i<N; i++){
    double sigma_y = m_sigma_y_Vec[i];
    double b = parVec[i+1]; 
    double sigma_u = parVec[N+i+1];      // note this parVec contains sigma_u
    double e = 2.*m_rVec[i]*m_rVec[i];
    double alpha = 1./(2.*e);
    double beta = alpha / (sigma_u*sigma_u);
    double y = m_ran->Gaus(mu + b, sigma_y);
    double u = m_ran->Gaus(b, sigma_u);
    double v = m_rgsl->Gamma(alpha, 1./beta);  // gsl's convention for beta
    double s = sqrt(v);
    m_yVec.push_back(y);
    m_uVec.push_back(u);
    m_sVec.push_back(s);
    m_vVec.push_back(v);
  }

}

// profiled values of sigma_u
vector<double> Eoe::sigma_u_HatHatVec(vector<double> thetaVec){
  vector<double> suhhVec(m_N);
  for (int i=0; i<m_N; i++){
    double r = this->r(i);
    double e = 2.*r*r;
    double v = this->v(i);
    double umb = this->u(i) - thetaVec[i];
    suhhVec[i] = sqrt( (v + e*umb*umb)/(1. + e) );
  }
  return suhhVec;
}

// Estimate start values and step sizes using ordinary LS and linear approx.
void Eoe::setStartPar(){

  int L = this->L();
  int N = this->N();
  int npar = this->npar();

// use ordinary LS for start value of mu, use estimates of biases
// but otherwise ignore systematics
  double wSum = 0.;
  vector<double> w(L);
  for (int i=0; i<L; i++){
    double sy = this->sigma_y(i);
    double su = this->s(i);
    w[i] = 1./(sy*sy);
    wSum += w[i];
  }
  double sigmaMuHatLS = 1./sqrt(wSum);    // needed for step size

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

// Approx. profiled theta based on linear (small r) approximation
  vector<double> thetaHHLVec = this->thetaHatHatLinVec(muHatLS);
  for (int i=0; i<N; i++){
    m_startPar[i+1] = thetaHHLVec[i];
  }

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

}

vector<double> Eoe::thetaHatHatLinVec(double mu) {
  vector<double> muVec(1);
  muVec[0] = mu;
  vector<double> thetaHHLVec = this->thetaHatHatLinVec(muVec);
  return thetaHHLVec;
}

// Approx. profiled theta based on linear (small r) approximation
vector<double> Eoe::thetaHatHatLinVec(vector<double> muVec) {
  int L = this->L();
  int M = this->M();      // here this should be 1
  int N = this->N();
  double mu = muVec[0];
  vector<double> thetaHHLVec(N);
  if ( muVec.size() != 1 ) {
    cout << "thetaHatHatLinVec: multiple PoI not supported, ret 0" << endl;
    for (int i=0; i<N; i++){
      thetaHHLVec[i] = 0.;
    }
    return thetaHHLVec;
  } 

  TMatrixDSym A(N);
  TVectorD B(N);
  for (int i=0; i<N; i++){
    for (int j=0; j<N; j++){
      A[i][j] = 0.;
      for (int k=0; k<L; k++){
        A[i][j] +=  m_R[k][i]*m_R[k][j] / (m_sigma_y_Vec[k]*m_sigma_y_Vec[k]);
      }
      if (i == j) {
        A[i][j] +=  (1. + 2.*m_rVec[i]*m_rVec[i]) / m_vVec[i];
      }
    }
  }

  for (int i=0; i<N; i++){
    B[i] = (1. + 2.*m_rVec[i]*m_rVec[i]) * m_uVec[i]  / m_vVec[i];
    for (int j=0; j<L; j++){
      B[i] +=  (this->m_yVec[j] - this->phi(j, mu)) * m_R[j][i] / 
	(m_sigma_y_Vec[j]*m_sigma_y_Vec[j]);
    }
  }

  TDecompSVD svd(A);
  Bool_t svdOK;
  TVectorD thetaHatHatLin = svd.Solve(B, svdOK);
  if ( !svdOK ) {
    cout << "getStartPar failed to find thetaHatHat, setting to zero" << endl;
    thetaHatHatLin.ResizeTo(N);
    for (int i=0; i<N; i++){
      thetaHatHatLin[i] = 0.;
    }
  }

// copy into output vector and return
  for (int i=0; i<N; i++){
    thetaHHLVec[i] = thetaHatHatLin[i];
  }
  return thetaHHLVec;

}