#!/usr/bin/env python
# coding=utf8

import numpy as np
import sympy as sp


class ModelRoutine:
    def __init__(self, matrix, args, pars, ufunc,
                 reduced=False):
        self.pars = list(pars) + ['periodic']
        self.matrix = matrix
        self.args = args
        self._ufunc = ufunc

    def __repr__(self):
        return sp.Matrix(self.matrix.tolist()).__repr__()


class F_Routine(ModelRoutine):
    """Compute the right hand side of the dynamical system
      :math:`\\frac{\\partial U}{\\partial t} = F(U)`

      Parameters
      ----------
      fields : triflow.Fields
          triflow fields container generated by a triflow.Model containing the actual state of the dependent variables and helper functions.
      pars : dict
          dictionnary with the different physical parameters of the model and the 'periodic' key.

      Returns
      -------
      numpy.ndarray
          flat array containing the right hand side of the dynamical system.
      """  # noqa

    def __call__(self, fields, pars):
        uargs = [fields['x'], *[fields[key] for key in self.args]]
        pargs = [pars[key] + fields.x * 0 if key != 'periodic' else pars[key]
                 for key
                 in self.pars]
        F = self._ufunc(*uargs, *pargs)
        return F

    def diff_approx(self, fields, pars, eps=1E-8):
        nvar, N = len(fields.dependent_variables), fields.size
        fpars = {key: pars[key] for key in self.pars}
        fpars['dx'] = (fields['x'][-1] - fields['x'][0]) / fields['x'].size
        J = np.zeros((N * nvar, N * nvar))
        indices = np.indices(fields.uarray.shape)
        for i, (var_index, node_index) in enumerate(zip(*map(np.ravel,
                                                             indices))):
            fields_plus = fields.copy()
            fields_plus.uarray[var_index, node_index] += eps
            fields_moins = fields.copy()
            fields_moins.uarray[var_index, node_index] -= eps
            Fplus = self(fields_plus, pars)
            Fmoins = self(fields_moins, pars)
            J[i] = (Fplus - Fmoins) / (2 * eps)

        return J.T


class J_Routine(ModelRoutine):
    """Compute the right hand side of the dynamical system
      :math:`\\frac{\\partial U}{\\partial t} = F(U)`

      Parameters
      ----------
      fields : triflow.Fields
          triflow fields container generated by a triflow.Model containing the actual state of the dependent variables and helper functions.
      pars : dict
          dictionnary with the different physical parameters of the model and the 'periodic' key.
      sparse : bool, optional, default True
          whether should the matrix returned as dense or sparse form.

      Returns
      -------
      scipy.sparse.CSC or numpy.ndarray: sparse or dense form (depending of the `sparse` argument) of the Jacobian approximation of the dynamical system right hand side.
      """  # noqa

    def __call__(self, fields, pars, sparse=True):
        uargs = [fields['x'], *[fields[key] for key in self.args]]
        pargs = [pars[key] + fields['x'] * 0
                 if key != 'periodic' else pars[key]
                 for key
                 in self.pars]
        J = self._ufunc(*uargs, *pargs)

        return J if sparse else J.todense()
