from copy import deepcopy
from scipy.special import binom
import numpy as np
import time
from numpy.polynomial.polynomial import Polynomial


class Function:
    def __iadd__(self, _f):
        f1 = deepcopy(self)
        return FunctionSum(f1, _f)

    def __imul__(self, _f):
        f1 = deepcopy(self)
        return FunctionScalarMul(f1, _f)

    def compose(self, _f):
        f1 = deepcopy(self)
        return FunctionComposition(f1, _f)

    def __add__(self, _f):
        f1 = deepcopy(self)
        return FunctionSum(f1, _f)

    def deriv_once(self):
        return self.deriv()

    def deriv(self, _deg=1):
        raise NotImplementedError()


class ConstFunction(Function):
    def __init__(self, _val, _exec_time, _dim):
        self.dim_ = _dim
        self.val_ = _val
        self.T_ = _exec_time

    def __call__(self, _t):
        _t = np.atleast_1d(_t)
        return np.ones((self.dim_, len(_t))) * self.val_

    def deriv(self, _deg=1):
        if _deg == 0:
            return self
        return ZeroFunction(self.T_, self.dim_)


class ZeroFunction(ConstFunction):
    def __init__(self, _exec_time, _dim):
        ConstFunction.__init__(self, 0.0, _exec_time, _dim)

    def __call__(self, _t):
        _t = np.atleast_1d(_t)
        return np.zeros((self.dim_, len(_t)))


class OneFunction(ConstFunction):
    def __init__(self, _dim):
        ConstFunction.__init__(self, 1.0, _dim)


class FunctionSum(Function):
    def __init__(self, _f1, *args):
        self.dim_ = _f1.dim_
        self.T_ = _f1.T_

        self.function_array_ = []

        for fun in (_f1, ) + args:
            assert fun.dim_ == self.dim_
            assert np.abs(fun.T_ - self.T_) < 1.0e-4
            if isinstance(fun, FunctionSum):
                self.function_array_ += fun.function_array_
            else:
                self.function_array_.append(deepcopy(fun))

    def __call__(self, _t):
        result = self.function_array_[0](_t)
        for fun in self.function_array_[1:]:
            result += fun(_t)
        return result

    def deriv(self, _deg=1):
        if _deg == 0:
            return self
        fun_ders = [fun.deriv(_deg) for fun in self.function_array_]
        return FunctionSum(*fun_ders)


class FunctionScalarMul(Function):
    def __init__(self, _f1, _scalarfun):
        assert _scalarfun.dim_ == 1
        self.dim_ = _f1.dim_
        self.f1_ = deepcopy(_f1)
        self.scalarfun_ = deepcopy(_scalarfun)

        self.T_ = _f1.T_
        assert np.abs(_f1.T_ - _scalarfun.T_) < 1.0e-4

    def __call__(self, _t):
        _t = np.atleast_1d(_t)
        f1 = self.f1_(_t)
        scalar = self.scalarfun_(_t)
        return np.einsum('i,ij ->ij', scalar.ravel(), f1)

    def deriv(self, _deg=1):
        if _deg == 0:
            return self
        result = self
        for _ in range(_deg):
            result = result.deriv_once()

        return result

    def deriv_once(self):
        fun1 = self.f1_
        sfun = self.scalarfun_

        term_1 = FunctionScalarMul(fun1.deriv(), sfun)
        term_2 = FunctionScalarMul(fun1, sfun.deriv())

        return FunctionSum(term_1, term_2)


class FunctionComposition(Function):
    def __init__(self, _f1, _sfun):
        self.f1_ = _f1
        assert _sfun.dim_ == 1
        self.scalarfun_ = _sfun
        self.T_ = _sfun.T_

        assert _sfun(_sfun.T_) <= _f1.T_ + 1.0e-7,\
            'Domain of the parametrization' +\
            ' is larger than the domain of the funciton'
        self.dim_ = _f1.dim_

    def __call__(self, _t):
        return self.f1_(self.scalarfun_(_t).ravel())

    def deriv(self, _deg=1):
        if _deg == 0:
            return self
        result = self
        for _ in range(_deg):
            result = result.deriv_once()

        return result

    def deriv_once(self):
        f1 = FunctionComposition(self.f1_.deriv(), self.scalarfun_)
        f2 = self.scalarfun_.deriv()
        return FunctionScalarMul(f1, f2)


class FunctionConcat(Function):
    def __init__(self, _f1, _f2):
        self.T_ = _f2.T_
        self.dim_ = _f1.dim_
        self.func_array_ = [_f1, _f2]
        self.f1_ = deepcopy(_f1)
        self.f2_ = deepcopy(_f2)
        assert _f1.dim_ == _f2.dim_

    def __call__(self, _time):
        _time = np.atleast_1d(_time)
        time_prev = _time[_time < self.f1_.T_]
        time_post = _time[_time >= self.f1_.T_]
        result = np.vstack([
            self.f1_(time_prev).reshape(-1, self.dim_),
            self.f2_(time_post).reshape(-1, self.dim_)
        ])
        return result

    def deriv(self, _ord=1):
        return FunctionConcat(self.f1_.deriv(_ord), self.f2_.deriv(_ord))


class ScalarPolWrapper(Function):
    def __init__(self, _pol):
        self.T_ = _pol.domain[-1]
        self.dim_ = 1
        self.pol_ = _pol

    def __call__(self, _time):
        _time = np.atleast_1d(_time)
        return self.pol_(_time).reshape(-1, 1)

    def deriv(self, _ord=1):
        return ScalarPolWrapper(self.pol_.deriv(_ord))


class Identity(ScalarPolWrapper):
    def __init__(self, _exec_time):
        pol = Polynomial(
            [0, 1], domain=[0, _exec_time], window=[0, _exec_time])
        ScalarPolWrapper.__init__(self, pol)
