"""
This file provides proximity operators and projectors for the
Python Phase Retrieval (pyPR) toolbox.
"""

from pyprUtilities import *
import sys
import numpy as np
import pyShearLab2D as pySL



def approximateMagnitudeProjection(input_data, Psi):
    """
    Calculates an approximate projection onto a regularized set of measurements. 
    This method can easily be extend to use input_data.noise.type and then 
    calculate different error distances accordingly.    
    
    Input:
        
        input_data:     Input_data structure.
        Psi:            Current iterate / wave function.
    """
    if calculateKLDistance(input_data, Psi) >= input_data["noise"]["error"]:
        outputWave = magnitudeProjection(input_data, Psi)
    else:
        outputWave = Psi
    return outputWave


def constraintProximityOperator(input_data, Psi, gamma):
    """
    Applies Projection/Proximity Operator depending on the given constraint.
    
    Input:
        
        input_data  Input_data structure.
        Psi         Object to be projected.
        gamma       Parameter (only for proximity operators)
    """
    proximityOperatorName = retrieveProximityName(input_data)

    if (proximityOperatorName == "proxSupportConstraint") or (proximityOperatorName == "proxSupportRangeConstraint") or (proximityOperatorName == "proxRangeConstraint"):
        outputPsi = eval(proximityOperatorName+"(input_data, Psi)")
    elif proximityOperatorName == "proxShearlet":
        #outputPsi = eval(proximityOperatorName, input_data, Psi, gamma)
        outputPsi = eval(proximityOperatorName+"(input_data, Psi, gamma)")
    else:
        sys.exit("The provided proximityOperatorName " + proximityOperatorName + " is not supported." +
                 "Check constraintProximityOperator in pyprProxOperators.py for details.")
    return outputPsi


def magnitudeProjection(input_data, Psi):
    """
    Calculates the projection onto a magnitude constraint set.
    
    Input:
    
        input_data  Input_data structure.
        Psi         Object / wave function to be projected.
    """

    propagatedWave = wavePropagator(input_data, Psi, 1)
    projectedWave = MagProj(input_data["problem"]["data"], propagatedWave)
    outputWave = wavePropagator(input_data, projectedWave, -1)

    return outputWave



def proxRangeConstraint(input_data, Psi):
    """
    Implements range constraint for different types of objects, e.h. positivity,
    negative phase, and others.
    
    Input:
        
        input_data  Input_data structure.
        Psi         Object / wave function to be projected.
    """

    if input_data["problemType"] == "phase":

        PsiPhase = np.angle(Psi)

        if "range" not in input_data["problem"]["constraint"]:

            # no info, hence negativitiy assumed.
            PsiPhase [PsiPhase >= 0] = 0
        else:
            if input_data["problem"]["constraint"]["range"] == "positive":

                PsiPhase [PsiPhase <= 0] = 0

        return np.exp(1j * PsiPhase)

    elif input_data["problemType"] == "amplitude":

        PsiAmplitude = np.abs(Psi)
        PsiAmplitude[ PsiAmplitude >= 1 ] = 1

        return PsiAmplitude

    elif input_data["problemType"]["mixed"]:

        PsiPhase = np.angle(Psi)
        PsiAmplitude = np.abs(Psi)

        if "range" not in input_data["problem"]["constraint"]:

            PsiPhase[ PsiPhase <= 0 ] = 0
            PsiAmplitude[ PsiAmplitude >= 1 ] = 1
        else:
            if input_data["problem"]["constraint"]["range"] == "positive":
                PsiPhase[ PsiPhase >= 0 ] = 0
                PsiAmplitude[ PsiAmplitude >= 1] = 1

        return PsiAmplitude * np.exp(1j * PsiPhase)

    elif input_data["problemType"] == "real":

        if "range" not in input_data["problem"]["constraint"]:

            Psi[np.real(Psi) <= 0] = 0                  # positive
            Psi = np.real(Psi)                          # and real-valued

            return Psi

    else:
        sys.exit("No suitable problem type specified in input_data['problemType'].")


def proxShearlet(input_data, Psi, gamma):

    """
    Implements the thresholding of shearlet coefficients for different kinds
    of objects (phase, amplitude, mixed and real objects) with general
    thresholding functions specified in the retrieveProximityName function.
    
    If the object is, e.g., a phase object, the shearlet decomposition and the
    thresholding is applide to the phase, e.g. to np.angle(Psi). Similarly this
    is done for mixed or amplitude objects.
    
    This can be combined with a range constraint (such as negativity for the
    phase) which will be applied before the shearlet decomposition. This is
    mathematically not completely consistent but makes sense in practice.
    
    Input:
        
        input_data  Input_data structure.
        Psi         Object / wave function the prox operator is applied to.
        gamma       Prox operator parameter (step length), threshold parameter.
    """

    shrinkageOperator = retrieveShrinkageName(input_data)

    if input_data["problemType"] == "phase":

        PsiPhase = np.angle(Psi)

        if "range" in input_data["problem"]["constraint"]:

            if input_data["problem"]["constraint"]["range"] == "positive":
                PsiPhase [ PsiPhase <= 0 ] = 0
            else:
                PsiPhase [ PsiPhase >= 0 ] = 0
        else:
            PsiPhase[ PsiPhase >= 0 ] = 0

        coefficientsPhase = pySL.SLsheardec2D(PsiPhase, input_data["system"])
        coefficientsPhase = eval(shrinkageOperator+"(input_data,coefficientsPhase, gamma)")
        PsiPhase = pySL.SLshearrec2D(coefficientsPhase, input_data["system"])

        return np.exp(1j*PsiPhase)

    elif input_data["problemType"] == "amplitude":

        if input_data["problem"]["constraint"]["name"] == "shearlet thresholding and range":
            Psi [ Psi >= 1 ] = 1

        PsiAmplitude = np.log(np.real(Psi))

        coefficientsAmplitude = pySL.SLsheardec2D(PsiAmplitude, input_data["system"])
        coefficientsAmplitude = eval(shrinkageOperator+"(input_data, coefficientsAmplitude, gamma)")
        PsiAmplitude = pySL.SLshearrec2D(coefficientsAmplitude, input_data["system"])

        return np.exp(PsiAmplitude)

    elif input_data["problemType"] == "mixed":

        PsiAmplitude = np.abs(Psi)
        PsiPhase = np.angle(Psi)

        coefficientsAmplitude = pySL.SLsheardec2D(PsiAmplitude, input_data["system"])
        coefficientsPhase = pySL.SLsheardec2D(PsiPhase, input_data["system"])

        coefficientsAmplitude = eval(shrinkageOperator+"(input_data, coefficientsAmplitude, gamma)")
        coefficientsPhase = eval(shrinkageOperator+"(input_data, coefficientsPhase, gamma)")

        PsiAmplitude = pySL.SLshearrec2D(coefficientsAmplitude, input_data["system"])
        PsiPhase = pySL.SLshearrec2D(coefficientsPhase, input_data["system"])

        if input_data["constraint"]["name"] == "shearlet thresholding and range":
            PsiAmplitude[ PsiAmplitude >= 1] = 1
            PsiPhase[ PsiPhase >= 0 ] = 0       # extend this for specific range information

        return PsiAmplitude * np.exp(1j * PsiPhase)

    elif input_data["problemType"] == "real":

        PsiReal = np.real(Psi)
        coefficientsReal = pySL.SLsheardec2D(PsiReal, input_data["system"])
        coefficientsReal = eval(shrinkageOperator+"(input_data, coefficientsReal, gamma)")

        PsiReal = pySL.SLshearrec2D(coefficientsReal, input_data["system"])

        if input_data["problem"]["constraint"]["name"] == "shearlet thresholding and range":

            PsiReal = np.real(PsiReal)
            PsiReal[ PsiReal <= 0 ] = 0

        return PsiReal

    else:
        sys.exit("No suitable problem type specified in input_data.problem.type. See proxShearlet in pyprProxOperators.py for implementation details.")

# Implements support constraints for different objects.
# Input:    input_data  - General input_data structure.
#           Psi         - Object to be projected.
#
# Output:   outputPsi   - Object after projection onto support constraint
#
# TODO: Figure out if we want to use np.abs or np.log for the amplitude part.
# TODO: It porbably only differs in the scaling which should not matter for
# TODO: support constraint but will for something like thresholding.
#
# TODO: Would it make sense to have different supports in mixed objects for phase/ampl.?
#

def proxSupportConstraint(input_data, Psi):
    """
    Implements support constraints for different objects.
    
    Input:
        
        input_data   Input_data structure.
        Psi          Object / wave function to be projected.    
    """

    if input_data["problemType"] == "phase":

        PsiPhase = np.angle(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsiPhase = np.zeros(PsiPhase.shape)

        outPsiPhase = PsiPhase[suppPsi]

        return np.exp(1j * outPsiPhase)

    elif input_data["problemType"] == "amplitude":

        PsiAmpl = np.log(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsi = np.zeros(Psi.shape)

        outPsi[suppPsi] = PsiAmpl[suppPsi]

        return np.exp(outPsi)


    elif input_data["problemType"] == "mixed":

        PsiPhase = np.angle(Psi)
        PsiAmpl = np.abs(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsiPhase = np.zeros(Psi.shape)
        outPsiAmpl = np.zeros(Psi.shape)

        outPsiPhase[suppPsi] = PsiPhase[suppPsi]
        outPsiAmpl[suppPsi] = PsiAmpl[suppPsi]

        return outPsiAmpl * np.exp(1j * outPsiPhase)

    elif input_data["problemType"] == "real":

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsi = np.zeros(Psi.shape)

        outPsi[suppPsi] = Psi[suppPsi]

        return outPsi

    else:
        sys.exit("Error: No suitable problem type specified in input_data['problemType'']. See proxSupportConstraint in pyprProxOperators.py for details.")


# Implements support + range constraints for different objects.
# Input:    input_data  - General input_data structure.
#           Psi         - Object to be projected.
#
# Output:   outputPsi   - Object after projection onto support + range constraint
#
# TODO: Figure out if we want to use np.abs or np.log for the amplitude part.
# TODO: It porbably only differs in the scaling which should not matter for
# TODO: support constraint but will for something like thresholding or range constraint.
#
# TODO: Would it make sense to have different supports in mixed objects for phase/ampl.?
#
def proxSupportRangeConstraint(input_data, Psi):

    if input_data["problemType"] == "phase":

        PsiPhase = np.angle(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsiPhase = np.zeros(PsiPhase.shape)

        # If no range type is specified we go with negative phase.

        if "range" in input_data["problem"]["constraint"]:

            if input_data["problem"]["constraint"]["range"] == "positive":
                PsiPhase[PsiPhase <= 0] = 0
            else:
                PsiPhase[PsiPhase >= 0] = 0

        else:
            PsiPhase[PsiPhase >= 0] = 0

        outPsiPhase[suppPsi] = PsiPhase[suppPsi]

        return np.exp(1j * outPsiPhase)

    elif input_data["problemType"] == "amplitude":

        PsiAmpl = np.log(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsi = np.zeros(Psi.shape)

        # TODO: Extend to different amplitude constraints? Does that make sense?
        PsiAmpl [ PsiAmpl >= 1 ] = 1

        outPsi[suppPsi] = PsiAmpl[suppPsi]

        return np.exp(outPsi)


    elif input_data["problemType"] == "mixed":

        PsiPhase = np.angle(Psi)
        PsiAmpl = np.abs(Psi)

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsiPhase = np.zeros(Psi.shape)
        outPsiAmpl = np.zeros(Psi.shape)

        # if no range type is specified we go with negative phase
        if "range" in input_data["problem"]["constraint"]:

            if input_data["problem"]["constraint"]["range"] == "positive":
                PsiPhase[PsiPhase <= 0] = 0
            else:
                PsiPhase[PsiPhase >= 0] = 0
        else:
            PsiPhase[PsiPhase >= 0] = 0

        # TODO: Extend to different amplitude constraints? Does that make sense?
        PsiAmpl[PsiAmpl >= 1] = 1

        outPsiPhase[suppPsi] = PsiPhase[suppPsi]
        outPsiAmpl[suppPsi] = PsiAmpl[suppPsi]

        return outPsiAmpl * np.exp(1j * outPsiPhase)

    elif input_data["problemType"] == "real":

        suppPsi = input_data["problem"]["constraint"]["support"]

        outPsi = np.zeros(Psi.shape)

        Psi = np.real(Psi)

        # TODO: This is the non-negativitiy. Could this be extended in a meaningful way?
        Psi[Psi <= 0] = 0

        outPsi[suppPsi] = Psi[suppPsi]

        return outPsi

    else:
        sys.exit("Error: No suitable problem type specified in input_data['problemType'']. See proxSupportRangeConstraint in pyprProxOperators.py for details.")