import numpy as np
from .proxoperators import ProxOperator


class P_S_real(ProxOperator):
    """
    Projection subroutine for projecting onto support constraints
    """

    def __init__(self, config):
        """
        Initialization
        
        Parameters
        ----------
        config : dict - Dictionary containing the problem configuration. It must contain the following mapping:
        'support_idx' : array_like - vector of indeces of the nonzero elements of the array
        """
        self.support_idx = config['support_idx']

    def work(self, u):
        """
        Parameters
        ----------
        u : array_like - array to be projected
        
        Returns
        -------
        p_S : array_like - the projection 
        """
        support_idx = self.support_idx - 1
        p_S = np.zeros(u.shape, dtype=u.dtype)
        p_S[support_idx] = np.real(u[support_idx])
        return p_S