#include <R.h>
#include <Rmath.h>
#include <string.h>
#include <omp.h>

#define MIN(x, y) (((x) < (y)) ? (x) : (y))

extern int dgemv_(char* trans, int *m, int *n, double* alpha, double* A, int* lda,
		  double* x, int* incx, double* beta, double* y, int* incy);
extern double ddot_(int* n, double* dx, int* incx, double* dy, int* incy);

extern int daxpy_(int* n, double* da, double* dx, int* incx, double* dy, int* incy);

void lossdistrib(double *p, int *np, double *w, double *S, int *N, int *defaultflag, double *q);

double shockprob(double p, double rho, double Z, int give_log);

void lossdistrib_Z(double *p, int *np, double *w, double *S, int *N, int *defaultflag,
		   double *rho, double *Z, double *wZ, int *nZ, double *q);

void lossdistrib_truncated(double *p, int *np, double *w, double *S, int *N,
			   int *T, int *defaultflag, double *q);

void lossdistrib_joint( double *p, int *np, double *w, double *S, int *N, int *defaultflag, double *q);

void recovdist(double *dp, double *pp, int *n, double *w, double *S, int *N, double *q);

void lossdistrib_prepay_joint(double *dp, double *pp, int *ndp, double *w,
			      double *S, int *N, int *defaultflag, double *q);
double dqnorm(double x);

double dshockprob(double p, double rho, double Z);

void shockprobvec2(double p, double rho, double* Z, int nZ, double *q);

double shockseverity(double S, double Z, double rho, double p);

void fitprob(double* Z, double* w, int* nZ, double* rho, double* p0, double* result);

void stochasticrecov(double* R, double* Rtilde, double* Z, double* w, int* nZ, double* rho,
		     double* porig, double* pmod, double* q);

void lossdistrib_prepay_joint_Z(double *dp, double *pp, int *ndp, double *w,
				double *S, int *N, int *defaultflag, double *rho,
				double *Z, double *wZ, int *nZ, double *q);
void lossdistrib_joint_Z(double *dp, int *ndp, double *w,
			 double *S, int *N, int *defaultflag, double *rho,
			 double *Z, double *wZ, int *nZ, double *q);

void BClossdist(double *SurvProb, int *dim1, int *dim2, double *issuerweights,
		double *recov, double *Z, double *w, int *n, double *rho, int *N,
		int *defaultflag, double *L, double *R);

void lossdistrib(double *p, int *np, double *w, double *S, int *N, int *defaultflag, double *q) {
  /* recursive algorithm with first order correction for computing
     the loss distribution.
     p vector of default probabilities
     np length of p
     S vector of severities (should be same length as p)
     N number of ticks in the grid
     defaultflat if true compute the default distribution
     q the loss distribution */

  int i, j, d1, d2, M;
  double lu, d, p1, p2, sum;
  double *qtemp;

  lu = 1./(*N-1);
  qtemp = Calloc(*N, double);
  q[0] = 1;
  M = 1;
  for(i=0; i<(*np); i++){
    d = (*defaultflag)? w[i]/lu : S[i] * w[i]/ lu;
    d1 = floor(d);
    d2 = ceil(d);
    p1 = p[i] * (d2-d);
    p2 = p[i] - p1;
    memcpy(qtemp, q, MIN(M, *N) * sizeof(double));
    for(j=0; j < MIN(M, *N); j++){
      q[j] = (1-p[i]) * q[j];
    }
    for(j=0; j < MIN(M, *N-d2); j++){
      q[d1+j] += p1 * qtemp[j];
      q[d2+j] += p2 * qtemp[j];
    };
    M+=d2;
  }

  /* correction for weight loss */
  if(M>*N){
    sum = 0;
    for(j=0; j<MIN(M, *N); j++){
      sum += q[j];
    }
    q[*N-1] += 1-sum;
  }
  Free(qtemp);
}

void lossdistrib_Z(double *p, int *np, double *w, double *S, int *N, int *defaultflag,
		   double *rho, double *Z, double *wZ, int *nZ, double *q){
  int i, j;
  double* pshocked = malloc(sizeof(double) * (*np) * (*nZ));

#pragma omp parallel for private(j)
  for(i = 0; i < *nZ; i++){
    for(j = 0; j < *np; j++){
      pshocked[j + (*np) * i] = shockprob(p[j], *rho, Z[i], 0);
    }
    lossdistrib(pshocked + (*np) * i, np, w, S + (*np) * i, N,
		defaultflag, q + (*N) * i);
  }
}

void lossdistrib_truncated(double *p, int *np, double *w, double *S, int *N,
			   int *T, int *defaultflag, double *q) {
  /* recursive algorithm with first order correction for computing
     the loss distribution.
     p vector of default probabilities
     np length of p
     S vector of severities (should be same length as p)
     N number of ticks in the grid
     T where to truncate the distribution
     defaultflat if true computes the default distribution
     q the loss distribution */

  int i, j, d1, d2, M;
  double lu, d, p1, p2, sum;
  double *q1, *q2;

  lu = 1./(*N-1);
  q1 = Calloc( *T, double);
  q2 = Calloc( *T, double);
  q[0] = 1;
  M = 1;
  for(i=0; i<(*np); i++){
    d = (*defaultflag)? w[i] / lu : S[i] * w[i] / lu;
    d1 = floor(d);
    d2 = ceil(d);
    p1 = p[i] * (d2-d);
    p2 = p[i] - p1;
    for(j=0; j < MIN(M, *T); j++){
      q1[j] = p1 * q[j];
      q2[j] = p2 * q[j];
      q[j] = (1-p[i]) * q[j];
    }
    for(j=0; j < MIN(M, *T-d1); j++){
      q[d1+j] += q1[j];
    };
    for(j=0; j < MIN(M, *T-d2); j++){
      q[d2+j] += q2[j];
    };
    M += d2;
  }
  Free(q1);
  Free(q2);
}

void lossdistrib_joint( double *p, int *np, double *w, double *S, int *N, int *defaultflag, double *q) {
  /* recursive algorithm with first order correction
     computes jointly the loss and recovery distribution
     p vector of default probabilities
     np length of p
     w vector of issuer weights (length np)
     S vector of severities (should be same length as p)
     N number of ticks in the grid
     defaultflag if true computes the default distribution
     q the joint probability distribution */

  int i, j, k, m, n;
  int Mx, My;
  double lu, x, y, sum;
  double alpha1, alpha2, beta1, beta2;
  double w1, w2, w3, w4;
  double *qtemp;

  lu = 1./(*N-1);
  qtemp = Calloc( (*N) * (*N), double);
  q[0] = 1;
  Mx=1;
  My=1;
  for(k=0; k<(*np); k++){
    x = (*defaultflag)? w[k] /lu : S[k] * w[k] / lu;
    y = (1-S[k]) * w[k] / lu;
    i = floor(x);
    j = floor(y);
    alpha1 = i + 1 - x;
    alpha2 = 1 - alpha1;
    beta1 = j + 1 - y;
    beta2 = 1 - beta1;
    w1 = alpha1 * beta1;
    w2 = alpha1 * beta2;
    w3 = alpha2 * beta2;
    w4 = alpha2 * beta1;

    for(n=0; n<MIN(My, *N); n++){
      memcpy(qtemp+n*(*N), q+n*(*N), MIN(Mx, *N) * sizeof(double));
    }
    for(n=0; n<MIN(My, *N); n++){
      for(m=0; m<MIN(Mx, *N); m++){
	q[m+(*N)*n] = (1-p[k])* q[m+(*N)*n];
      }
    }
    for(n=0; n < MIN(My, *N-j-1); n++){
      for(m=0; m < MIN(Mx, *N-i-1); m++){
	q[i+m+(*N)*(j+n)] += w1 * p[k] * qtemp[m+(*N)*n];
	q[i+m+(*N)*(j+1+n)] += w2 * p[k] * qtemp[m+(*N)*n];
	q[i+1+m+(*N)*(j+1+n)] += w3 * p[k] * qtemp[m+(*N)*n];
	q[i+1+m+(*N)*(j+n)] += w4 * p[k] *qtemp[m+(*N)*n];
      }
    }
    Mx += i+1;
    My += j+1;
  }
  /* correction for weight loss */
  if(Mx>*N || My>*N){
    sum = 0;
    for(m=0; m < MIN(Mx, *N); m++){
      for(n=0; n < MIN(My, *N); n++){
  	sum += q[m+n*(*N)];
      }
    }
    q[MIN(*N, Mx)*MIN(My,*N)-1] += 1 - sum;
  }
  Free(qtemp);
}

void recovdist(double *dp, double *pp, int *n, double *w, double *S, int *N, double *q) {
  /* recursive algorithm with first order correction for computing
     the recovery distribution in case of prepayment.
     dp vector of default probabilities
     pp vector of prepay probabilities
     n length of p
     S vector of severities (should be same length as p)
     w vector of weights
     N number of ticks in the grid
     q the loss distribution */

  int i, j, d1l, d1u, d2l, d2u;
  int M;
  double lu, d1, d2, dp1, dp2, pp1, pp2, sum;
  double *qtemp;

  lu = 1./(*N - 1);
  qtemp = Calloc( (*N), double);
  q[0] = 1;
  M=1;
  for(i=0; i<(*n); i++){
    d1 = w[i] * (1-S[i]) /lu;
    d2 = w[i]/lu;
    d1l = floor(d1);
    d1u = d1l + 1;
    d2l = floor(d2);
    d2u = d2l + 1;
    dp1 = dp[i] * (d1u - d1);
    dp2 = dp[i] - dp1;
    pp1 = pp[i] * (d2u - d2);
    pp2 = pp[i] - pp1;
    memcpy(qtemp, q, MIN(M, *N) * sizeof(double));
    for(j = 0; j< MIN(M, *N); j++){
      q[j] = (1-dp[i]-pp[i]) * q[j];
    }
    for(j=0; j < MIN(M, *N-d2u); j++){
      q[d1l+j] += dp1 * qtemp[j];
      q[d1u+j] += dp2 * qtemp[j];
      q[d2l+j] += pp1 * qtemp[j];
      q[d2u+j] += pp2 * qtemp[j];
    };
    M += d2u;
  }
  /* correction for weight loss */
  if(M>*N){
    sum = 0;
    for(j=0; j<MIN(M, *N); j++){
      sum += q[j];
    }
    q[*N-1] += 1-sum;
  }
  Free(qtemp);
}

void lossdistrib_prepay_joint(double *dp, double *pp, int *ndp, double *w,
			      double *S, int *N, int *defaultflag, double *q) {
  int i, j1, j2, k, m, n;
  double lu, x, y1, y2, sum;
  double alpha1, alpha2, beta1, beta2;
  double dpw1, dpw2, dpw3, dpw4;
  double ppw1, ppw2, ppw3;
  double temp;
  double *qtemp;
  int Mx, My;

  lu = 1./(*N-1);
  qtemp = Calloc((*N) * (*N), double);
  q[0] = 1;
  Mx=1;
  My=1;
  for(k=0; k<(*ndp); k++){
    y1 = (1-S[k]) * w[k]/lu;
    y2 = w[k]/lu;
    x = (*defaultflag)? y2: y2-y1;
    i = floor(x);
    j1 = floor(y1);
    j2 = floor(y2);
    alpha1 = i + 1 - x;
    alpha2 = 1 - alpha1;
    beta1 = j1 + 1 - y1;
    beta2 = 1 - beta1;
    dpw1 = alpha1 * beta1 * dp[k];
    dpw2 = alpha1 * beta2 * dp[k];
    dpw3 = alpha2 * beta2 * dp[k];
    dpw4 = alpha2 * beta1 * dp[k];

    /* by default distribution, we mean names fractions of names that disappeared
       either because of default or prepayment */
    for(n=0; n<MIN(My, *N); n++){
      memcpy(qtemp+n*(*N), q+n*(*N), MIN(Mx, *N) * sizeof(double));
    }
    for(n=0; n<MIN(My, *N); n++){
      for(m=0; m<MIN(Mx, *N); m++){
	q[m+(*N)*n] = (1-dp[k]-pp[k]) * q[m+(*N)*n];
      }
    }
    if(*defaultflag){
      ppw1 = alpha1 * alpha1 * pp[k];
      ppw2 = alpha1 * alpha2 * pp[k];
      ppw3 = alpha2 * alpha2 * pp[k];
      for(n=0; n < MIN(My, *N-j2-1); n++){
	for(m=0; m < MIN(Mx, *N-i-1); m++){
	  q[i+m+(*N)*(j1+n)] += dpw1 * qtemp[m+(*N)*n];
	  q[i+m+(*N)*(j1+1+n)] += dpw2 * qtemp[m+(*N)*n];
	  q[i+1+m+(*N)*(j1+1+n)] += dpw3 * qtemp[m+(*N)*n];
	  q[i+1+m+(*N)*(j1+n)] += dpw4 * qtemp[m+(*N)*n];

	  q[i+m+(*N)*(j2+n)] += ppw1 * qtemp[m+(*N)*n];
	  temp = ppw2 * qtemp[m+(*N)*n];
	  q[i+m+(*N)*(j2+1+n)] += temp;
	  q[i+1+m+(*N)*(j2+1+n)] += ppw3 * qtemp[m+(*N)*n];
	  q[i+1+m+(*N)*(j2+n)] += temp;
	}
      }
    }else{
      for(n=0; n < MIN(My, *N-j2-1); n++){
	for(m=0; m < MIN(Mx, *N-i-1); m++){
	  q[i+m+(*N)*(j1+n)] += dpw1 * qtemp[m+(*N)*n];
	  q[i+m+(*N)*(j1+1+n)] += dpw2 * qtemp[m+(*N)*n];
	  q[i+1+m+(*N)*(j1+1+n)] += dpw3 * qtemp[m+(*N)*n];
	  q[i+1+m+(*N)*(j1+n)] += dpw4 * qtemp[m+(*N)*n];
	  q[m+(*N)*(j2+n)] +=  pp[k]*(j2+1-y2) * qtemp[m+(*N)*n];
	  q[m+(*N)*(j2+1+n)] +=  pp[k]*(y2-j2) * qtemp[m+(*N)*n];
	}
      }
    }
    Mx += i + 1;
    My += j2 + 1;
  }
  /* correction for weight loss */
  if(Mx>*N || My>*N){
    sum = 0;
    for(m=0; m < MIN(Mx, *N); m++){
      for(n=0; n < MIN(My, *N); n++){
	sum += q[m+n*(*N)];
      }
    }
    q[MIN(*N, Mx)*MIN(My,*N)-1] += 1 - sum;
  }
  Free(qtemp);
}

double shockprob(double p, double rho, double Z, int give_log){
  return( pnorm( (qnorm(p, 0, 1, 1, 0) - sqrt(rho) * Z)/sqrt(1 - rho), 0, 1, 1, give_log));
}

double dqnorm(double x){
  return 1/dnorm(qnorm(x, 0, 1, 1, 0), 0, 1, 0);
}

double dshockprob(double p, double rho, double Z){
  return( dnorm((qnorm(p, 0, 1, 1, 0) - sqrt(rho) * Z)/sqrt(1-rho), 0, 1, 0) * dqnorm(p)/sqrt(1-rho) );
}

void shockprobvec2(double p, double rho, double* Z, int nZ, double *q){
  /* return a two column vectors with shockprob in the first column
   and dshockprob in the second column*/
  int i;
#pragma omp parallel for
  for(i = 0; i < nZ; i++){
    q[i] = shockprob(p, rho, Z[i], 0);
    q[i + nZ] = dshockprob(p, rho, Z[i]);
  }
}

double shockseverity(double S, double Z, double rho, double p){
  return( exp(shockprob(S * p, rho, Z, 1) - shockprob(p, rho, Z, 1)) );
}

/* void addandmultiply(double *X, double alpha, double *Y, int n) { */
/*   int i; */
/*   for(i = 0; i<n; i++){ */
/*     Y[i] += alpha*X[i]; */
/*   } */
/* } */

void fitprob(double* Z, double* w, int* nZ, double* rho, double* p0, double* result){
  double eps = 1e-12;
  int one = 1;
  double *q = malloc( 2 * (*nZ) * sizeof(double));
  double dp, p, phi;

  if(*p0 == 0){
    *result = 0;
  }else{
    shockprobvec2(*p0, *rho, Z, *nZ, q);
    dp = (ddot_(nZ, q, &one, w, &one) - *p0)/ddot_(nZ, q + *nZ, &one, w, &one);
    p = *p0;
    while(fabs(dp) > eps){
      phi = 1;
      while( (p - phi * dp) < 0 || (p - phi * dp) > 1){
	phi *= 0.8;
      }
      p -= phi * dp;
      shockprobvec2(p, *rho, Z, *nZ, q);
      dp = (ddot_(nZ, q, &one, w, &one) - *p0)/ddot_(nZ, q + *nZ, &one, w, &one);
    }
    *result = p;
  }
  free(q);
}

void stochasticrecov(double* R, double* Rtilde, double* Z, double* w, int* nZ, double* rho,
		     double* porig, double* pmod, double* q){
  double ptemp, ptilde;
  int i;
  if(*porig==0){
    for(i = 0; i < *nZ; i++){
      q[i] = *R;
    }
  }else{
    ptemp = (1 - *R) / (1 - *Rtilde) * *porig;
    fitprob(Z, w, nZ, rho, &ptemp, &ptilde);
    #pragma omp parallel for
    for(i = 0; i < *nZ; i++){
      q[i] = fabs(1 - (1 - *Rtilde) * exp( shockprob(ptilde, *rho, Z[i], 1) - shockprob(*pmod, *rho, Z[i], 1)));
    }
  }
}

void lossdistrib_prepay_joint_Z(double *dp, double *pp, int *ndp, double *w,
				double *S, int *N, int *defaultflag, double *rho,
				double *Z, double *wZ, int *nZ, double *q) {
  int i, j;
  double* dpshocked = malloc(sizeof(double) * (*ndp) * (*nZ));
  double* ppshocked = malloc(sizeof(double) * (*ndp) * (*nZ));
  int N2 = (*N) * (*N);
  double* qmat = malloc(sizeof(double) * N2 * (*nZ));

  double alpha = 1;
  double beta = 0;
  int one = 1;

#pragma omp parallel for private(j)
  for(i = 0; i < *nZ; i++){
    for(j = 0; j < *ndp; j++){
      dpshocked[j + (*ndp) * i] = shockprob(dp[j], *rho, Z[i], 0);
      ppshocked[j + (*ndp) * i] = shockprob(pp[j], *rho, -Z[i], 0);
    }
    lossdistrib_prepay_joint(dpshocked + (*ndp) * i, ppshocked + (*ndp) * i, ndp, w,
			     S + (*ndp) * i, N, defaultflag, qmat + N2 * i);
  }

  dgemv_("n", &N2, nZ, &alpha, qmat, &N2, wZ, &one, &beta, q, &one);

  free(dpshocked);
  free(ppshocked);
  free(qmat);
}

void lossdistrib_joint_Z(double *dp, int *ndp, double *w,
			 double *S, int *N, int *defaultflag, double *rho,
			 double *Z, double *wZ, int *nZ, double *q) {
  int i, j;
  double* dpshocked = malloc(sizeof(double) * (*ndp) * (*nZ));
  int N2 = (*N) * (*N);
  double* qmat = malloc(sizeof(double) * N2 * (*nZ));

  double alpha = 1;
  double beta = 0;
  int one = 1;

#pragma omp parallel for private(j)
  for(i = 0; i < *nZ; i++){
    for(j = 0; j < *ndp; j++){
      dpshocked[j + (*ndp) * i] = shockprob(dp[j], *rho, Z[i], 0);
    }
    lossdistrib_joint(dpshocked + (*ndp) * i, ndp, w, S + (*ndp) * i, N,
		      defaultflag, qmat + N2 * i);
  }

  dgemv_("n", &N2, nZ, &alpha, qmat, &N2, wZ, &one, &beta, q, &one);

  free(dpshocked);
  free(qmat);
}

void BClossdist(double *SurvProb, int *dim1, int *dim2, double *issuerweights,
		double *recov, double *Z, double *w, int *n, double *rho, int *N,
		int *defaultflag, double *L, double *R) {
  /*
    computes the loss and recovery distribution over time with a flat gaussiancorrelation
    inputs:
      Survprob: matrix of size dim1 x dim2. dim1 is the number of issuers and dim2 number of time steps
      recov: vector of recoveries (length dim1)
      issuerweights: vector of issuer weights (length dim2)
      Z: vector of factor values (length n)
      w: vector of factor weights (length n)
      rho: correlation beta
      N: number of ticks in the grid
      defaultflag: if true, computes the default distribution
    outputs:
      L: matrix of size (N, dim2)
      R: matrix of size (N, dim2)
   */
  int t, i, j;
  double g;
  double *gshocked, *Rshocked, *Sshocked, *Lw, *Rw;
  int one = 1;

  gshocked = Calloc((*dim1), double);
  Rshocked = Calloc((*dim1), double);
  Sshocked = Calloc((*dim1), double);
  Lw = malloc((*N) * sizeof(double));
  Rw = malloc((*N) * sizeof(double));

  for(t=0; t < (*dim2); t++) {
    for(i=0; i < *n; i++){
      for(j=0; j < (*dim1); j++){
	g = 1 - SurvProb[j + (*dim1) * t];
	gshocked[j] = shockprob(g, *rho, Z[i], 0);
	Sshocked[j] = shockseverity(1-recov[j], Z[i], *rho, g);
	Rshocked[j] = 1 - Sshocked[j];
      }
      /* reset Lw and Rw to 0 */
      memset(Lw, 0, *N * sizeof(double));
      memset(Rw, 0, *N * sizeof(double));
      lossdistrib(gshocked, dim1, issuerweights, Sshocked, N, defaultflag, Lw);
      lossdistrib(gshocked, dim1, issuerweights, Rshocked, N, defaultflag, Rw);
      /* addandmultiply(Lw, w[i], L + t * (*N), *N); */
      /* addandmultiply(Rw, w[i], R + t * (*N), *N); */
      daxpy_(N, w + i, Lw, &one, R + t * (*N), &one);
      daxpy_(N, w + i, Rw, &one, R + t * (*N), &one);
    }
  }
  Free(gshocked);
  Free(Rshocked);
  Free(Sshocked);
  free(Lw);
  free(Rw);
}