import getdp as GetDP
import numpy as np

class GetDPWave:
    """A GetDP solver using complex arithmetic for time-harmonic wave problems

    Only one instance of this class makes sens

    The .pro file should use the following variables:
    angularFreqRe -- the real part of the angular frequency to use
    angularFreqIm -- the imaginary part of the angular frequency to use
    b -- the right hand side (RHS) to use
    x -- the solution vector computed by GetDP
    imposeRHS -- should the imposed RHS be taken into account?
    viewOnly -- should x be used for post-processing only
    fileName -- post-processing file name
    """

    def __init__(self, pro, mesh, resolution):
        """Instanciates a new SolverGetDP with a full '-solve'

        Keyword arguments:
        pro -- the .pro file to use
        mesh -- the .msh file to use
        resolution -- the resolution (from the .pro file) to use
        """
        # Save
        self.__pro        = pro
        self.__mesh       = mesh
        self.__resolution = resolution

        # Generate DoFs and initialise RHS and solution vectors in GetDP
        GetDP.GetDP(["getdp",  self.__pro,
                     "-msh",   self.__mesh,
                     "-solve", self.__resolution,
                     "-v", "2"])

    def solution(self):
        """Returns the solution"""
        return self.__toNumpy(GetDP.GetDPGetNumber("x"))

    def size(self):
        """Returns the number of degrees of freedom"""
        return self.solution().shape[0]

    def solve(self, b, w):
        """Solves with b as RHS and w as complex angular frequency"""
        GetDP.GetDPSetNumber("angularFreqRe", np.real(w).tolist())
        GetDP.GetDPSetNumber("angularFreqIm", np.imag(w).tolist())
        GetDP.GetDPSetNumber("b", self.__toGetDP(b))
        GetDP.GetDPSetNumber("imposeRHS", 1)
        GetDP.GetDP(["getdp", self.__pro,
                     "-msh",  self.__mesh,
                     "-cal"])
        GetDP.GetDPSetNumber("imposeRHS", 0)

    def view(self, x, posName, fileName):
        """Generates a post-processing view

        Keyword arguments:
        x -- the solution vector to use
        posName -- the post-operation to use (from the .pro file)
        fileName -- the post-precessing file name

        This method generates a linear system
        """
        GetDP.GetDPSetNumber("x", self.__toGetDP(x))
        GetDP.GetDPSetNumber("viewOnly", 1)
        GetDP.GetDPSetString("fileName", fileName)
        GetDP.GetDP(["getdp", self.__pro,
                     "-msh",  self.__mesh,
                     "-cal",
                     "-pos",  posName])
        GetDP.GetDPSetNumber("viewOnly", 0)

    @staticmethod
    def __toNumpy(vGetDP):
        """Takes a GetDP list and returns a numpy array"""
        size   = vGetDP.size() / 2
        vNumpy = np.empty((size, 1), dtype=complex)

        for i in range(size):
            vNumpy[i] = complex(vGetDP[i*2], vGetDP[i*2 + 1])

        return vNumpy

    @staticmethod
    def __toGetDP(vNumpy):
        """Takes a numpy array and returns a GetDP list"""
        size   = vNumpy.shape[0]
        vGetDP = list()

        for i in range(size):
            vGetDP.append(float(np.real(vNumpy[i])))
            vGetDP.append(float(np.imag(vNumpy[i])))

        return vGetDP