import sys
import petsc4py
import numpy as np

petsc4py.init(sys.argv)
from petsc4py import PETSc

def fill_function(tao, x, f, points_original):
    x_values = x.getArray(readonly=True)
    diff_vectors = points_original.ravel() - x_values
    f_values = np.square(diff_vectors)
    # f_values[10] = 0
    f.setValues(np.arange(f_values.size), f_values)
    f.assemble()

def fill_jacobian(tao, x, J, points_original):
    x_values = x.getArray(readonly=True)
    points_original_flat = points_original.ravel()
    deriv_values = -2*(points_original_flat - x_values)
    # deriv_values[10] = 0
    for i in range(x_values.size):
        J.setValue(i, i, deriv_values[i])
    # print(deriv_values)
    J.assemble()

# ---------------------------------------------------------------------------------------------

if __name__ == '__main__':
    # Initialize original grid points
    grid_dim = 10
    grid_spacing = 100
    num_points = grid_dim * grid_dim
    points_original = np.zeros(shape=(num_points, 2), dtype=np.float64)
    for i in range(grid_dim):
        for j in range(grid_dim):
            points_original[i*grid_dim+j] = (i*grid_spacing, j*grid_spacing)

    # Compute transformed grid points
    affine_mat = np.array([[-0.5, -0.86, 100], [0.86, -0.5, 100]]) # createAffineMatrix(120, 1, 1, 100, 100)
    points_transformed = np.matmul(affine_mat[:2,:2], points_original.T).T + affine_mat[:2,2]

    # Initialize PETSc objects
    num_unknown = points_transformed.size
    mat_shape = (num_unknown, num_unknown)
    A = PETSc.Mat()
    A.createAIJ(size=mat_shape, comm=PETSc.COMM_WORLD)
    A.setUp()
    x, f = A.createVecs()

    options = PETSc.Options()
    options.setValue("-tao_monitor", "")
    options.setValue("-tao_view", "")
    options.setFromOptions()

    tao = PETSc.TAO()
    tao.create(PETSc.COMM_WORLD)
    tao.setType("brgn")
    tao.setFromOptions()

    tao.setResidual(fill_function, f, args=(points_original,))
    tao.setJacobian(fill_jacobian, A, args=(points_original,))

    x.setArray(points_transformed.ravel())
    tao.setInitial(x)
    tao.solve(x)