# SPDX-License-Identifier: MPL-2.0
import pytest 
import numpy as np

from aare import _aare #import the C++ module
from aare import corner 
from conftest import test_data_path


def test_cluster_vector_can_be_converted_to_numpy():
    cv = _aare.ClusterVector_Cluster3x3i()
    arr = np.array(cv, copy=False)
    assert arr.shape == (0,)  # 4 for x, y, size, energy and 9 for the cluster data


def test_ClusterVector(): 
    """Test ClusterVector""" 

    clustervector = _aare.ClusterVector_Cluster3x3i()
    assert clustervector.cluster_size_x == 3
    assert clustervector.cluster_size_y == 3
    assert clustervector.item_size() == 4+9*4
    assert clustervector.frame_number == 0
    assert clustervector.size == 0

    cluster = _aare.Cluster3x3i(0,0,np.ones(9, dtype=np.int32))
    
    clustervector.push_back(cluster) 
    assert clustervector.size == 1

    with pytest.raises(TypeError):  # Or use the appropriate exception type
        clustervector.push_back(_aare.Cluster2x2i(0,0,np.ones(4, dtype=np.int32)))

    with pytest.raises(TypeError): 
        clustervector.push_back(_aare.Cluster3x3f(0,0,np.ones(9, dtype=np.float32))) 

def test_Interpolator(): 
    """Test Interpolator"""

    ebins = np.linspace(0,10, 20, dtype=np.float64)
    xbins = np.linspace(0, 5, 30, dtype=np.float64)
    ybins = np.linspace(0, 5, 30, dtype=np.float64)

    etacube = np.zeros(shape=[29, 29, 19], dtype=np.float64)
    interpolator = _aare.Interpolator(etacube, xbins, ybins, ebins)

    assert interpolator.get_ietax().shape == (29,29,19)
    assert interpolator.get_ietay().shape == (29,29,19)
    clustervector = _aare.ClusterVector_Cluster3x3i()
    
    cluster = _aare.Cluster3x3i(1,1, np.ones(9, dtype=np.int32))
    clustervector.push_back(cluster) 

    [u,v] = interpolator.transform_eta_values(_aare.Etai())

    assert u == 0
    assert v == 0

    interpolated_photons = interpolator.interpolate(clustervector)

    assert interpolated_photons.size == 1

    assert interpolated_photons[0]["x"] == 0.5 
    assert interpolated_photons[0]["y"] == 0.5
    assert interpolated_photons[0]["energy"] == 4 #eta_sum = 4, dx, dy = -1,-1 m_ietax = 0, m_ietay = 0

    clustervector = _aare.ClusterVector_Cluster2x2i()

    cluster = _aare.Cluster2x2i(1,1, np.ones(4, dtype=np.int32))
    clustervector.push_back(cluster) 

    interpolated_photons = interpolator.interpolate(clustervector)

    assert interpolated_photons.size == 1

    assert interpolated_photons[0]["x"] == 0.5
    assert interpolated_photons[0]["y"] == 0.5
    assert interpolated_photons[0]["energy"] == 4



def test_calculate_eta(): 
    """Calculate Eta""" 
    cluster = _aare.Cluster3x3i(0,0, np.ones(9, dtype=np.int32)) 

    eta2 = _aare.calculate_eta2(cluster)

    assert eta2.x == 0.5
    assert eta2.y == 0.5
    assert eta2.c == corner.cTopLeft
    assert eta2.sum == 4


def test_max_sum(): 
    """Max 2x2 Sum"""
    cluster = _aare.Cluster3x3i(5,5,np.array([1, 1, 1, 2, 3, 1, 2, 2, 1], dtype=np.int32))

    max_sum = cluster.max_sum_2x2()

    assert max_sum[0] == 9
    assert max_sum[1] == 2


def test_cluster_finder(): 
    """Test ClusterFinder""" 

    clusterfinder = _aare.ClusterFinder_Cluster3x3i([100,100])

    #frame = np.random.rand(100,100)
    frame = np.zeros(shape=[100,100])

    clusterfinder.find_clusters(frame)

    clusters = clusterfinder.steal_clusters(False) #conversion does not work

    assert clusters.size == 0


def test_2x2_reduction(): 
    """Test 2x2 Reduction"""
    cluster = _aare.Cluster3x3i(5,5,np.array([1, 1, 1, 2, 3, 1, 2, 2, 1], dtype=np.int32))

    reduced_cluster = _aare.reduce_to_2x2(cluster) 

    assert reduced_cluster.x == 5
    assert reduced_cluster.y == 5
    assert (reduced_cluster.data == np.array([[2, 3], [2, 2]], dtype=np.int32)).all()
    
    
def test_3x3_reduction(): 
    """Test 3x3 Reduction"""
    cluster = _aare.Cluster5x5d(5,5,np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 2.0, 2.0, 3.0,
                                   1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.double))

    reduced_cluster = _aare.reduce_to_3x3(cluster) 

    assert reduced_cluster.x == 5
    assert reduced_cluster.y == 5
    assert (reduced_cluster.data == np.array([[2.0, 1.0, 1.0], [2.0, 3.0, 1.0], [2.0, 1.0, 1.0]], dtype=np.double)).all()