/* VerifyThis 2016 - Challenge 1: Matrix Multiplication
 * Consider the following pseudocode algorithm, which is naive 
 * implementation of matrix multiplication. For simplicity we assume that
 * the matrices are square.

 int[][] matrixMultiply(int[][] A, int[][] B) {
	int n = A.length;

	// initialise C
	int[][] C = new int[n][n];

	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];
       			}
   		}
	}
	return C;
 }

 * Tasks.
 * (1) Provide a specification to describe the behaviour of this algorithm,
 * and prove that it correctly implements its specification.
 * (2) Show that matrix multiplication is associative, i.e., the order in
 * which matrices are multiplied can be disregarded: A(BC) = (AB)C. To show
 * this, you should write a program that performs the two different 
 * computations, and then prove that the result of the two computations is
 * always the same.
 * (3) In the literature, there exist many proposals for more efficient 
 * matrix multiplication algorithms. Strassen’s algorithm was one of the
 * first. The key idea of the algorithm is to use a recurisive algorithm
 * that reduces the number of multiplications on submatrices (from 8 to 7),
 * see https://en.wikipedia.org/wiki/Strassen_algorithm for an explanation.
 * A relatively clean Java implementation (and Python and C++) can be found
 * here: https://martin-thoma.com/strassen-algorithm-in-python-java-cpp/.
 * Prove that the naive algorithm above has the same behaviour as Strassen’s
 * algorithm. Proving it for a restricted case, like a 2x2 matrix should be
 * straightforward, the challenge is to prove it for arbitrary matrices with
 * size 2^n.
 *
 * Author: Stephen Siegel
 */

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

// upper bound on N, the size of the matrices
$input int BOUND = 4; // can go up to 16 if you have a few minutes
$assume(BOUND >= 1);
$input int N=4; // 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 associativity: (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);
}