import pytest
import numpy as np

import cristallina.analysis

__author__ = "Alexander Steppke"

	
def test_minimal_2d_gaussian():

    image = np.array([[0,0,0,0,0],
            [0,0,0,0,0],
            [0,0,1,0,0],
            [0,0,0,0,0],
            [0,0,0,0,0],
            ])
    
    center_x, center_y, result = cristallina.analysis.fit_2d_gaussian(image)
    assert np.allclose(center_x, 2.0, rtol=1e-04)
    assert np.allclose(center_y, 2.0, rtol=1e-04)



def test_2d_gaussian():

        # define normalized 2D gaussian
        def gauss2d(x=0, y=0, mx=0, my=0, sx=1, sy=1):
            return (1 / (2 * np.pi * sx * sy) * np.exp(-((x - mx) ** 2 / (2 * sx**2.0) + (y - my) ** 2 / (2 * sy**2))))

        x = np.arange(0, 150, 1)
        y = np.arange(0, 100, 1)
        x, y = np.meshgrid(x, y)  

        z = gauss2d(x, y, mx=40, my=50, sx=20, sy=40)
        
        center_x, center_y, result = cristallina.analysis.fit_2d_gaussian(z)
        assert np.allclose(center_x, 40, rtol=1e-04)
        assert np.allclose(center_y, 50, rtol=1e-04)


def test_2d_gaussian_rotated():

        # define normalized 2D gaussian
        def gauss2d_rotated(x=0, y=0, center_x=0, center_y=0, sx=1, sy=1, rotation=0):

            sr = np.sin(rotation)
            cr = np.cos(rotation)

            center_x_rot = center_x * cr - center_y * sr
            center_y_rot = center_x * sr + center_y * cr

            x_rot = x * cr - y * sr
            y_rot = x * sr + y * cr

            return (1 / (2 * np.pi * sx * sy) * np.exp(-((x_rot - center_x_rot) ** 2 / (2 * sx**2.0) + (y_rot - center_y_rot) ** 2 / (2 * sy**2))))

        x = np.arange(0, 150, 1)
        y = np.arange(0, 100, 1)
        x, y = np.meshgrid(x, y)  

        z = 100*gauss2d_rotated(x, y, center_x=40, center_y=50, sx=10, sy=20, rotation=0.5)
        
        center_x, center_y, result = cristallina.analysis.fit_2d_gaussian_rotated(z, vary_rotation=True, plot=False)
        assert np.allclose(center_x, 40, rtol=1e-04)
        assert np.allclose(center_y, 50, rtol=1e-04)
        assert np.allclose(result.params['rotation'].value, 0.5, rtol=1e-02)