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

#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];

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

// 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];
      }
    }
  }
}

/* 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];
      }
    }
  }
}

// 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);
  $assume(isPowerOf2(N));
  float R1[N][N], R2[N][N];
  matrixMultiply(N, R1, A0, B0);
  strassenR(N, R2, A0, B0);
  $assert($equals(&R1, &R2));
}