#include "libscl.h"
#include "kronprd.h"
using namespace scl;
using namespace std;

int main(int argc, char** argp, char** envp)
{
  const INTEGER L = 3;
  const INTEGER M = 3;
  const INTEGER n = 5000;
  const INTEGER drop = 13;

  const INTEGER K = 1 + M*+L;
  
  const REAL minus_log_root_two_pi = -9.1893853320467278e-01;

  ifstream fin;
  ofstream fout;
  string filename;

  if (L == 1) {
    filename = "10s0s0s0s1000000.dbg";
  }
  else if (L == 3) {
    filename = "30s0s0s0s1000000.dbg";
  }
  else error("Error, tstvar, illegal value for L");

  fin.open(filename.c_str());

  string line;
  realmat b0_snp(M,1,0.0);
  realmat B_snp(M,M*L,0.0);
  realmat R0_snp(M*(M+1)/2,1,0.0);
  string::size_type pos;
  while(getline(fin,line)) {

    pos = line.find("MEAN FUNCTION START VALUES FOR b0");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M; ++i) {
        INTEGER discard;
        fin >> b0_snp[i] >> discard;
      }
    }

    pos = line.find("MEAN FUNCTION START VALUES FOR B");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M*M*L; ++i) {
        INTEGER discard;
        fin >> B_snp[i] >> discard;
      }
    }

    pos = line.find("VARIANCE FUNCTION START VALUES FOR Rparms");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M*(M+1)/2; ++i) {
        INTEGER discard;
        fin >> R0_snp[i] >> discard;
      }
    }
  }

  fin.clear(); fin.close();

  realmat b0 = b0_snp;
  realmat B = B_snp;
  realmat R0 = R0_snp;

  realmat R_sim(M,M);
  for (INTEGER i=1; i<=M*(M+1)/2; ++i) R_sim[i] = R0[i];
  sym2gen(R_sim.begin(),M);

  realmat V_sim = R_sim*T(R_sim);

  INT_32BIT seed = 897345;

  realmat data(n,M);
  realmat y0(L,M);
  for (INTEGER t=1; t<=L; ++t) {
    realmat e(M,1);
    for (INTEGER i=1; i<=M; ++i) e[i] = unsk(seed);
    realmat y = b0 + R_sim*e;
    for (INTEGER j=1; j<=M; ++j) y0(t,j) = y[j];
  }

  cout << '\n';
  cout << starbox("/Dimensions//");
  cout << '\n';
  cout << "\t n = " << n << '\n';
  cout << "\t M = " << M << '\n';
  cout << "\t L = " << L << '\n';
  cout << "\t drop = " << drop << '\n';
  cout << '\n';

  cout << '\n';
  cout << starbox("/Parameter true values//");
  cout << '\n';
  cout << "\t b0 = " << b0 << '\n';
  cout << "\t B = " << B << '\n';
  cout << "\t R0 = " << R0 << '\n';
  cout << "\t R_sim = " << R_sim << '\n';
  cout << "\t V_sim = " << V_sim << '\n';
  cout << "\t y0 = " << y0 << '\n';
  cout << '\n';

  for (INTEGER j=1; j<=M; ++j) {
    for (INTEGER t=1; t<=L; ++t) {
      data(t,j) = y0(t,j);
    }
  }

  realmat e(M,1);
  realmat x(M*L,1);
  realmat y(M,1);
  for (INTEGER t=L+1; t<=n; ++t) {
    for (INTEGER k=1; k<=L; ++k) {
      for (INTEGER j=1; j<=M; ++j) x[M*(k-1)+j] = data(t-k,j);
    }
    for (INTEGER i=1; i<=M; ++i) e[i] = unsk(seed);
    y = b0 + B*x + R_sim*e;
    for (INTEGER j=1; j<=M; ++j) data(t,j) = y[j];
  }

  //cerr << data;

  // debug varcoef

  /*
  filename = "data.txt";
  fout.open(filename.c_str());
  if (!fin) error("Error, cannot open " + filename);

  for (INTEGER i=1; i<=n; ++i) {
    for (INTEGER j=1; j<=M; ++j) {
      fout << fmt('e',26,17,data(i,j));
    }
    fout << '\n';
  }

  fout.clear(); fout.close();
  */

  realmat data0(n-drop+L,M);

  for (INTEGER t=1; t<=n-drop+L; ++t) {
    for (INTEGER j=1; j<=M; ++j) {
      data0(t,j) = data(drop-L+t,j);
    }
  }

  intvec idx(M);
  for (INTEGER j=1; j<=M; ++j) idx[j] = j;

  realmat S_ls, C;

  INTEGER rank = varcoef(data0,idx,L,b0,B,S_ls,C);

  realmat S_mle = REAL(data0.nrow()-1-rank)/REAL(data0.nrow()-1)*S_ls;

  realmat V_mle = kronprd(C,S_mle);

  realmat sdev(M,K);
  for (INTEGER i=1; i<=M*K; ++i) sdev[i] = sqrt(V_mle(i,i));

  realmat R = S_mle;
  gen2sym(R.begin(),idx.size());
  factor(R.begin(),idx.size());
  upr2gen(R.begin(),idx.size());

  REAL log_likelihood = 0.0;

  REAL log_detR = 0.0;
  for (INTEGER i=1; i<=M; ++i) log_detR += log(fabs(R(i,i)));
  realmat Rinv = R;
  drinv(Rinv);

  for (INTEGER t=drop+1; t<=n; ++t) {
    for (INTEGER j=1; j<=M; ++j) y[j] = data(t,j);
    for (INTEGER k=1; k<=L; ++k) {
      for (INTEGER j=1; j<=M; ++j) x[M*(k-1)+j] = data(t-k,j);
    }
    realmat z = Rinv*(y - b0 - B*x);
    REAL sumsqr = 0.0;
    for (INTEGER i=1; i<=M; ++i) sumsqr += z[i]*z[i];
    REAL e = REAL(M)*minus_log_root_two_pi - log_detR - 0.5*sumsqr;
    log_likelihood += e;
  }

  REAL lrt = -2.0*log_likelihood; 

  cout << '\n';
  cout << starbox("/varcoef results//");
  cout << '\n';
  cout << "\t b0 = " << b0 << '\n';
  cout << "\t B = " << B << '\n';
  cout << "\t R = " << R << '\n';
  cout << "\t S_ls = " << S_ls << '\n';
  cout << "\t S_mle = " << S_mle << '\n';
  cout << '\n';
  cout << "\t -2 log_likelihood = " << lrt << fmt('e',27,16,lrt) << '\n';
  cout << '\n';
  cout << "\t rank = " << rank << '\n';
  cout << '\n';
  intvec jdx(K,-1); jdx[1] = 1;
  cout << "\t sdev b0 = " << sdev("",jdx) << '\n';
  jdx[1] = -1;
  for (INTEGER j=2; j<=K; ++j) jdx[j] = j;
  cout << "\t sdev B = " << sdev("",jdx) << '\n';
  cout << '\n';

  // check atainst SNP

  if (L == 1) {
    filename = "10s0s0s0s1000000.fit";
  }
  else if (L == 3) {
    filename = "30s0s0s0s1000000.fit";
  }
  else error("Error, tstvar, illegal value for L");

  fin.open(filename.c_str());

  R0_snp.resize(M,M,0.0);

  while(getline(fin,line)) {

    pos = line.find("MEAN FUNCTION START VALUES FOR b0");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M; ++i) {
        INTEGER discard;
        fin >> b0_snp[i] >> discard;
      }
    }

    pos = line.find("MEAN FUNCTION START VALUES FOR B");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M*M*L; ++i) {
        INTEGER discard;
        fin >> B_snp[i] >> discard;
      }
    }

    pos = line.find("VARIANCE FUNCTION START VALUES FOR Rparms");
    if(pos != string::npos) {
      for (INTEGER i=1; i<=M*(M+1)/2; ++i) {
        INTEGER discard;
        fin >> R0_snp[i] >> discard;
      }
    }
  }


  fin.clear(); fin.close();
  sym2gen(R0_snp.begin(),M);
  realmat S_snp = R0_snp*T(R0_snp);

  R = S_snp;
  gen2sym(R.begin(),idx.size());
  factor(R.begin(),idx.size());
  upr2gen(R.begin(),idx.size());

  log_likelihood = 0.0;

  log_detR = 0.0;
  for (INTEGER i=1; i<=M; ++i) log_detR += log(fabs(R(i,i)));
  Rinv = R;
  drinv(Rinv);

  for (INTEGER t=drop+1; t<=n; ++t) {
    for (INTEGER j=1; j<=M; ++j) y[j] = data(t,j);
    for (INTEGER k=1; k<=L; ++k) {
      for (INTEGER j=1; j<=M; ++j) x[M*(k-1)+j] = data(t-k,j);
    }
    realmat z = Rinv*(y - b0_snp - B_snp*x);
    REAL sumsqr = 0.0;
    for (INTEGER i=1; i<=M; ++i) sumsqr += z[i]*z[i];
    REAL e = REAL(M)*minus_log_root_two_pi - log_detR - 0.5*sumsqr;
    log_likelihood += e;
  }

  lrt = -2.0*log_likelihood;

  realmat U,w,V;
  svd(S_snp,U,w,V);
  for (INTEGER i=1; i<=M; ++i) w[i] = sqrt(w[i]);
  realmat R_snp = U*diag(w)*T(V);

  cout << starbox("/snp results//");
  cout << '\n';
  cout << "\t b0 = " << b0_snp << '\n';
  cout << "\t B = " << B_snp << '\n';
  cout << "\t R = " << R_snp << '\n';
  cout << "\t S = " << S_snp << '\n';
  cout << '\n';
  cout << "\t -2 log_likelihood = " << lrt << fmt('e',27,16,lrt) << '\n';
  cout << '\n';

  return 0;
}
