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

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

void lossdistrib(double *p, int *np, double *w, double *S, int *N, 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
     q the loss distribution */

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

  lu = 1./(*N-1);
  q1 = malloc( *N * sizeof(double));
  q2 = malloc( *N * sizeof(double));
  q[0] = 1;
  for(i=0; i<(*np); i++){
    d = S[i] * w[i]/ lu;
    d1 = floor(d);
    d2 = ceil(d);
    p1 = p[i] * (d2-d);
    p2 = p[i] - p1;

    for(j=0; j<d1; j++){
      q1[j]=0;
    };
    for(j=0; j<*N-d1; j++){
      q1[d1+j] = p1*q[j];
    };
    for(j=0; j<d2; j++){
      q2[j]=0;
    };
    for(j=0; j<*N-d2; j++){
      q2[d2+j] = p2*q[j];
    };
    for(j=0; j<*N; j++){
      q[j] = q1[j] + q2[j] + (1-p[i]) * q[j];
    };
  }
  /* correction for weight loss */
  /* sum = 0; */
  /* for(j=0; j<*N; j++){ */
  /*   sum += q[j]; */
  /* } */
  /* q[*N-1] += 1-sum; */
  free(q1);
  free(q2);
}

void lossdistrib_truncated(double *p, int *np, double *w, double *S, int *N, int *T, 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
     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, sizeof(double));
  q2 = calloc( *T, sizeof(double));
  q[0] = 1;
  M = 1;
  for(i=0; i<(*np); i++){
    d = 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);
}

inline void aux2(double *q, double *qtemp, double p, int N, int i, int j){
  int m, n;
  for(m=i; m<N; m++){
    for(n=j; n<N; n++){
      qtemp[m+N*n]+= p * q[(m-i)+N*(n-j)];
    }
  }
}

void lossdistrib_joint( double *p, int *np, double *w, double *S, int *N, 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
     q the joint probability distribution */

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

  lu = 1/(*N-1);
  qtemp = malloc( (*N) * (*N) * sizeof(double));
  q[0] = 1;
  for(k=0; k<(*np); k++){
    x = 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;
    /* init qtemp to 0 */
    for(l=0; l < (*N)*(*N); l++){
      qtemp[l] = 0;
    }
    aux2(q, qtemp, w1 * p[k], *N, i, j);
    aux2(q, qtemp, w2 * p[k], *N, i, j+1);
    aux2(q, qtemp, w3 * p[k], *N, i+1, j+1);
    aux2(q, qtemp, w4 * p[k], *N, i+1, j);
    for(l=0; l < (*N) * (*N); l++){
      q[l] = qtemp[l] + (1-p[k]) * q[l];
    }
  }
  /* correction for weight loss */
  sum = 0;
  for(k=0; k < (*N)*(*N); k++){
    sum+=q[k];
  }
  q[(*N)*(*N)-1]+= 1 - sum;
  free(qtemp);
}

inline void aux(double *q, double *qtemp, double p, int N, int d) {
  /* helper function for the lossdistrib functions */
  int j;
  for(j=d; j<N; j++){
    qtemp[j] += p * q[j-d];
  };
}

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 loss distribution.
     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;
  double lu, d1, d2, dp1, dp2, pp1, pp2, sum;
  double *qtemp;

  lu = 1./(*N - 1);
  qtemp = malloc( (*N) * sizeof(double));
  q[0] = 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;
    for(j=0; j<(*N); j++){
      qtemp[j]=0;
    }
    aux(q, qtemp, dp1, *N, d1l);
    aux(q, qtemp, dp2, *N, d1u);
    aux(q, qtemp, pp1, *N, d2l);
    aux(q, qtemp, pp2, *N, d2u);
    for(j=0; j<(*N); j++){
      q[j] = qtemp[j] + (1-dp[i]-pp[i]) * q[j];
    };
  }
  /* correction for weight loss */
  sum = 0;
  for(j=0; j<(*N); j++){
    sum += q[j];
  }
  q[*N-1] += 1-sum;
  free(qtemp);
}

void lossdistrib_prepay_joint(double *dp, double *pp, int *n, double *w, double *S, int *N, double *q) {
  int i, j1, j2, k, l;
  double lu, x, y1, y2, sum;
  double alpha1, alpha2, beta1, beta2;
  double w1, w2, w3, w4;
  double *qtemp;

  lu = 1./(*N-1);
  qtemp = malloc( (*N) * (*N) * sizeof(double));
  q[0] = 1;
  for(k=0; k<(*n); k++){
    x = S[k] * w[k]/lu;
    y1 = (1-S[k]) * w[k]/lu;
    y2 = w[k]/lu;
    i = floor(x);
    j1 = floor(y1);
    j2 = floor(y2);
    alpha1 = i + 1 - x;
    alpha2 = 1 - alpha1;
    beta1 = j1 + 1 - y1;
    beta2 = 1 - beta1;
    w1 = alpha1 * beta1;
    w2 = alpha1 * beta2;
    w3 = alpha2 * beta2;
    w4 = alpha2 * beta1;
    /* init qtemp to 0 */
    for(l=0; l < (*N)*(*N); l++){
      qtemp[l] = 0;
    }
    aux2(q, qtemp, w1 * dp[k], *N, i, j1);
    aux2(q, qtemp, w2 * dp[k], *N, i, j1+1);
    aux2(q, qtemp, w3 * dp[k], *N, i+1, j1+1);
    aux2(q, qtemp, w4 * dp[k], *N, i+1, j1);
    aux2(q, qtemp, pp[k]*(j2+1-y2), *N, 0, j2);
    aux2(q, qtemp, pp[k]*(y2-j2), *N, 0, j2+1);
    for(l=0; l < (*N)*(*N); l++){
      q[l] = qtemp[l] + (1-dp[k]-pp[k]) * q[l];
    }
  }
  /* correction for weight loss */
  sum = 0;
  for(k=0; k<(*N)*(*N); k++){
    sum+=q[k];
  }
  q[(*N)*(*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 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 BClossdist(double *SurvProb, int *dim1, int *dim2, double *issuerweights,
		double *recov, double *Z, double *w, int *n, double *rho, int *N, int *T, 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
      T: number of probabilities to compute
    outputs:
      L: matrix of size (T, dim2)
      R: matrix of size (T, dim2)
   */
  int t, i, j;
  double g;
  double *gshocked, *Rshocked, *Sshocked, *Lw, *Rw;

  gshocked = calloc((*dim1), sizeof(double));
  Rshocked = calloc((*dim1), sizeof(double));
  Sshocked = calloc((*dim1), sizeof(double));
  Lw = malloc((*T) * sizeof(double));
  Rw = malloc((*T) * 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 */
      for(j=0; j<*T; j++){
	Lw[j]=0;
	Rw[j]=0;
      }
      lossdistrib_truncated(gshocked, dim1, issuerweights, Sshocked, N, T, Lw);
      lossdistrib_truncated(gshocked, dim1, issuerweights, Rshocked, N, T, Rw);
      addandmultiply(Lw, w[i], L + t * (*T), *T);
      addandmultiply(Rw, w[i], R + t * (*T), *T);
    }
  }
  free(gshocked);
  free(Rshocked);
  free(Sshocked);
  free(Lw);
  free(Rw);
}