# -*- coding: utf-8 -*-
#   This work is part of the Core Imaging Library (CIL) developed by CCPi 
#   (Collaborative Computational Project in Tomographic Imaging), with 
#   substantial contributions by UKRI-STFC and University of Manchester.

#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import numpy as np
from cil.framework import ImageData
from cil.optimisation.operators import LinearOperator

class DiagonalOperator(LinearOperator):

    r"""DiagonalOperator 

    Performs an element-wise multiplication, i.e., `Hadamard Product <https://en.wikipedia.org/wiki/Hadamard_product_(matrices)#:~:text=In%20mathematics%2C%20the%20Hadamard%20product,elements%20i%2C%20j%20of%20the>`_ 
    of a :class:`DataContainer` `x` and :class:`DataContainer` `diagonal`, `d` .

    .. math:: (D\circ x) = \sum_{i,j}^{M,N} D_{i,j} x_{i, j}

    In matrix-vector interpretation, if `D` is a :math:`M\times N` dense matrix and is flattened, we have a :math:`M*N \times M*N` vector.
    A sparse diagonal matrix, i.e., :class:`DigaonalOperator` can be created if we add the vector above to the main diagonal.
    If the :class:`DataContainer` `x` is also flattened, we have a :math:`M*N` vector.
    Now, matrix-vector multiplcation is allowed and results to a :math:`(M*N,1)` vector. After reshaping we recover a :math:`M\times N` :class:`DataContainer`. 

    Parameters
    ----------
    diagonal : DataContainer
        DataContainer with the same dimensions as the data to be operated on.
    domain_geometry : ImageGeometry
        Specifies the geometry of the operator domain. If 'None' will use the diagonal geometry directly. default=None .

    """

    
    def __init__(self, diagonal, domain_geometry=None):

        if domain_geometry is None:
            domain_geometry = diagonal.geometry.copy()

        super(DiagonalOperator, self).__init__(domain_geometry=domain_geometry, 
                                    range_geometry=domain_geometry)
        self.diagonal = diagonal

        
    def direct(self,x,out=None):
        
        "Returns :math:`D\circ x` "
        
        if out is None:
            return self.diagonal * x
        else:
            self.diagonal.multiply(x,out=out)
    

    def adjoint(self,x, out=None):
        
        "Returns :math:`D\circ x` "
        
        return self.direct(x, out=out)

  
    def calculate_norm(self, **kwargs):
        
        " Returns the operator norm of DiagonalOperator which is the maximum element in the `diagonal`."
        
        return self.diagonal.max()
