/* sd_code.c

   James S. Plank
   plank@cs.utk.edu
   http://web.eecs.utk.edu/~plank

   Professor
   EECS Department
   University of Tennessee
   Knoxville, TN 37996

   January, 2013.

   Part of the SD coding library.  Please see Technical Report UT-CS-13-704
   http://web.eecs.utk.edu/~plank/plank/papers/FAST-2013-SD.html
 */

#include <strings.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

#include "gf_complete.h"

#define talloc(type, num) (type *) malloc(sizeof(type)*(num))

usage(char *s)
{
  fprintf(stderr, "usage: sd_code n m s r w size IO(0|1)  a_0 a_1 ... a_{m+s-1} d_0 .. d_{m-1} s_0 .. s_{s-1}\n\n");
  fprintf(stderr, "Reads in the codeword (size bytes per block) on stdin, kills disks d_x and sectorx s_y ... and decodes, printing the bytes on stdout\n");
  if (s != NULL) fprintf(stderr, "%s\n", s);
  exit(1);
}

void
timer_start (double *t)
{
    struct timeval  tv;

    gettimeofday (&tv, NULL);
    *t = (double)tv.tv_sec + (double)tv.tv_usec * 1e-6;
}

double
timer_split (const double *t)
{
    struct timeval  tv;
    double  cur_t;

    gettimeofday (&tv, NULL);
    cur_t = (double)tv.tv_sec + (double)tv.tv_usec * 1e-6;
    return (cur_t - *t);
}

void print_data(int n, int r, int size, uint8_t **array)
{
  int i, j;

  for (i = 0; i < n*r; i++) {
    for (j = 0; j < size; j++) {
      printf(" %02x", array[i][j]);
    }
    printf("\n");
  }
}

int invert_matrix(int *mat, int *inv, int rows, gf_t *gf)
{
  int cols, i, j, k, x, rs2;
  int row_start, tmp, inverse;
 
  cols = rows;

  k = 0;
  for (i = 0; i < rows; i++) {
    for (j = 0; j < cols; j++) {
      inv[k] = (i == j) ? 1 : 0;
      k++;
    }
  }

  /* First -- convert into upper triangular  */
  for (i = 0; i < cols; i++) {
    row_start = cols*i;

    /* Swap rows if we ave a zero i,i element.  If we can't swap, then the 
       matrix was not invertible  */

    if (mat[row_start+i] == 0) { 
      for (j = i+1; j < rows && mat[cols*j+i] == 0; j++) ;
      if (j == rows) return -1;
      rs2 = j*cols;
      for (k = 0; k < cols; k++) {
        tmp = mat[row_start+k];
        mat[row_start+k] = mat[rs2+k];
        mat[rs2+k] = tmp;
        tmp = inv[row_start+k];
        inv[row_start+k] = inv[rs2+k];
        inv[rs2+k] = tmp;
      }
    }
 
    /* Multiply the row by 1/element i,i  */
    tmp = mat[row_start+i];
    if (tmp != 1) {
      inverse = gf->divide.w32(gf, 1, tmp);
      for (j = 0; j < cols; j++) { 
        mat[row_start+j] = gf->multiply.w32(gf, mat[row_start+j], inverse);
        inv[row_start+j] = gf->multiply.w32(gf, inv[row_start+j], inverse);
      }
    }

    /* Now for each j>i, add A_ji*Ai to Aj  */
    k = row_start+i;
    for (j = i+1; j != cols; j++) {
      k += cols;
      if (mat[k] != 0) {
        if (mat[k] == 1) {
          rs2 = cols*j;
          for (x = 0; x < cols; x++) {
            mat[rs2+x] ^= mat[row_start+x];
            inv[rs2+x] ^= inv[row_start+x];
          }
        } else {
          tmp = mat[k];
          rs2 = cols*j;
          for (x = 0; x < cols; x++) {
            mat[rs2+x] ^= gf->multiply.w32(gf, tmp, mat[row_start+x]);
            inv[rs2+x] ^= gf->multiply.w32(gf, tmp, inv[row_start+x]);
          }
        }
      }
    }
  }

  /* Now the matrix is upper triangular.  Start at the top and multiply down  */

  for (i = rows-1; i >= 0; i--) {
    row_start = i*cols;
    for (j = 0; j < i; j++) {
      rs2 = j*cols;
      if (mat[rs2+i] != 0) {
        tmp = mat[rs2+i];
        mat[rs2+i] = 0; 
        for (k = 0; k < cols; k++) {
          inv[rs2+k] ^= gf->multiply.w32(gf, tmp, inv[row_start+k]);
        }
      }
    }
  }
  return 0;
}


main(int argc, char **argv)
{
  gf_t gfs, gfm;
  int i, j, n, m, r, s, w, c, size, nf;
  int *matrix;
  int *encoder, *inverse;
  int *erased, *eindex;
  int alpha, k, coef, symbol, l;
  int *alphas, alpha_counter;
  uint8_t **array;
  uint8_t **syndromes;
  double timer, split;
  int IO;

  if (argc <= 8) usage(NULL);
  if (sscanf(argv[1], "%d", &n) == 0 || n <= 0) usage("Bad n");
  if (sscanf(argv[2], "%d", &m) == 0 || m <= 0) usage("Bad m");
  if (sscanf(argv[3], "%d", &s) == 0 || s <= 0) usage("Bad s");
  if (sscanf(argv[4], "%d", &r) == 0 || r <= 0) usage("Bad r");
  if (sscanf(argv[5], "%d", &w) == 0 || w <= 0) usage("Bad w");
  if (sscanf(argv[6], "%d", &size) == 0 || size <= 0) usage("Bad size");
  if (sscanf(argv[7], "%d", &IO) == 0 || IO < 0 || IO > 1) usage("Bad IO");
  if (size % 8 != 0) usage("Size has to be a multiple of 8\n");
  if (argc != 8 + 2*(m+s)) usage("Wrong number of arguments");

  if (w == 16 || w == 32) {
    if (!gf_init_hard(&gfm, w, GF_MULT_SPLIT_TABLE, GF_REGION_ALTMAP | GF_REGION_SSE, GF_DIVIDE_DEFAULT, 0, 4, w, NULL, NULL)) {
      printf("Bad gf spec\n");
      exit(1);
    }
  } else if (w == 8) {
    if (!gf_init_easy(&gfm, w)) {
      printf("Bad gf spec\n");
      exit(1);
    }
  } else {
    printf("Not supporting w = %d\n", w);
  }

  if (!gf_init_easy(&gfs, w)) {
    printf("Bad gf spec\n");
    exit(1);
  }
  alphas = talloc(int, m+s);

  for (i = 0; i < m+s; i++) {
    if (sscanf(argv[8+i], "%d", alphas+i) == 0 || alphas[i] <= 0) usage("Bad a_i");
  }

  nf = r*m + s;

  erased = talloc(int, n*r);
  for (i = 0; i < n*r; i++) erased[i] = 0;
  eindex = talloc(int, nf);

  for (i = 0; i < m; i++) {
    if (sscanf(argv[8+m+s+i], "%d", &j) != 1 || j < 0 || j >= n) usage("Bad d_x");
    if (erased[j]) usage("Duplicate failed disk");
    while (j < n*r) {
      erased[j] = 1;
      j += n;
    }
  }
  for (i = 0; i < s; i++) {
    if (sscanf(argv[8+2*m+s+i], "%d", &j) != 1 || j < 0 || j >= n*r) usage("Bad s_y");
    if (erased[j]) usage("Duplicate failed sector");
    erased[j] = 1;
  }

  j = 0;
  for (i = 0; i < n*r; i++) {
    if (erased[i]) {
      eindex[j] = i;
      j++;
    }
  }

  matrix = talloc(int, n*r*(m*r+s));
  encoder = talloc(int, (m*r+s)*(m*r+s));
  inverse = talloc(int, (m*r+s)*(m*r+s));

  if (!IO) srand48(time(0));

  array = talloc(uint8_t *, n*r);
  for (i = 0; i < n*r; i++) array[i] = talloc(uint8_t, size);

  for (i = 0; i < n*r; i++) {
    for (j = 0; j < size; j++) {
      if (!IO) {
        array[i][j] = lrand48()%256;
      } else {
        if (scanf("%x", &c) != 1) {
          fprintf(stderr, "Bad input at block %d byte %d\n", i, j);
        } else array[i][j] = c; 
      }
    }
  }
  syndromes = talloc(uint8_t *, (m*r+s));
  for (i = 0; i < m*r+s; i++) {
    syndromes[i] = talloc(uint8_t, size);
    bzero(syndromes[i], size);
  }

  /* Create the parity check matrix.  First the m*r rows corresponding to C_{i,j}. */

  alpha_counter = 0;
  for (i = 0; i < m; i++) {
    coef = 1;
    for (k = 0; k < r; k++) {
      for (j = 0; j < n; j++) {
        matrix[(k*m+i)*(r*n)+(k*n)+j] = coef;
        coef = gfs.multiply.w32(&gfs, coef, alphas[alpha_counter]);
      }
    }
    alpha_counter++;
  }

  /* Then the final s rows corresponding to S_x */

  for (i = 0; i < s; i++) {
    coef = 1;
    for (j = 0; j < n*r; j++) {
      matrix[(n*r)*(m*r+i)+j] = coef;
      coef = gfs.multiply.w32(&gfs, coef, alphas[alpha_counter]);
    }
    alpha_counter++;
  }

  /* Erase the bad blocks (turn them into zeros, and create a decoding matrix
     from the columns of the parity check matrix that correspond to the failed
     blocks. */

  c = 0;
  for (i = 0; i < n*r; i++) {
    if (erased[i]) {
      bzero(array[i], size);
      for (j = 0; j < (r*m+s); j++) {
        encoder[j*(r*m+s)+c] = matrix[j*n*r+i];
      }
      eindex[c] = i;
      c++;
    }
  }

  timer_start(&timer);
  if (invert_matrix(encoder, inverse, r*m+s, &gfs) == -1) {
    printf("Can't invert\n");
    exit(0);
  }

  for (i = 0; i < n*r; i++) {
    if (!erased[i]) {
      for (j = 0; j < m*r+s; j++) {
        coef = matrix[j*n*r+i];
        if (coef != 0) {
          gfm.multiply_region.w32(&gfm, array[i], syndromes[j], coef, size, 1);
        }
      }
    }
  }

  for (i = 0; i < m*r+s; i++) {
    for (j = 0; j < m*r+s; j++) {
      coef = inverse[j*(m*r+s)+i];
      if (coef != 0) {
        gfm.multiply_region.w32(&gfm, syndromes[i], array[eindex[j]], coef, size, 1);
      }
    }
  }

  split = timer_split(&timer);

  fprintf(stderr, "Timer: %.6lf\n", split); 
  if (IO) print_data(n, r, size, array);
}