# cPy - exercises with numpy
# provided by Anthony Vogelaar, Perth Scotland
# rev 0.10  2024-07-11  initial release

import  os;                 os.system('clear')
import  time
import  string
import  numpy               as  np
import  matplotlib.pyplot   as  plt

# ========================= Python vs NumPy in multiply ==================

def test1():
    L = 1000000

    # Python
    a = list(range(L))
    b = list(range(L))
    t1 = time.time()
    c = []
    for x, y in zip(a, b):
        c.append(x*y)
    t2 = time.time() - t1; print('Python:', t2)

    # NumPy
    a = np.arange(L)
    b = np.arange(L)
    t3 = time.time()
    c = a * b
    t4 = time.time() - t3; print('NumPy:', t4)
    print('Python / Numpy:', t2/t4)
    # print(c)

# ========================= Python vs NumPy in broadcast ==================
    
def test2():
    L = 1000000

    # Python
    a = list(range(L))
    t1 = time.time()
    for i in range(len(a)):
        a[i] *= 10
    t2 = time.time() - t1; print('Python:', t2)
    # print(a)

    # NumPy
    a = np.arange(L)
    t3 = time.time()
    a *= 10
    t4 = time.time() - t3; print('Numpy:', t4)
    print('Python / Numpy:', t2/t4)
    # print(a)

# =========================== User manual ============================
     
# URL: https://numpy.org/doc/stable/user/index.html#user
#    1.1 Why is Numpy fast
#    1.2 Who else uses NumPy

# ============================= Basics ===============================

# ndarray = homogeneous multidimensional
# indexed by a tuple. In NumPy a dimension is called axEs

# $ . py & python
# import numpy as np
# a 3D point a = np.array([1, 2, 3])
#    has 1 axis, three elements in the axis --> the length is 3
#
# b = np.array([[1., 0., 1.], [0., 1., 2.]])
# ndarray alias array  numpy.array != array.array (1 dimensional)

# array attributes:
def test3():
    b = np.array([[1., 0., 1.], 
                  [0., 1., 2.]])
    print('ndarray.ndim    :', b.ndim)
    # print('ndarray.shape   :', b.shape)
    # print('ndarray.size    :', b.size)              # total number of elements
    # print('ndarray.dtype   :', b.dtype)             # see page 90. int=numpy.int_,  bool=numpy.bool, 
    #                                                 # float=numpy.float64, complex=numpy.complex128
    # print('ndarray.itemsize:', b.itemsize)
    # print('ndarray.data    :', b.data)              # buffer address
    # print('type(b)         :', type(b))         

# ========================== Array creation  =========================

def test4():
    b = np.array([1.2, 3.5, 5.1]);                      print(b.dtype, b)
    # b = np.array([1, 3.5], dtype=int);                  print('\n', b.dtype, b)
    # b = np.array([[1, 2], [3, 4+2j]], dtype=complex);   print('\n', b.dtype, b)
    # b = np.zeros((3, 2));                               print('\n', b.dtype, b)
    # b = np.ones((2, 3, 4), dtype=np.int16);             print('\n', b.dtype, b)
    # b = np.empty((2, 2));                               print('\n', b.dtype, b)
    # b = np.arange(10, 30, 5);                           print('\n', b.dtype, b)
    # b = np.arange(10, 30, 5.1);                         print('\n', b.dtype, b)
    # b = np.linspace(1, 10, 10);                         print('\n', b.dtype, b)
    # x = np.linspace(0, 2*np.pi, 5);                     print('\n', x.dtype, x)
    # f = np.sin(x);                                      print('\n', f)

# ============================= Functions ============================

# Arithmetic operators on arrays apply elementwise. A new array is created and filled with the result.
# + - * / ** etc.
# acting in place: +=, -=, *=, /= etc.

# unary operators are implemented as methods
def test5():
    rg = np.random.default_rng(1)
    a = rg.random((2, 3));  print(a)
    # print('a.sum()      :', a.sum())
    # print('a.sum(axis=0):', a.sum(axis=0))          # sum of each column
    # print('a.sum(axis=1):', a.sum(axis=1))          # sum of each row
    # print('a.min()      :', a.min())
    # print('a.max()      :', a.max())

# universal functions
# URL: https://numpy.org/doc/stable/reference/ufuncs.html

# ========================= Indexing / Slicing =======================

# One-dimensional arrays can be indexed, sliced and iterated over
# Multidimensional arrays can have one index per axis. These indices are given in a tuple    
def test6():
    a = np.arange(10)**3;       print(a)
    print('a[2]   :', a[2])
    # print('a[-2]  :', a[-2])
    # print('a[2:5] :', a[2:5])
    # print('a[:8:2]:', a[:8:2])
    # print('a[::-1]:', a[::-1])
    # for i in a: print(f'{i**(1 / 3.):.3f}', end='  ')
    # print('\n')

    # def f(x, y): return 10 * x + y
    # b = np.fromfunction(f, (5, 4), dtype=int); print(b, '\n')
    # print('b[2, 3]  :', b[2, 3])
    # print('b[0:5, 1]:', b[0:5, 1])      # each row in the second column
    # print('b[1:3, :]:', b[1:3, :])      # each column in the second and third row
    # print('b[-1]    :', b[-1])          # missing indices are considered complete slices
    # print('\nfor row in b:')
    # for row in b: print(row)
    # print('\nfor element in b.flat:')
    # for element in b.flat: print(element, end=' ')
    print()

# ========================= Shape manipulation =======================
    
def test7():
    rg = np.random.default_rng(None)
    a = np.floor(10 * rg.random((3, 4))); print(a)
    # print('a.ravel()      :', a.ravel())                    # .flat = copy, ravel=view
    # print('a.reshape(2, 6):', '\n', a.reshape(2, 6))
    # print('a.T            :', '\n', a.T)                    # transposed
    print('np.hsplit(a, 2):', '\n', np.hsplit(a, 2))

# ============================ Numpy operations ========================

def test8():

    scr1 = np.array(list(string.ascii_uppercase[:20])).reshape(4, 5)

    def test0():
        os.system('clear')
        print(scr1)


    def test1():
        for i in range(-4, 5):
            os.system('clear')
            print(np.tril(scr1, i))
            time.sleep(2)


    def test2():
        for i in range(5, -4, -1):
            os.system('clear')
            print(np.triu(scr1, i))
            time.sleep(2)

    scr2 = np.rot90(scr1)

    def test3():
        for i in range(-5, 4):
            os.system('clear')
            print(np.rot90(np.tril(scr2, i), axes=(1,0)))
            time.sleep(2)


    def test4():
        for i in range(4, -5, -1):
            os.system('clear')
            print(np.rot90(np.triu(scr2, i), axes=(1,0)))
            time.sleep(2)

    scr3 = np.ones((4, 5), dtype=bool)

    def test5():
        for i in range(-4, 5):
            os.system('clear')
            # print(np.tril(scr3, i))
            print(np.logical_not(np.tril(scr3, i)))
            time.sleep(2)


    test5()

# ============================ Matrices ================================
    
# see: https://www.mathsisfun.com/algebra/matrix-multiplying.html

def test9():

    def test1():

        # for matrix multiplication: the number of columns in the first matrix
        #  must be equal to the number of rows in the second matrix. 
        m1 = np.array([[1, 2, 3],
                       [4, 5, 6]])
        m2 = np.array([[10, 20],
                       [30, 40],
                       [50, 60]])
        
        t1 = time.time()
        m3 = m1 @ m2
        t2 = time.time()

        print(m1) 
        print('\n', m2)
        print('\n', m3)
        print(f'duration [2*3] @ [3*2]: {(t2-t1)*1e3:.3f} mSec\n')


    def test2():
        n = 3
        m1 = np.arange(1,  2*n+1).reshape(2, n)
        print(m1)
        m2 = 10 * np.arange(1, 2*n+1).reshape(n, 2)
        print(m2)

        t1 = time.time()
        m3 = m1 @ m2
        t2 = time.time()

        print(m1) 
        print('\n', m2)
        print('\n', m3)
        print(f'duration [2*3] @ [3*2]: {(t2-t1)*1e3:.3f} mSec\n')


    def test3():
        t1 = time.time()
        n = 65_0000_000     # 62e6 * 2 * 8 * 3 = 29 GB
        m1 = np.arange(1,  2*n+1).reshape(2, n)
        m2 = 10 * np.arange(1, 2*n+1).reshape(n, 2)

        t2 = time.time()
        m3 = m1 @ m2
        t3 = time.time()
        print(m3)

        print(f'preparation: {(t2-t1)*1e3:.0f} mSec')
        print(f'execute    : [2*{n}] @ [{n}*2]: {(t3-t2)*1e3:.0f} mSec\n')


    test3()

# GPU usage:
# https://pypi.org/project/pycuda/
# https://developer.nvidia.com/cunumeric
    
# ======================= One dimensional plot =========================
    
def test10():
    a = np.array([2, 1, 5, 7, 4, 6, 8, 14, 10, 9, 18, 20, 22])
    plt.plot(a)
    plt.show()

# ======================= Two dimensional plot =========================
    
def test11():
    x = np.linspace(0, 4*np.pi, 200)
    # y = np.linspace(0, 10, 20)
    y = np.sin(x)
    print(x, y)
    plt.plot(x, y, 'purple') # line
    # plt.plot(x, y, 'o')
    plt.show()

# ======================= Three dimensional plot =========================
    
def test12():
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    x = np.arange(-5, 5, 0.05);     #  print('x:', x.size, '\n', x)
    y = np.arange(-5, 5, 0.05);     # print('y:', y.size, '\n', y)
    x, y = np.meshgrid(x, y);       # print('x, y:', x.size, '\n', x, y)
    r = np.sqrt(x**2 + y**2);       # print('r:', r.size, '\n', r)
    z = np.sin(r);                  # print('z:', z.size, '\n', z)
    print('size:', z.size)
    # ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='viridis')
    ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='plasma')
    plt.show()

# ============================= Tests ================================

test12()

# =============================== End =================================
