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

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

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

  N = ceil(1./(*lu) + 1);
  q1 = malloc( N * sizeof(double));
  q2 = malloc( N * sizeof(double));
  q[0] = 1;
  for(i=0; i<(*np); i++){
    d = S[i]/(*np * (*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 lossdistrib2( double *p, int *np, double *S, double *lu, double *q) {
  /* recursive algorithm with first order correction
     computes jointly the loss and recovery distribution
     p vector of default probabilities
     np length of p
     S vector of severities (should be same length as p)
     lu loss unit
     q the joint probability distribution */

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

  N = ceil(1./(*lu) + 1);
  qtemp = malloc( N * N * sizeof(double));
  q[0] = 1;
  for(k=0; k<(*np); k++){
    x = S[k]/(*np * (*lu));
    y = (1-S[k])/(*np * (*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;
    }
    for(m=i; m<N; m++){
      for(n=j; n<N; n++){
	qtemp[m+N*n]+= w1 * p[k]* q[(m-i)+ N*(n-j)];
      }
    }
    for(m=i; m<N; m++){
      for(n=j+1; n<N; n++){
	qtemp[m+N*n]+= w2 * p[k]* q[(m-i)+ N*(n-j-1)];
      }
    }
    for(m=i+1; m<N; m++){
      for(n=j+1; n<N; n++){
	qtemp[m+N*n]+= w3 * p[k]* q[(m-i-1)+ N*(n-j-1)];
      }
    }
    for(m=i+1; m<N; m++){
      for(n=j; n<N; n++){
	qtemp[m+N*n]+= w4 * p[k]* q[(m-i-1)+ N*(n-j)];
      }
    }
    for(m=0; m<N; m++){
      for(n=0; n<N;n++){
	q[m+N*n] = (1-p[k]) * q[m+N*n]+ qtemp[m+N*n];
      }
    }
  }
  /* correction for weight loss */
  sum = 0;
  for(l=0; l<N*N; l++){
    sum+=q[l];
  }
  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 *recov, double *Z, double *w, int *n, double *rho, double *lu, double *L, double *R) {

  int t, i, j, N;
  double g;
  double *gshocked, *Rshocked, *Sshocked, *Lw, *Rw;

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