# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2021 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.

import numpy as np
import pytest

import nifty8 as ift

from ..common import list2fixture

dtuple = ift.DomainTuple.make((ift.RGSpace(2, 0.2), ift.RGSpace(3, 0.3), 
                                ift.RGSpace(4, 0.4), ift.RGSpace(5, 0.5)))

pmp = pytest.mark.parametrize

domain = list2fixture([dtuple])
spaces = list2fixture((None, (2,), (1, 3), (1, 2, 3), (0, 1, 2, 3)))

def test_matrix_product_endomorphic(domain, spaces, n_tests=4):
    mat_shape = ()
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
        mat_shape += domain.shape
    mat_shape = mat_shape*2
    
    for i in range(n_tests):
        mat = ift.random.current_rng().standard_normal(mat_shape)
        mat = mat + 1j*ift.random.current_rng().standard_normal(mat_shape)
        op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
        ift.extra.check_linear_operator(op)
    print(f'Domain shape={domain.shape}, spaces={spaces}, '+
    f'matrix shape={mat_shape}, target=domain (endomorphic)')

def test_matrix_product_spaces(domain, spaces, n_tests=4):
    mat_shape = (7, 8)
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
            mat_shape += domain.shape
        
    for i in range(n_tests):
        mat = ift.random.current_rng().standard_normal(mat_shape)
        mat = mat + 1j*ift.random.current_rng().standard_normal(mat_shape)
        op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
        ift.extra.check_linear_operator(op)
    print(f'Domain shape={domain.shape}, spaces={spaces}, '+
    f'matrix shape={mat_shape}, target shape={op.target.shape}')

def test_matrix_product_flatten(domain, n_tests=4):
    appl_shape = (ift.utilities.my_product(domain.shape),)
    mat_shape = appl_shape * 2
    for i in range(n_tests):
        mat = ift.random.current_rng().standard_normal(mat_shape)
        mat = mat + 1j*ift.random.current_rng().standard_normal(mat_shape)
        op = ift.MatrixProductOperator(domain, mat, spaces=None, flatten=True)
        ift.extra.check_linear_operator(op)
    print(f'flatten=True. Domain shape={domain.shape}, matrix shape={mat_shape}')

# the below function demonstrates the only error that cannot be caught
# when the operator is initialized. It is caused due to the matrix having
# too few dimensions to stand in the places of summed over axes of the domain
# as explained in the operator's documentation.
def test_matrix_product_invalid_shapes(domain):
    mat_shape = ()
    spaces = (2,)
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
            mat_shape += domain.shape
    with pytest.raises(ValueError):
        mat = ift.random.current_rng().standard_normal(mat_shape)
        op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
        ift.extra.check_linear_operator(op)
    print('ValueError raised because positions of unused subspaces of '+
          'domain are changed.\n'+
          f'Domain shape={domain.shape}, spaces={spaces}, matrix shape={mat_shape}')
    mat_shape = ()
    spaces = (3,)
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
            mat_shape += domain.shape
    mat = ift.random.current_rng().standard_normal(mat_shape)
    op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
    ift.extra.check_linear_operator(op)
    print('No errors raised because positions of unused subspaces of '+
          'domain are not changed.\n'+
          f'Domain shape={domain.shape}, spaces={spaces}, '+
          f'matrix shape={mat_shape}, target shape= {op.target.shape}')
    mat_shape = (7,)
    spaces = (1, 2)
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
            mat_shape += domain.shape
    with pytest.raises(ValueError):
        mat = ift.random.current_rng().standard_normal(mat_shape)
        op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
        ift.extra.check_linear_operator(op)
    print('ValueError raised because positions of unused subspaces of '+
          'domain are changed.\n'+
          f'Domain shape={domain.shape}, spaces={spaces}, matrix shape={mat_shape}')
    mat_shape = (7,)
    spaces = (1, 3)
    if spaces != None:
        for i in spaces:
            mat_shape += domain[i].shape
    else:
            mat_shape += domain.shape
    mat = ift.random.current_rng().standard_normal(mat_shape)
    op = ift.MatrixProductOperator(domain, mat, spaces=spaces)
    ift.extra.check_linear_operator(op)
    print('No errors raised because positions of unused subspaces of '+
          'domain are not changed.\n'+
          f'Domain shape={domain.shape}, spaces={spaces}, '+
          f'matrix shape={mat_shape}, target shape={op.target.shape}')