from .lin_op import LinOp
import numpy as np
import cv2

from proximal.halide.halide import *
from proximal.utils.utils import *

class warp(LinOp):
    """Warp using a homography.
    """
    def __init__(self, arg, H, implem = None ):
        self.H = H.copy()

        #Compute inverse
        self.Hinv = np.zeros(H.shape)
        if len(H.shape) > 2 :
            for j in range(self.H.shape[2]):
                self.Hinv[:,:,j] = np.linalg.pinv(H[:,:,j])
        else :
            self.Hinv = np.linalg.pinv(H)

        #Check for the shape
        if len(H.shape) < 2 or len(H.shape) > 3:
            raise Exception('Error, warp supports only up to 4d inputs (expects first 3 to be image).')

        #Has to have third dimension
        if len(arg.shape) != 3:
            raise Exception('Images must have third dimension')

        shape = arg.shape
        if len(H.shape) == 3:
            shape += (H.shape[2],)

        #Temp array for halide
        self.tmpfwd = np.zeros( (shape[0], shape[1], shape[2] if (len(shape) > 2) else 1, H.shape[2] if (len(H.shape) > 2) else 1),
                                dtype=np.float32, order='FORTRAN');
        self.tmpadj = np.zeros( (shape[0], shape[1], shape[2] if (len(shape) > 2) else 1), dtype=np.float32, order='FORTRAN');

        #Halide homographies
        if len(H.shape) == 2:
            self.Hf = np.asfortranarray( H[...,np.newaxis].astype(np.float32) ) #Third axis for halide
            self.Hinvf = np.asfortranarray( self.Hinv[...,np.newaxis].astype(np.float32) ) #Third axis for halide
        else:
            self.Hf = np.asfortranarray( H.astype(np.float32) )
            self.Hinvf = np.asfortranarray( self.Hinv.astype(np.float32) )

        super(warp, self).__init__([arg], shape, implem)

    def forward(self, inputs, outputs):
        """The forward operator.

        Reads from inputs and writes to outputs.
        """

        if self.implementation == Impl['halide'] :

            #Halide implementation
            tmpin = np.asfortranarray( inputs[0].astype(np.float32) )
            Halide('A_warp.cpp').A_warp( tmpin, self.Hinvf, self.tmpfwd ) #Call
            np.copyto(outputs[0], np.reshape( self.tmpfwd, self.shape ) )

        else:

            #CV2 version
            inimg = inputs[0]
            if len(self.H.shape) == 2:

                warpedInput = cv2.warpPerspective(np.asfortranarray(inimg),  self.H.T, inimg.shape[1::-1], flags=cv2.INTER_LINEAR ,
                                    borderMode=cv2.BORDER_CONSTANT, borderValue=0.) #Necessary due to array layout in opencv
                np.copyto( outputs[0], warpedInput )

            else:
                for j in range(self.H.shape[2]):
                    warpedInput = cv2.warpPerspective(np.asfortranarray(inimg), self.H[:,:,j].T, inimg.shape[1::-1], flags=cv2.INTER_LINEAR ,
                                                    borderMode=cv2.BORDER_CONSTANT, borderValue=0.) #Necessary due to array layout in opencv

                    np.copyto( outputs[0][:,:,:,j], warpedInput )

    def adjoint(self, inputs, outputs):
        """The adjoint operator.

        Reads from inputs and writes to outputs.
        """


        if self.implementation == Impl['halide'] :

            #Halide implementation
            if len(self.H.shape) == 2:
                tmpin = np.asfortranarray( inputs[0][..., np.newaxis].astype(np.float32) )
            else:
                tmpin = np.asfortranarray( inputs[0].astype(np.float32) )

            Halide('At_warp.cpp').At_warp( tmpin, self.Hf, self.tmpadj ) #Call
            np.copyto(outputs[0], self.tmpadj )

        else:

            #CV2 version
            inimg = inputs[0]
            if len(self.H.shape) == 2:
                # + cv2.WARP_INVERSE_MAP
                warpedInput = cv2.warpPerspective(np.asfortranarray(inimg), self.Hinv.T, inimg.shape[1::-1], flags=cv2.INTER_LINEAR,
                                    borderMode=cv2.BORDER_CONSTANT, borderValue=0.)
                np.copyto( outputs[0], warpedInput )

            else:
                outputs[0][:] = 0.0
                for j in range(self.H.shape[2]):
                    warpedInput = cv2.warpPerspective(np.asfortranarray(inimg[:,:,:,j]), self.Hinv[:,:,j].T, inimg.shape[1::-1], flags=cv2.INTER_LINEAR,
                                                    borderMode=cv2.BORDER_CONSTANT, borderValue=0.) #Necessary due to array layout in opencv
                    outputs[0] += warpedInput

    # TODO what is the spectral norm of a warp?
