"""
This package contains algorithms that use proximity or projection
operators to obtain a solution to a feasibility or general minimization
problem based on two proximity/projection operators. The neccessary
input variables are described exemplary for the MAP algorithm but are
the same for all algorithms. Some algorithms may need additional parameters
such as the RAAR algorithm uses beta_k for a relaxation.

The output is also described for the MAP algorithm and should be the
same. In the following, P_A and P_B denote the projections onto A or B
or the prox operators with respect to a function A or B, respectively.
Hence, if we have prox operators, MAP is: x^(k+1) = prox_B prox_A x^k.

For more details please refer to the documentation.
"""

import numpy as np
from pyprUtilities import *
from pyprProxOperators import *


def MAP(input_data):

    """
    MAP: Method of Alternating Projections.
    Algorithm:

        x^(k+1) = P_B P_A x^k

    Input:

        input_data: A dictionary containing all relevant information
                    to run the algorithm. The following fields are
                    neccessary while there are optional parameters
                    you can but do not have to define (uses default
                    values if not supplied). The the documentation
                    for more details.
        ["data"]:   Only for simulated data. This is something
                    like the phase function from which we
                    generate and optical transfer function (OTF).
        ["problemType"]:    Either "phase", "amplitude", "mixed" or
                            "real".
                            A phase object is of the form
                            Psi = exp(i * g)
                            for some real-valued g.
        ["algorithm"]: Dictionary containg parameters for the
                        algorithm, the following are neccessary:
        [..]["name"]:   Name of the algorithm, currently "MAP" or
                        "RAAR".
        [..]["error"]:  Error to be achieved, used as stopping
                        criterion.
        [..]["maxit"]:  Maximum number of iterations.

        ["Psi"]:    Wave function or object from which measurements
                    are created. Only neccessary for simulated data.
        ["problem"]: Dictionary describing problem parameters,
                     such as the underlying transform. neccessary
                     parameters:
        [..]["propagator"]: "Fourier" for Fourier transform,
                        "Fresnel" for Fresnel transform.
        [..]["Fx"]:     Fresnel number in x direction. Only needed
                        if propagator is "Fresnel".
        [..]["Fy"]:     Fresnel number in y direction. Only needed
                        if propagator is "Fresnel".
        [..]["data"]:  Contains the measured data, e.g. |F(Psi)|
                        where F denotes the propagator and Psi
                        the OTF.
        [..]["constraint"]: Dictionary containing information on the
                        constraint (such as support, range,
                        sparsity).
                        The following are mandatory fields:
        [..][..]["name"]:  "support", "support and range", "shearlet
                        thresholding", "shearlet thresholding and
                        range". See documentation for their
                        meaning.
        [..][..]["info"]:   Additional parameter, e.g. type of range
                        or type of thresholding.
        [..][..]["parameter"]: "constant" or "dynamic", see
                            documentation, needed only for
                            "shearlet thresholding" and
                            "shearlet thresholding and range".
        [..][..]["gamma0"]: (Initial) Parameter for thresholding,
                        only needed for "shearlet thresholding"
                        and "shearlet thresholding and range".
        [..][..]["support"]: Indicies for the support set, only
                        needed for constraint "support" and
                        "support and range".
        ["initialGuess"]: The start-value for the algorithm.
                        Usually one uses
                        F^-1(input_data["problem"]["data"]).
                        This can be computed using the
                        wavePropagator function.
        ["system"]:     Shearlet system. Only needed when using
                        "shearlet thresholding" or "shearlet
                        thresholding and range". Can be computed
                        using the pyShearLab toolbox.
        If you want to simulate noisy data, the following is
        also neccessary:
        ["noise"]:  Dictionary containt parameters for noise
                    generation.
        [..]["type"]: "poisson" (currently the only option)
        [..]["level"]: The level or intensity of the noise.
                       For Poisson, this is proportional to
                       the number of photons / exposure time.
        [..]["error"]: Bound on the error of the data used for
                        the approximate projection.

    Output:

        output_data: A dictionary containg the following:
        ["problem"]: Copied from input_data, see above.
        ["algorithm"]: Copied from input_data, see above.
        ["reconstruction"]: The reconstructed OTF / wave function.
        ["error"]: An array that contains the computed error
                    for each iteration step.
    """

    x = input_data["initialGuess"]
    k = 1

    output_data = {}
    output_data["error"] = np.zeros(input_data["algorithm"]["maxit"])

    error = np.inf

    while (k < input_data["algorithm"]["maxit"]) and (error > input_data["algorithm"]["error"]):

        gamma = computeShrinkageParameter(input_data, k)

        print("Constraint: " + input_data["problem"]["constraint"]["name"] + " Iteration: " + str(k) + " Gamma: "
              + str(gamma) + " Error: " + str(error))

        tmp1 = magnitudeProjection(input_data, x)
        x_old = x.copy()
        x = constraintProximityOperator(input_data, tmp1, gamma)

        error = computeError(input_data, x, x_old, tmp1)

        output_data["error"][k] = error
        k += 1

    output_data["error"] = output_data["error"][:k-1]

    output_data["reconstruction"] = x
    # "copy" information from the input data for post-processing
    output_data["problem"] = input_data["problem"]
    output_data["algorithm"] = input_data["algorithm"]

    if 'physicalParameters' in input_data:
        output_data["physicalParameters"] = input_data["physicalParameters"]

    return output_data


def RAAR(input_data):

    """
    RAAR: Relaxed Averaged Alternating Reflections.
    Algorithm:

            x^(k+1) = beta_k/2 (R_B R_A+I)x^(k) + (1-beta_k)P_A x^(k)
            where R_A = 2*P_A - I, I = Identity

    Input: See MAP.

    Output: See MAP.
    """


    x = input_data["initialGuess"]
    tmp1 = 2*magnitudeProjection(input_data, x) - x
    k = 1
    output_data = {}
    output_data["error"] = np.zeros(input_data["algorithm"]["maxit"])
    error = np.inf

    while (k < input_data["algorithm"]["maxit"]) and (error > input_data["algorithm"]["error"]):

        beta = computeRelaxationParameter(input_data, k)
        gamma = computeShrinkageParameter(input_data, k)

        print("Constraint: " + input_data["problem"]["constraint"]["name"] + " Iteration: " + str(k) + " Gamma: " + str(
            gamma) + " Beta: " + str(beta) + " Error: " + str(error))

        tmp3 = constraintProximityOperator(input_data, tmp1, gamma)

        x_old = x.copy()
        tmp_u = (beta*(2*tmp3-tmp1) + (1-beta)*tmp1)/2  # raar update rule
        x = tmp_u.copy()
        tmp1 = 2*magnitudeProjection(input_data, x) - x # reflector next iteration

        error = computeError(input_data, x, x_old, tmp3)
        output_data["error"][k-1] = error

        k += 1

    output_data["error"] = output_data["error"][:k-1]
    output_data["reconstruction"] = x
    # "copy" information from the input data for post-processing
    output_data["problem"] = input_data["problem"]
    output_data["algorithm"] = input_data["algorithm"]
    if 'physicalParameters' in input_data:
        output_data["physicalParameters"] = input_data["physicalParameters"]

    return output_data