"""
Small experiment with TPSA.

Sergey Tomin.
"""

import numpy as np
from scipy.special import binom
from scipy.special import factorial2, factorial
import copy
import time
from numba import jit


@jit(nopython=True)
def nb_factorial(nn):
    res = 1
    for ii in range(2, nn + 1):
        res *= ii
    return res
print(nb_factorial(10))


@jit(nopython=True)
def nb_binom(n, k):
    return nb_factorial(n)/(nb_factorial(k) * nb_factorial(n - k))


@jit(nopython=True)
def mult(res, a, b, order):
    """
    Multiplication res = a*b

    :param res: result array
    :param a: array [order, order]
    :param b: array [order, order]
    :param order: size of the array
    :return:
    """
    res[0, 0] = a[0, 0] * b[0, 0]
    for m in range(order):
        for n in range(order):
            tmp = 0.
            for s in range(m+1):
                for t in range(n+1):
                    tmp += a[s, t] * b[m - s, n - t] * nb_binom(m, s) * nb_binom(n, t)
            res[m, n] = tmp
    return res

@jit(nopython=True)
def mult3(res, a, b, order):
    """
    Multiplication res = a*b

    :param res: result array
    :param a: array [order, order]
    :param b: array [order, order]
    :param order: size of the array
    :return:
    """
    res[0, 0, 0] = a[0, 0, 0] * b[0, 0, 0]
    for p in range(order):
        for m in range(order):
            for n in range(order):
                tmp = 0.
                for s in range(m+1):
                    for t in range(n+1):
                        tmp += a[s, t] * b[m - s, n - t] * nb_binom(m, s) * nb_binom(n, t)
                res[m, n] = tmp
    return res


@jit(nopython=True)
def inverse(a, order):
    """
    Inversion: b = 1/a

    :param a:
    :return: b
    """
    #M, N = np.shape(a)
    b = np.zeros((order, order))
    b[0, 0] = 1. / a[0, 0]
    for m in range(order):
        for n in range(order):
            if n == 0 and m == 0:
                continue
            tmp = 0.
            for s in range(m + 1):
                for t in range(n + 1):
                    tmp += a[s, t] * b[m - s, n - t] * nb_binom(m, s) * nb_binom(n, t)
            b[m, n] = -tmp / a[0, 0]
    return b


def tpsa_exp(a):
    a_ceil = a.ceil()
    tmp = 0
    pow_a = 1
    for k in range(2*a.order-1):

        tmp = tmp + pow_a/factorial(k)
        pow_a =  a_ceil * pow_a
    return np.exp(a.v[0, 0])*tmp


def tpsa_sin(a):
    a_ceil = a.ceil()
    pow_array = [1, 1*a_ceil]
    sin_terms = 0
    cos_terms = 0

    for k in range(0, 2*a.order-1):
        sin_terms = sin_terms + (-1) ** k / factorial(2 * k) * pow_array[0]
        cos_terms = cos_terms + (-1) ** k / factorial(2 * k + 1) * pow_array[1]
        pow_array[0] = pow_array[1] * a_ceil
        pow_array[1] = pow_array[0] * a_ceil
    return np.sin(a.v[0, 0]) * sin_terms + np.cos(a.v[0, 0]) * cos_terms


def tpsa_cos(a):
    a_ceil = a.ceil()
    pow_array = [1, 1*a_ceil]
    sin_terms = 0
    cos_terms = 0

    for k in range(0, 2*a.order-1):
        sin_terms = sin_terms + (-1) ** k / factorial(2 * k) * pow_array[0]
        cos_terms = cos_terms + (-1) ** k / factorial(2 * k + 1) * pow_array[1]
        pow_array[0] = pow_array[1] * a_ceil
        pow_array[1] = pow_array[0] * a_ceil
    return np.cos(a.v[0, 0])* sin_terms - np.sin(a.v[0, 0]) * cos_terms


def tpsa_ln(a):
    a_floor = a.floor()
    a_floor.v[0, 0] = np.log(a.v[0, 0])
    a_ceil = a.ceil()
    a_ceil_n = a_ceil/a.v[0, 0]
    tmp = 0
    pow = 1
    for k in range(1, 2*a.order-1):
        pow = pow * a_ceil_n
        tmp = tmp + (-1)**(k+1)/k * pow
    return a_floor + tmp


def tpsa_sqrt(a):
    a_ceil = a.ceil()
    a_ceil_n = a_ceil/a.v[0, 0]
    tmp = 1 + 0.5 * a_ceil_n
    pow = a_ceil_n
    for k in range(2, 2*a.order-1):
        pow = pow * a_ceil_n
        tmp = tmp - (-1)**(k) * factorial2(2 * k - 3)/factorial2(2*k) * pow
    return np.sqrt(a.v[0, 0])*tmp


class PSC:
    def __init__(self, v):
        self.v = np.array(v, dtype=np.float64)

    def floor(self):
        m = copy.deepcopy(self)
        m.v[1:] *= 0
        return m

    def ceil(self):
        m = copy.deepcopy(self)
        m.v[0] = 0
        return m

    @property
    def order(self):
        return len(self.v)

    def __mul__(self, other):
        m = copy.deepcopy(self)
        if other.__class__ == PSC:
            if m.order != other.order:
                raise ValueError("not equal dimensions")
            other_v = np.array(other.v)
        elif other.__class__ == np.ndarray:
            if len(other) != m.order:
                raise ValueError("not equal dimensions")
            other_v = np.copy(other)
        else:
            other_v = np.zeros(len(self.v))
            other_v[0] = other
        m.v[0] = self.v[0]*other_v[0]
        for i in range(1, len(m.v)):
            tmp = 0
            for n in range(i+1):
                tmp += binom(i, n)*self.v[n]*other_v[i - n]
            m.v[i] = tmp
        return m

    def __rmul__(self, other):
        return self.__mul__(other)

    def __add__(self, other):
        m = copy.deepcopy(self)
        if other.__class__ == PSC:
            v = np.array(other.v)
        else:
            v = np.zeros(len(self.v))
            v[0] = other
        m.v = np.array(self.v) + np.array(v)
        return m

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        m = copy.deepcopy(self)
        if other.__class__ == PSC:
            v = np.array(other.v)
        else:
            v = np.zeros(len(self.v))
            v[0] = other
        m.v = np.array(self.v) - np.array(v)
        return m

    def __rsub__(self, other):
        return self.__sub__(other)

    def __truediv__(self, other):
        if other.__class__ == PSC:
            b = np.zeros(len(other.v))
            b[0] = 1/other.v[0]
            for n in range(1, len(other.v)):
                tmp = 0
                for k in range(1, n+1):
                    tmp += binom(n, k)*other.v[k]*b[n-k]
                b[n] = -tmp/other.v[0]
            return self.__mul__(b)
        else:
            return self.__mul__(1./other)

    def __rtruediv__(self, other):
        return self.__truediv__(other)

    def __str__(self):
        return str(self.v)


class PSC2(object):

    def __init__(self, a=0., size=2):
        self.v = np.zeros((size, size))
        self.v[0, 0] = a

    def floor(self):
        m = copy.deepcopy(self)
        m.v[1:, 1:] *= 0
        return m

    def ceil(self):
        m = copy.deepcopy(self)
        m.v[0, 0] = 0
        return m

    @property
    def order(self):
        return np.shape(self.v)[0]

    def __mul__(self, other):
        obj = copy.deepcopy(self)
        if isinstance(other, PSC2):
            if obj.order != other.order:
                raise ValueError("not equal dimensions")
            other_v = np.array(other.v)
        elif other.__class__ == np.ndarray:
            if len(other) != obj.order:
                raise ValueError("not equal dimensions")
            other_v = np.copy(other)
        elif isinstance(other, (int, float, complex)):
            other_v = np.zeros_like(self.v)
            other_v[0, 0] = other
        else:
            raise ValueError("Unknown object")

        M, N = np.shape(obj.v)
        obj.v = mult(obj.v, a=self.v, b=other_v, order=M)
        return obj

    def __rmul__(self, other):
        return self.__mul__(other)

    def __add__(self, other):
        m = copy.deepcopy(self)
        if isinstance(other, PSC2):
            v = np.array(other.v)
        else:
            v = np.zeros(np.shape(self.v))
            v[0, 0] = other
        m.v = np.array(self.v) + np.array(v)
        return m

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        m = copy.deepcopy(self)
        if isinstance(other, PSC2):
            v = np.array(other.v)
        else:
            v = np.zeros_like(self.v)
            v[0, 0] = other
        m.v = np.array(self.v) - np.array(v)
        return m

    def __rsub__(self, other):
        return self.__sub__(other)

    def one_div_obj_py(self, obj):
        b = np.zeros_like(self.v)
        b[0, 0] = 1. / obj.v[0, 0]
        M, N = np.shape(self.v)
        for m in range(M):
            for n in range(N):
                if n == 0 and m == 0:
                    continue
                tmp = 0.
                for s in range(m + 1):
                    for t in range(n + 1):
                        tmp += obj.v[s, t] * b[m - s, n - t] * binom(m, s) * binom(n, t)

                b[m, n] = -tmp / obj.v[0, 0]
        obj_c = copy.deepcopy(obj)
        obj_c.v = b
        return obj_c

    def one_div_obj(self, obj):
        obj_c = copy.deepcopy(obj)
        obj_c.v = inverse(obj.v, order=np.shape(obj.v)[0])
        return obj_c

    def __truediv__(self, other):
        if isinstance(other, PSC2):
            b = self.one_div_obj(other)
            return self.__mul__(b)
        else:
            return self.__mul__(1./other)

    def __rtruediv__(self, other):
        obj = copy.deepcopy(self)
        b = self.one_div_obj(obj)
        return b.__mul__(other)

    def __str__(self):
        line = "PSC\n"
        return line + "v = " + str(self.v)


class PSC3(object):

    def __init__(self, a=0., size=2):
        self.v = np.zeros((size, size, size))
        self.v[0, 0, 0] = a

    def floor(self):
        m = copy.deepcopy(self)
        m.v[1:, 1:, 1:] *= 0
        return m

    def ceil(self):
        m = copy.deepcopy(self)
        m.v[0, 0, 0] = 0
        return m

    @property
    def order(self):
        return np.shape(self.v)[0]

    def __mul__(self, other):
        obj = copy.deepcopy(self)
        if isinstance(other, PSC3):
            if obj.order != other.order:
                raise ValueError("not equal dimensions")
            other_v = np.array(other.v)
        elif other.__class__ == np.ndarray:
            if len(other) != obj.order:
                raise ValueError("not equal dimensions")
            other_v = np.copy(other)
        elif isinstance(other, (int, float, complex)):
            other_v = np.zeros_like(self.v)
            other_v[0, 0] = other
        else:
            raise ValueError("Unknown object")

        M, N = np.shape(obj.v)
        obj.v = mult(obj.v, a=self.v, b=other_v, order=M)
        return obj

    def __rmul__(self, other):
        return self.__mul__(other)

    def __add__(self, other):
        m = copy.deepcopy(self)
        if isinstance(other, PSC2):
            v = np.array(other.v)
        else:
            v = np.zeros(np.shape(self.v))
            v[0, 0] = other
        m.v = np.array(self.v) + np.array(v)
        return m

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        m = copy.deepcopy(self)
        if isinstance(other, PSC2):
            v = np.array(other.v)
        else:
            v = np.zeros_like(self.v)
            v[0, 0] = other
        m.v = np.array(self.v) - np.array(v)
        return m

    def __rsub__(self, other):
        return self.__sub__(other)

    def one_div_obj_py(self, obj):
        b = np.zeros_like(self.v)
        b[0, 0] = 1. / obj.v[0, 0]
        M, N = np.shape(self.v)
        for m in range(M):
            for n in range(N):
                if n == 0 and m == 0:
                    continue
                tmp = 0.
                for s in range(m + 1):
                    for t in range(n + 1):
                        tmp += obj.v[s, t] * b[m - s, n - t] * binom(m, s) * binom(n, t)

                b[m, n] = -tmp / obj.v[0, 0]
        obj_c = copy.deepcopy(obj)
        obj_c.v = b
        return obj_c

    def one_div_obj(self, obj):
        obj_c = copy.deepcopy(obj)
        obj_c.v = inverse(obj.v, order=np.shape(obj.v)[0])
        return obj_c

    def __truediv__(self, other):
        if isinstance(other, PSC2):
            b = self.one_div_obj(other)
            return self.__mul__(b)
        else:
            return self.__mul__(1./other)

    def __rtruediv__(self, other):
        obj = copy.deepcopy(self)
        b = self.one_div_obj(obj)
        return b.__mul__(other)

    def __str__(self):
        line = "PSC\n"
        return line + "v = " + str(self.v)



class X(PSC2):
    def __init__(self, a, size):
        PSC2.__init__(self, a=a, size=size)
        self.v[0, 1] = 1.


class Y(PSC2):
    def __init__(self, a, size):
        PSC2.__init__(self, a=a, size=size)
        self.v[1, 0] = 1.


def series(psc, dx, dy):
    f = 0.
    for i in range(psc.order):
        for j in range(psc.order):
            f += dx**j * dy**i * psc.v[i, j]/factorial(i)/factorial(j)
    return f



# Example

a = PSC([0, 1, 0, 0, 0, 0])
x = X(0.11, size=5)
y = Y(0.15, size=5)

print(f"X0 = {x.v[0,0]}, Y0 = {y.v[0,0]}")



#print(a2*a2*a2 + b2*b2*b2 + a2*a2*a2* b2*b2*b2)

#print((2*b2 + a2)/a2)
#print(b2*a2)
#print(tpsa_exp(b2*a2))
#print(tpsa_sqrt(b2*a2))
start = time.time()
#print(tpsa_sin(x*y))
#print(tpsa_sqrt(tpsa_exp(tpsa_sin(x*y) + y)*x))
t = tpsa_sqrt(tpsa_exp(tpsa_sin(x*y) + y)/ x)
print("exec = ", time.time() - start, " sec")
#print(t.__class__)
#print(t.order)
print("coef = ", t)
x0 = 0.1
y0 = 0.2
start = time.time()
print("check with direct calc: ", np.sqrt(np.exp(np.sin(x0*y0) + y0)/x0), series(t, dx=x0 - x.v[0,0], dy=y0 - y.v[0,0]))
print("exec series = ", time.time() - start, " sec")
#print("exp = ", tpsa_exp(2*a)/a)
#print("ln = ", tpsa_ln(a))
#print("sin = ", tpsa_sin(a))
#print("sqrt = ", tpsa_sqrt(a), np.sqrt(a.v[0]), 0.5*1/np.sqrt(a.v[0]), -0.25*a.v[0]**(-3/2))
#print(np.sin(a.v[0]), np.cos(a.v[0]) )