func matrix calculate_weight (matrix x, matrix y) {
  matrix result <2, 1>;
  result = inv(x' * x) * x' * y;
  return result;
}

func void gradient_descent (matrix x, matrix y) {
  int n = 3;
  matrix y_current <3, 1>;
  matrix sum_c <1,1>;
  
  double learning_rate = 0.0001;
  double w_current = 0.0;
  double b_current = 0.0;
  double w_gradient = 0.0;
  double b_gradient = 0.0;

  int i;
  int j;
  double cost;
  
  for (i = 0; i < 250000; i = i + 1) {
    y_current = (w_current * x) + b_current;
    cost = 0.0;
    for(j = 0; j < n; j = j + 1){
	    cost = cost + (y[j,0] - y_current[j,0]) * (y[j,0] - y_current[j,0]);
    }
    cost = cost / n; 
    sum_c = sum_col(x .* (y - y_current));
    w_gradient = -(2.0/n) * sum_c[0,0];
    sum_c = sum_col(y - y_current);
    b_gradient = -(2.0/n) * sum_c[0,0];
    w_current = w_current - (learning_rate * w_gradient);
    b_current = b_current - (learning_rate * b_gradient);
  }
  print("Approach 2: Solving equation by gradient descent:");
  print("After 250000 iterations, the final cost is");
  print(cost);
  print("W is");
  print(w_current);
  print(b_current); 
}

func int main() { 
  matrix x <3, 1> = [1;2;3];
  matrix y <3, 1> = [4;5;6];
  matrix ones <3, 1> = fill(3, 1, 1.0);
  matrix x_new <3, 2>;
  matrix w <2, 1>;
  matrix y_lr <3, 1>;
  
  /* analytical solution for linear regression */
  /* x_pre * w_pre + b = y 
     [x_pre : ones] * w = y
     x * w = y
     x' * x * w = x' * y
     w = inv(x' * x) * x' * y */
  x_new = [x, ones];
  w = calculate_weight(x_new, y);
  y_lr = x_new * w;

  print("Given X and Y, linear regression is finding a w so that Y = W * X");
  print("Approach 1: Solving equation analytically:");
  print("W is");
  print(w);
  print("The predicted y is:");
  print(y_lr);

  /* gradient descent for linear regreesion */
  gradient_descent(x, y);
  return 0;
}
