// JepNumpy.java
// Andrew Davison, ad@coe.psu.ac.th, July 2023

/* Calls the Numpy module in Python. This is a 
   third-part module for fast array manipulation
   (https://numpy.org/).

   Examples include:
     * making matricies
     * matrix product
     * maximums
     * means
     * sorting
     * array resizing 

  Usage:
     > compile JepNumpy
     > run JepNumpy 

  Based on Python test code in numpyTests.py
*/

import java.util.*;

import jep.*;
import jep.python.*;


public class JepNumpy 
{

  public static void main(String[] args) throws JepException 
  {
    System.out.println("Starting Jep...");
    Interpreter jep = new SharedInterpreter();

    System.out.println("\nMake matricies ----------");
    float[] f = new float[] { 1.0f, 2.1f, 3.3f, 
                              4.5f, 5.6f, 6.7f };
    NDArray<float[]> nd = new NDArray<>(f, 3, 2);
    jep.set("nd", nd);
    jep.exec("print('nd=\\n', nd)");
    jep.exec("print('nd[0]=', nd[0])");

    System.out.println("\nProducts ----------");

    jep.exec("import numpy as np");

    jep.exec("A = np.array([[1, 1],[0, 1]])");
    jep.exec("B = np.array([[2, 0],[3, 4]])");
    jep.exec("print('A@B=\\n', A@B )");

    NDArray<?> ndProd = (NDArray<?>)jep.getValue("A.dot(B)");
    int[] prod = (int[]) ndProd.getData();
    System.out.println("prod = " + Arrays.toString(prod));

    int[][] prod2 = PyUtils.reshapeInt2D(prod, 2, 2);
    System.out.println("prod reshaped = ");
    PyUtils.printInt2D(prod2);

    NDArray<int[]> iarray = new NDArray<>(
                new int[] {1, 2, 1, 1},
                new int[] {2, 2});  // dimensions
    // System.out.println("iarray: " + iarray);
    prod = (int[]) iarray.getData();
    System.out.println("prod: " + Arrays.toString(prod));
    jep.set("iarray", iarray);
    jep.exec("print('iarray=\\n',iarray)");
    jep.exec("print('A@iarray=\\n',A@iarray)");

    jep.exec(
         "matrix1 = [\n" +
         "    [3, 4, 2],\n" +
         "    [5, 1, 8],\n" +
         "    [3, 1, 9]\n" +
         "]\n" +
         "matrix2 = [\n" +
         "    [3, 7, 5],\n" +
         "    [2, 9, 8],\n" +
         "    [1, 5, 8]\n" +
         "]\n" +
         "print('matrix dot =\\n', np.dot(matrix1, matrix2))"
      );

    System.out.println("\nMaximums -----");
    // 5x4 array holding 0 to 19 in row order
    jep.exec("arr1 = np.arange(20).reshape(5, 4)");
    jep.exec("print('arr1=\\n', arr1)");

    // find max of entire array
    jep.exec("print('arr1 max =', np.argmax(arr1))");

    // find max index for each row
    NDArray<?> nda1 = (NDArray<?>)jep.getValue(
                        "np.argmax(arr1, axis=1)");
    long[] rarr1 = (long[]) nda1.getData();
    System.out.println("row idx maxs = " + 
                              Arrays.toString(rarr1));


    System.out.println("\nMeans -----");
    jep.exec(
          "list = [\n" +
          "    np.array([3, 2, 8, 9]),\n" +
          "    np.array([4, 12, 34, 25, 78]),\n" +
          "    np.array([23, 12, 67])\n" +
          "]\n" +
          "result = []\n" +
          "for i in range(len(list)):\n" +
          "    result.append(np.mean(list[i]))\n"
       );

    List<?> means = (List<?>) jep.getValue("result");
    System.out.println("Means = " + means);
    System.out.println("2nd mean = " + ((double)means.get(1)));


    System.out.println("\nSorting -----");
    NDArray<int[]> arr2 = new NDArray<>(
                new int[] {3, 7, 1,
                          10, 3, 2,
                           5, 6, 7},
                new int[] {3, 3});  // dimensions
    jep.set("arr2", arr2);
    jep.exec("print('arr2=\\n',arr2)");

    // Sort the whole array
    jep.exec("print('Sorted=\\n',np.sort(arr2))");
    
    // Sort along each column
    NDArray<?> sortND = (NDArray<?>)jep.getValue(
                        "np.sort(arr2, axis=0)");
    int[] dims = sortND.getDimensions();
    System.out.println("Dimensions = " + Arrays.toString(dims));
    int[] arr3 = (int[]) sortND.getData();
    System.out.println("Sorted down each column = " + 
                                Arrays.toString(arr3));

    int[][] arr3ColSort = PyUtils.reshapeInt2D(arr3, dims);
    System.out.println("Reshaped (into rows) = ");
    PyUtils.printInt2D(arr3ColSort);


    System.out.println("\nAdd row -----");
    jep.exec(
       "arr4 = np.array([\n" +
       "    [3, 2, 8],\n" +
       "    [4, 12, 34],\n" +
       "    [23, 12, 67]\n" +
       "])"
     );

    NDArray<int[]> nRow = new NDArray<>(
                new int[] {2, 1, 8},
                new int[] {1, 3});  // dimensions
    jep.set("nRow", nRow);
    jep.exec("arr4Row = np.vstack((arr4, nRow))");
    jep.exec("print('With new row =\\n', arr4Row)");

    jep.close();
  } // end of main()

} // end of JepNumpy class