import numpy as np
from clawpack import pyclaw
from clawpack import riemann
from unflatten import unflatten

from pySDC.core.Problem import ptype
from pySDC.implementations.datatype_classes import mesh, rhs_imex_mesh


class advection_2d_explicit(ptype):
    """
    Example implementing the forced 1D heat equation with Dirichlet-0 BC in [0,1]

    Attributes:
      solver: Sharpclaw solver
      state:  Sharclaw state
      domain: Sharpclaw domain
    """

    def __init__(self, cparams, dtype_u, dtype_f):
        """
        Initialization routine

        Args:
            cparams: custom parameters for the example
            dtype_u: particle data type (will be passed parent class)
            dtype_f: acceleration data type (will be passed parent class)
        """

        # these parameters will be used later, so assert their existence
        assert 'nvars' in cparams

        # add parameters as attributes for further reference
        for k, v in cparams.items():
            setattr(self, k, v)

        # invoke super init, passing number of dofs, dtype_u and dtype_f
        super(advection_2d_explicit, self).__init__(self.nvars, dtype_u, dtype_f)

        riemann_solver = riemann.advection_2D  # NOTE: This uses the FORTRAN kernels of clawpack
        self.solver = pyclaw.SharpClawSolver2D(riemann.advection_2D)
        self.solver.weno_order = 5
        self.solver.time_integrator = 'Euler'  # Remove later
        self.solver.kernel_language = 'Fortran'
        self.solver.bc_lower[0] = pyclaw.BC.periodic
        self.solver.bc_upper[0] = pyclaw.BC.periodic
        self.solver.bc_lower[1] = pyclaw.BC.periodic
        self.solver.bc_upper[1] = pyclaw.BC.periodic
        self.solver.cfl_max = 1.0
        assert self.solver.is_valid()

        x = pyclaw.Dimension(-1.0, 1.0, self.nvars[1], name='x')
        y = pyclaw.Dimension(0.0, 1.0, self.nvars[2], name='y')
        self.domain = pyclaw.Domain([x, y])

        self.state = pyclaw.State(self.domain, self.solver.num_eqn)
        # self.dx = self.state.grid.x.centers[1] - self.state.grid.x.centers[0]

        self.state.problem_data['u'] = 1.0
        self.state.problem_data['v'] = 0.0

        solution = pyclaw.Solution(self.state, self.domain)
        self.solver.setup(solution)

        self.xc, self.yc = self.state.grid.p_centers

    def solve_system(self, rhs, factor, u0, t):
        """
        Simple linear solver for (I-dtA)u = rhs

        Args:
            rhs: right-hand side for the nonlinear system
            factor: abbrev. for the node-to-node stepsize (or any other factor required)
            u0: initial guess for the iterative solver (not used here so far)
            t: current time (e.g. for time-dependent BCs)

        Returns:
            solution as mesh
        """

        # me = mesh(self.nvars)
        # me.values = LA.spsolve(sp.eye(self.nvars)-factor*self.A,rhs.values)

        return rhs

    def __eval_fexpl(self, u, t):
        """
        Helper routine to evaluate the explicit part of the RHS

        Args:
            u: current values (not used here)
            t: current time

        Returns:
            explicit part of RHS
        """

        fexpl = mesh(self.nvars)

        # Copy values of u into pyClaw state object
        self.state.q[0, :, :] = u.values[0, :, :]

        # Evaluate right hand side
        self.solver.before_step(self.solver, self.state)
        tmp = self.solver.dqdt(self.state)

        fexpl.values[0, :, :] = unflatten(tmp, 1, self.nvars[1], self.nvars[2])

        # Copy values of u into pyClaw state object
        # self.state.q[0,:,:] = u.values[1,:,:]

        # Evaluate right hand side
        # tmp = self.solver.dqdt(self.state)
        # fexpl.values[1,:,:] = tmp.reshape(self.nvars[1:])

        return fexpl

    def __eval_fimpl(self, u, t):
        """
        Helper routine to evaluate the implicit part of the RHS

        Args:
            u: current values
            t: current time (not used here)

        Returns:
            implicit part of RHS
        """

        fimpl = mesh(self.nvars, val=0.0)
        # fimpl.values = self.A.dot(u.values)

        return fimpl

    def eval_f(self, u, t):
        """
        Routine to evaluate both parts of the RHS

        Args:
            u: current values
            t: current time

        Returns:
            the RHS divided into two parts
        """

        f = rhs_imex_mesh(self.nvars)
        f.impl = self.__eval_fimpl(u, t)
        f.expl = self.__eval_fexpl(u, t)
        return f

    def u_exact(self, t):
        """
        Routine to compute the exact solution at time t

        Args:
            t: current time

        Returns:
            exact solution
        """

        me = mesh(self.nvars)
        me.values[0, :, :] = np.sin(2 * np.pi * self.xc) * np.sin(2 * np.pi * self.yc)
        # me.values[1,:,:] = np.sin(2*np.pi*self.xc)#*np.sin(2*np.pi*self.yc)

        return me
