/* Post-commit solution to matrixMultiplication, using CIVL.
 * Stephen Siegel
 */

#include <civlc.cvh>
#include <stdio.h>
#include <pointer.cvh>

// upper bound on N, the size of the matrices
$input int BOUND = 8; // can go up to 16 if you have a few minutes
$assume(BOUND >= 1);
$input int N; // the size of the matrices
$assume(1<=N && N<=BOUND);
// some arbitrary input matrices...
$input float A0[N][N];
$input float B0[N][N];
$input float C0[N][N];

// the "leaf size" for Strassen...
$input int LEAF_SIZE;
$assume (2 <= LEAF_SIZE && LEAF_SIZE <= N);

/* Part 1 */

// impl: C is "out" variable
void matrixMultiply(int n, float C[][], float A[][], float B[][]) {
  for (int i=0; i<n; i++)
    for (int j=0; j<n; j++)
      C[i][j] = 0.0;

  for (int i = 0; i < n; i++) {
    for (int k = 0; k < n; k++) {
      for (int j = 0; j < n; j++) {
        C[i][j] += A[i][k] * B[k][j];
      }
    }
  }
}

// Unfortunately no easy way to specify the sum...

// Some "helpers" for verification...

// Computes dot product of two vectors of length n...
float dot(int n, float u[], float v[]) {
  float sum = 0;

  for (int i=0; i<n; i++)
    sum += u[i]*v[i];
  return sum;
}

// gets the index-th column of matrix mat, putting it in
// contiguous memory starting from result.  Returns
// pointer to element 0 of result.
float * column(int n, float result[], float mat[][], int index) {
  for (int i=0; i<n; i++)
    result[i] = mat[i][index];
  return &result[0];
}


// checks multiplication is correct: entry i,j should
// be the dot product of i-th row of A and j-th column of B...
void testMult(int n) {
  float actual[n][n];
  float tmp[n];

  matrixMultiply(n, actual, A0, B0);
  for (int i=0; i<n; i++)
    for (int j=0; j<n; j++)
      $assert(dot(n, A0[i], column(n, tmp, B0, j)) == actual[i][j]);
}



/* Part 2 */

// tests accociativity: (A0*B0)*C0 = A0*(B0*C0)
void assoc(int n) {
  float T1[n][n], T2[n][n], R1[n][n], R2[n][n];

  matrixMultiply(n, T1, A0, B0);
  matrixMultiply(n, R1, T1, C0);
  matrixMultiply(n, T2, B0, C0);
  matrixMultiply(n, R2, A0, T2);
  $assert($equals(&R1, &R2));
}


/* Part 3 : Strassen */

// adds two nxn matrices.  C is "out" variable.
void add(int n, float C[][], float A[][], float B[][]) {
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      C[i][j] = A[i][j] + B[i][j];
}

// subtracts two nxn matrices.  C is "out" variable.
void subtract(int n, float C[][], float A[][], float B[][]) {
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      C[i][j] = A[i][j] - B[i][j];
}


// Strassen algorithm from
// https://martin-thoma.com/strassen-algorithm-in-python-java-cpp/
// I'm just going to assume n is a power of 2.
// There is no problem dealing with the general case but need more
// time!

// multiplies two nxn matrices, storing result in C
void strassenR(int n, float C[][], float A[][], float B[][]) {
  if (n <= LEAF_SIZE) {
    matrixMultiply(n, C, A, B);
  } else {
    // initializing the new sub-matrices
    int newSize = n / 2;
    float a11[newSize][newSize];
    float a12[newSize][newSize];
    float a21[newSize][newSize];
    float a22[newSize][newSize];

    float b11[newSize][newSize];
    float b12[newSize][newSize];
    float b21[newSize][newSize];
    float b22[newSize][newSize];

    float aResult[newSize][newSize];
    float bResult[newSize][newSize];

    // dividing the matrices in 4 sub-matrices:
    for (int i = 0; i < newSize; i++) {
      for (int j = 0; j < newSize; j++) {
        a11[i][j] = A[i][j]; // top left
        a12[i][j] = A[i][j + newSize]; // top right
        a21[i][j] = A[i + newSize][j]; // bottom left
        a22[i][j] = A[i + newSize][j + newSize]; // bottom right

        b11[i][j] = B[i][j]; // top left
        b12[i][j] = B[i][j + newSize]; // top right
        b21[i][j] = B[i + newSize][j]; // bottom left
        b22[i][j] = B[i + newSize][j + newSize]; // bottom right
      }
    }
    // Calculating p1 to p7:
    add(newSize, aResult, a11, a22);
    add(newSize, bResult, b11, b22);
    float p1[newSize][newSize];
    strassenR(newSize, p1, aResult, bResult);
    // p1 = (a11+a22) * (b11+b22)

    add(newSize, aResult, a21, a22); // a21 + a22
    float p2[newSize][newSize];
    strassenR(newSize, p2, aResult, b11); // p2 = (a21+a22) * (b11)

    subtract(newSize, bResult, b12, b22); // b12 - b22
    float p3[newSize][newSize];
    strassenR(newSize, p3, a11, bResult);
    // p3 = (a11) * (b12 - b22)

    subtract(newSize, bResult, b21, b11); // b21 - b11
    float p4[newSize][newSize];
    strassenR(newSize, p4, a22, bResult);
    // p4 = (a22) * (b21 - b11)

    add(newSize, aResult, a11, a12); // a11 + a12
    float p5[newSize][newSize];
    strassenR(newSize, p5, aResult, b22);
    // p5 = (a11+a12) * (b22)

    subtract(newSize, aResult, a21, a11); // a21 - a11
    add(newSize, bResult, b11, b12); // b11 + b12
    float p6[newSize][newSize];
    strassenR(newSize, p6, aResult, bResult);
    // p6 = (a21-a11) * (b11+b12)

    subtract(newSize, aResult, a12, a22); // a12 - a22
    add(newSize, bResult, b21, b22); // b21 + b22
    float p7[newSize][newSize];
    strassenR(newSize, p7, aResult, bResult);
    // p7 = (a12-a22) * (b21+b22)

    // calculating c21, c21, c11 e c22:
    float c12[newSize][newSize];
    add(newSize, c12, p3, p5); // c12 = p3 + p5
    float c21[newSize][newSize];
    add(newSize, c21, p2, p4); // c21 = p2 + p4

    add(newSize, aResult, p1, p4); // p1 + p4
    add(newSize, bResult, aResult, p7); // p1 + p4 + p7
    float c11[newSize][newSize];
    subtract(newSize, c11, bResult, p5);
    // c11 = p1 + p4 - p5 + p7

    add(newSize, aResult, p1, p3); // p1 + p3
    add(newSize, bResult, aResult, p6); // p1 + p3 + p6
    float c22[newSize][newSize];
    subtract(newSize, c22, bResult, p2);
    // c22 = p1 + p3 - p2 + p6

    // Grouping the results obtained in a single matrix:
    for (int i = 0; i < newSize; i++) {
      for (int j = 0; j < newSize; j++) {
        C[i][j] = c11[i][j];
        C[i][j + newSize] = c12[i][j];
        C[i + newSize][j] = c21[i][j];
        C[i + newSize][j + newSize] = c22[i][j];
      }
    }
  }
}

// test Strassen multiplication agrees with the regular one...
void testStrassen(int n) {
  float R1[n][n], R2[n][n];

  matrixMultiply(n, R1, A0, B0);
  strassenR(n, R2, A0, B0);
  $assert($equals(&R1, &R2));
}

// determines whether n is a power of 2
_Bool isPowerOf2(int n) {
  while (n>1) {
    if (n%2 != 0)
      return $false;
    n = n/2;
  }
  return $true;
}

/* main: runs the three tests */
int main() {
  $elaborate(N); // hint to verifier
  printf("N=%d\n", N);
  testMult(N);
  assoc(N);
  $assume(isPowerOf2(N));
  testStrassen(N);
}