import numpy as np
import scipy.sparse as sps

from porepy.numerics.ad.forward_mode import Ad_array
import porepy.numerics.ad.functions as af

"""
Explanation: Introduction
This tutorial gives a short overview of the AD-module included in PorePy. For an example where the AD module has been used to solve non-linear compressible flow, see the tutorial:  "compressible_flow_with_automatic_differentiation"
End of explanation
"""


x = Ad_array(2, 1)

"""
Explanation: Scalar AD-variables
We initiate a variable $x = 2$ by giving a pair (val, jac) to the Ad_array class. val is the value at which the function will be evaluated and jac =1 since $\frac{d x}{dx} = 1$.
End of explanation
"""


y = x**2 + 3

"""
Explanation: We can now define a function $y=x^2 + 3$
End of explanation
"""


print('y value is: ', y.val)
print('dy/dx is: ', y.jac)

"""
Explanation: To obtain the function value and the derivative we can call .val and .jac
End of explanation
"""


h = af.exp(y)
print('h value is: ', h.val)
print('dh/dx is: ', h.jac)

"""
Explanation: $y$ is also an AD variable as a function of $x$. We can use it to declare further functions, e.g., $h(x) = e^{y(x)}$. To take the exponential of an Ad_array we need to call the exponential function found in the AD module
End of explanation
"""


y = Ad_array(7, 4)
h = af.exp(y)
print('h value is: ', h.val)
print('dh/dx is: ', h.jac)

"""
Explanation: If we knew the value and jacobian of $y$ we could alternatively skip initiating $x$ and initiate $y$ directly:
End of explanation
"""


x = Ad_array(np.array([1,2,3]), sps.diags([1,1,1]))

"""
Explanation: Arrays of AD-variables
The Ad_array class also support arrays.
End of explanation
"""


A = sps.csc_matrix(np.array([[0,2,3],[4,0,6],[7,8,0]]))
y = A*x  + x**2

print('Analytic y value: ')
print(np.array([14, 26, 32]))
print('Analytic y jacobian:')
print(np.array([[2,2,3],[4,4,6],[7,8,6]]),'\n')
print('Ad y value: ')
print(y.val)
print('Ad y jacobian:')
print(y.jac.A)


"""
Explanation: As for the scalar case, it is straight forward to define functions using normal Python programming. Let us declare the function
$$y = Ax + x^2$$
which has the jacobian
$$ J(y) = A + 2 \text{diag}(x)$$
With this notation we mean $x^2 = [x_1^2, x_2^2, x_3^2]$, and $\text{diag}(x)$ is a matrix with $x$ on the diagonal and zeros elsewhere.
End of explanation
"""


