# This code originates from https://gitlab.in2p3.fr/betoule/imageproc by Marc
# Betoule.


import numpy as np
from scipy import sparse

from .gnomonic import inverse_gnomonic_dcenter as igp_dcenter
from .gnomonic import inverse_gnomonic_dradec as igp_dradec
from .gnomonic import inverse_gnomonic_projection as igp


class ParamFunc:
    def __call__(self, x, p, **keys):
        """Function call

        Should implement the evaluation of a parametric and vectorial
        function of x.

        Params:
        -------
        x: ndarray N, m
        p: FitParameters instance
        Return:
        ------
        f: ndarray M, N

        """
        pass

    def derivatives(self, x, p, **keys):
        """Function call and derivatives

        Should implement the evaluation of a parametric and vectorial
        function of x and of its derivatives wrt the variable and the
        parameters.

        Params:
        -------
        x: ndarray N, m
        p: FitParameters instance
        Return:
        ------
        f: ndarray (M,N)
        dfdx: ndarray (M,m,N) the derivatives of f wrt x
        dfdp: the stacked derivatives of f.flat wrt p, matrix of size (NxM, P) in sparse format i,j,val

        """
        pass

    def test_derivatives(self, x, p, **keys):
        Y0, dY0dx, (ii, jj, vals) = self.derivatives(x, p, **keys)
        N, M, nobj = dY0dx.shape
        dYdp = [
            sparse.coo_matrix((val, (ii, jj)), shape=(nobj, len(p.full))).tocsc()
            for val in vals
        ]
        err = []
        D = []
        Dpred = []
        for i in range(M):
            dx = 1e-8
            xdx = x.copy()
            xdx[:, i] += dx
            Yx = self(xdx, p, **keys)
            D.append(Yx - Y0)
            Dpred.append(dY0dx[:, i, :] * dx)
            err.append((Yx - Y0) - (dY0dx[:, i, :] * dx))
        for i in range(len(p.full)):
            p2 = p.copy()
            p2.full[i] += dx
            Yx = self(x, p2, **keys)
            Ypred = np.array(
                [np.array((dx * dY[:, i]).todense()).squeeze() for dY in dYdp]
            )
            D.append(Yx - Y0)
            Dpred.append(Ypred)
            err.append((Yx - Y0) - Ypred)
        return err, D, Dpred

    def stack(self, x, p, **keys):
        """Stack the result of vector function

        Can be useful if we want to use 1 dimensionnal fit routines
        """
        return self(x, p, **keys).flatten()

    def dstackdp(self, x, p, **keys):
        """Return the sparse matrix corresponding to the derivatives of the "stacked" function with respect to the named parameters."""
        X, D, (i, j, vals) = self.derivatives(x, p, **keys)
        M, N, L = D.shape
        ii = np.hstack([i + L * m for m in range(M)])
        jj = np.tile(j, M)
        val = np.hstack(vals)
        return X.flatten(), ii, jj, val


class StackedFunc(ParamFunc):
    def __init__(self, f):
        self.f = f

    def __call__(self, x, p, **keys):
        return self.f.stack(x, p, **keys)

    def get_struct(self):
        return self.f.get_struct()

    def derivatives(self, x, p, **keys):
        X, D, (i, j, vals) = self.f.derivatives(x, p, **keys)
        M, N, L = D.shape
        ii = np.hstack([i + L * m for m in range(M)])
        jj = np.tile(j, M)
        val = np.hstack(vals)
        return (
            X.flatten(),
            np.hstack([D[i, :, :] for i in range(M)])[None, ...],
            (ii, jj, [val]),
        )


class Polynomial(ParamFunc):
    def __init__(self, deg=2, coefsuffix="alpha"):
        self.deg = deg
        self.coefsuffix = coefsuffix
        self.coefnames = [coefsuffix + "x" * i for i in range(self.deg + 1)]

    def get_struct(self):
        return [(n, 1) for n in self.coefnames]

    def __call__(self, x, p):
        y = np.full(len(x), p[self.coefnames[-1]].full)
        for d in self.coefnames[-2::-1]:
            y *= x
            y += p[d].full
        return y

    def derivatives(self, x, p):
        y = np.ones(len(x))
        a = np.arange(len(x))
        d = self.coefnames[0]
        val = [y.copy()]
        i = [a]
        j = [np.full(len(x), p[d].indexof())]
        f = y * p[d].full
        df = 0
        for _d, d in enumerate(self.coefnames[1:], 1):
            i.append(a)
            j.append(np.full(len(x), p[d].indexof()))
            coef = p[d].full
            df += coef * y * _d
            y *= x
            val.append(y.copy())
            f += coef * y
        return f, df, (np.hstack(i), np.hstack(j), np.hstack(val))


class BiPol2D(ParamFunc):
    def __init__(self, deg=2, suffix=["alpha", "beta"], key="ccd", n=1):
        self.deg = deg
        self.suffix = suffix
        self.key = key
        self.coeffs = [(i, j) for j in range(0, deg + 1) for i in range(0, deg + 1 - j)]
        self.coeffnames = [
            suf + "x" * i + "y" * j for i, j in self.coeffs for suf in self.suffix
        ]
        self.n = n

    def get_struct(self):
        return [(name, self.n) for name in self.coeffnames]

    def __call__(self, xy, p, **keys):
        x, y = xy
        k = keys[self.key]
        f = [0] * len(self.suffix)
        for degx, degy in self.coeffs:
            monom = x**degx * y**degy
            for _s, suf in enumerate(self.suffix):
                name = suf + "x" * degx + "y" * degy
                coef = p[name].full[k]
                f[_s] = coef * monom + f[_s]
        return np.array(f)

    def derivatives(self, xy, p, **keys):
        x, y = xy
        k = keys[self.key]
        f = [0] * len(self.suffix)
        df = [np.zeros(xy.shape)] * len(self.suffix)
        i = np.tile(np.arange(len(x)), len(self.coeffnames))
        j = []
        vals = [list() for _s in self.suffix]
        for degx, degy in self.coeffs:
            monom = x**degx * y**degy
            for _s, suf in enumerate(self.suffix):
                name = suf + "x" * degx + "y" * degy
                j.append(p[name].indexof(k))
                coef = p[name].full[k]
                f[_s] = coef * monom + f[_s]
                for pouet, tralala in enumerate(self.suffix):
                    if pouet == _s:
                        vals[_s].append(monom)
                    else:
                        vals[pouet].append(np.zeros(len(monom)))
                df[_s] = (
                    np.array(
                        [
                            degx * coef * x ** (degx - 1) * y**degy,
                            degy * coef * x**degx * y ** (degy - 1),
                        ]
                    )
                    + df[_s]
                )
        for v in range(len(vals)):
            vals[v] = np.hstack(vals[v])
        return np.array(f), np.array(df), (i, np.hstack(j), vals)


class InverseGnomonic(ParamFunc):
    def __init__(self, param_names=["rac_rad", "decc_rad"], key="exp"):
        self.ra0, self.dec0 = param_names
        self.key = key

    def __call__(self, XY, p, **keys):
        piece = keys[self.key]
        rac = p[self.ra0].full[piece]
        decc = p[self.dec0].full[piece]
        X, Y = XY
        return np.array(igp(X, Y, rac, decc))

    def get_struct(self, n):
        return [(name, n) for name in [self.ra0, self.dec0]]

    def derivatives(self, XY, p, **keys):
        piece = keys[self.key]
        rac = p[self.ra0].full[piece]
        decc = p[self.dec0].full[piece]
        X, Y = XY
        N = len(X)
        (dradrac, draddecc), (ddecdrac, ddecddecc) = igp_dcenter(X, Y, rac, decc)
        valra = np.hstack([dradrac, draddecc])
        valdec = np.hstack([ddecdrac, ddecddecc])
        ii = np.tile(np.arange(N), 2)
        jj = np.hstack([p[self.ra0].indexof(piece), p[self.dec0].indexof(piece)])
        (dradX, dradY), (ddecdX, ddecdY) = igp_dradec(X, Y, rac, decc)
        return (
            self(XY, p, **keys),
            np.array([[dradX, dradY], [ddecdX, ddecdY]]),
            (ii, jj, (valra, valdec)),
        )


class D32(ParamFunc):
    def __init__(self, param_names=["rac_rad", "decc_rad"]):
        self.ra0, self.dec0 = param_names

    def __call__(self, XY, p, piece=0):
        X, Y = XY[0, :], XY[1, :]
        return np.array([X, 2 * Y, np.full(len(X), 3 * p["rac_rad"].full[piece])])

    def derivatives(self, XY, p, piece=0):
        N = XY.shape[1]
        ii = np.tile(np.arange(N), 1)
        jj = np.full(N, p["rac_rad"].indexof(piece))
        val = np.full(N, 3)
        return (
            self(XY, p, piece),
            np.array(
                [
                    [np.ones(N), np.zeros(N)],
                    [np.zeros(N), np.full(N, 2)],
                    [np.zeros(N), np.zeros(N)],
                ]
            ),
            (ii, jj, (np.zeros(N), np.zeros(N), val)),
        )


class D23(ParamFunc):
    def __init__(self, param_names=["rac_rad", "decc_rad"]):
        self.ra0, self.dec0 = param_names

    def __call__(self, XY, p, piece=0):
        X, Y, Z = XY
        return np.array([p["decc_rad"].full[piece] * X, Y * Z])

    def derivatives(self, XY, p, piece=0):
        N = XY.shape[1]
        X, Y, Z = XY
        ii = np.tile(np.arange(N), 1)
        jj = np.full(N, p["decc_rad"].indexof(piece))
        val = np.full(N, 3)
        return (
            self(XY, p, piece),
            np.array(
                [
                    [np.full(N, p["decc_rad"].full[piece]), np.zeros(N), np.zeros(N)],
                    [np.zeros(N), Z, Y],
                ]
            ),
            (ii, jj, (val, np.zeros(N))),
        )


class CompoundFunc(ParamFunc):
    def __init__(self, f, g):
        self.f = f
        self.g = g

    def __call__(self, x, p, **keys):
        return self.f(self.g(x, p, **keys), p, **keys)

    def derivatives(self, x, p, **keys):
        X, dgdx, (idgdp, jdgdp, dgdp) = self.g.derivatives(x, p, **keys)
        Y, dfdX, (idfdp, jdfdp, dfdp) = self.f.derivatives(X, p, **keys)
        M, N, L = dfdX.shape
        CR = np.matmul(dgdx.T, dfdX.T).T
        ii = np.hstack([idgdp, idfdp])
        jj = np.hstack([jdgdp, jdfdp])
        val = [
            np.hstack(
                [
                    np.sum([dgdp[i] * dfdX[j, i, idgdp] for i in range(N)], axis=0),
                    dfdp[j],
                ]
            )
            for j in range(M)
        ]
        return Y, CR, (ii, jj, val)
