#include <math.h>
#include <stdlib.h>

#include "toolsmith.h"
#include "Vector.h"
#include "Matrix.h"
#include "Matrixlt.h"
#include "McMath.h"
#include "McStats.h"


double binlink(const Vector &y, const Vector &beta, const Matrix &X, Vector &p)
{
   int i, n = X.nrow();
   double eta, sum = 0.0;

   for(i = 0; i < n; i++) {
      eta = exp(dot(beta, X[i]));
      p[i] = 1.0 / (1.0 / eta + 1.0);
      sum += y[i] * log(eta) + log(1.0 / (1.0 + eta));
   }

   return sum;
};


double cor(const Vector &x)
{
   return 1.0;
}

double cor(const Vector &x, const Vector &y)
{
   return cov(x, y) / sqrt(cov(x) * cov(y));
}

Matrix cor(const Matrix &X)
{
   int i, j, n = X.ncol();
   Matrix Sigma(n, n), Xt(t(X));

   for(i = 0; i < n; i++) {
      for(j = 0; j < i; j++) {
         Sigma[i][j] = Sigma[j][i] = cor(Xt[i], Xt[j]);
      }
      Sigma[i][i] = 1.0;
   }

   return Sigma;
}

Matrix cor(const Matrix &X, const Matrix &Y)
{
   int i, j, m = X.ncol(), n = Y.ncol();
   Matrix Sigma(m, n), Xt(t(X)), Yt(t(Y));

   for(i = 0; i < m; i++) {
      for(j = 0; j < n; j++) {
         Sigma[i][j] = cor(Xt[i], Yt[j]);
      }
   }

   return Sigma;
}

double cov(const Vector &x)
{
   Vector z(x - mean(x));

   return dot(z, z) / (x.length() - 1.0);
}

double cov(const Vector &x, const Vector &y)
{
   return dot(x - mean(x), y - mean(y)) / (x.length() - 1.0);
}

Matrix cov(const Matrix &X)
{
   int i, j, n = X.ncol();
   Matrix Sigma(n, n), Xt(t(X));

   for(i = 0; i < n; i++) {
      for(j = 0; j < i; j++) {
         Sigma[i][j] = Sigma[j][i] = cov(Xt[i], Xt[j]);
      }
      Sigma[i][i] = cov(Xt[i]);
   }

   return Sigma;
}

Matrix cov(const Matrix &X, const Matrix &Y)
{
   int i, j, m = X.ncol(), n = Y.ncol();
   Matrix Sigma(m, n), Xt(t(X)), Yt(t(Y));

   for(i = 0; i < m; i++) {
      for(j = 0; j < n; j++) {
         Sigma[i][j] = cov(Xt[i], Yt[j]);
      }
   }

   return Sigma;
}

int logistic(const Vector &y, const Matrix &X, Vector &beta, Matrix &Sigma)
{
   static int maxit = 50;
   static double tol = 1e-8;

   bool converged = false;
   int iter = 0, i1, i2, k = X.ncol(), n = X.nrow();
   double d1, logliky;
   Vector delta(k), p(n), vk(k);
   Matrix InfoInv(k, k), Mkk(k, k), Mnk(n, k), Score(n, k);

   logliky = binlink(y, beta, X, p);

   while(iter <= maxit) {

      for(i2 = 0; i2 < k; i2++) {
         d1 = 0.0;
         for(i1 = 0; i1 < n; i1++) {
            Score[i1][i2] = (y[i1] - p[i1]) * X[i1][i2];
            Mnk[i1][i2] = sqrt(p[i1] * (1.0 - p[i1])) * X[i1][i2];
            d1 += Score[i1][i2];
         }
         vk[i2] = d1;
      }

      prodt(Mnk, Mkk);
      cholinv(Mkk, InfoInv);

      prod(InfoInv, vk, delta);

      iter++;

      if(!converged) {
         for(i1 = 0; i1 < k; i1++) {
             vk[i1] = beta[i1] + delta[i1];
         }
         while((logliky = binlink(y, vk, X, p)) < logliky) {
            for(i1 = 0; i1 < k; i1++) {
               delta[i1] /= 2;
               vk[i1] = beta[i1] + delta[i1];
            }
         }
         d1 = 0.0;
         for(i1 = 0; i1 < k; i1++) {
            beta[i1] = vk[i1];
            if(fabs(delta[i1]) > d1) d1 = delta[i1];
         }
         converged = (d1 < tol);
      } else {
         break;
      }

   }

   if(!converged) {
      iter = 0;
      cout << "Warning in logistic(): parameter estimates failed to converge"
           << endl;
   }

   prod(Score, InfoInv, Mnk);
   prodt(Mnk, Sigma);

   return iter;
}

double mean(const Vector &x)
{
   return sum(x) / x.length();
}

Vector rnorm(int n, double mean, double sd)
{
   int i = 0;
   double u1, u2, w, c;
   Vector z(n);

   while(i < n) {
      w = 2.0;
      while(w > 1.0) {
         u1 = 2.0 * drand48() - 1.0;
         u2 = 2.0 * drand48() - 1.0;
         w = u1 * u1 + u2 * u2;
      }
      c = sd * sqrt(-2.0 * log(w) / w);
      z[i++] = c * u1 + mean;
      z[i++ % n] = c * u2 + mean;
   }

   return z;
}

Vector rnorm(int n, const Vector &mean, const Vector &sd)
{
   int i, nmean = mean.length(), nsd = sd.length();
   Vector z(n);

   if(nmean == 0 || nsd == 0)
      error("rnorm(int, Vector, Vector)", "NULL vector", ERREXIT);

   z = rnorm(n, 0.0, 1.0);
   for(i = 0; i < n; i++) {
      z[i] = z[i] * sd[i % nmean] + mean[i % nsd];
   }

   return z;
}

Vector rnorm(const Vector &mean, const Matrixlt &L)
{
   return L % rnorm(L.nrow(), 0.0, 1.0) + mean;
}

Matrix rwishart(int df, const Matrixlt &L)
// Odell and Feiveson algorithm for generating Wishart variates
{
   int i, j, n = L.nrow();
   Matrix W(n, n);
   Matrixlt Z(n);
   
   for(i = 0; i < n; i++) {
      Z[i][i] = sqrt(rgamma(0.5 * (df - i), 0.5));
      for (j = 0; j < i; j++) {
         Z[i][j] = rnorm(0.0, 1.0);
      }
   }
   W = L % Z;
   
   return W % t(W);
}