import pytest
import numpy as np
import unittest.mock

import cristallina.analysis
import cristallina.utils

__author__ = "Alexander Steppke"


def test_joblib_memory():
    """We need joblib for fast caching of intermediate results in all cases. So we check
    if the basic function caching to disk works.
    """

    def calc_example(x):
        return x**2

    calc_cached = cristallina.analysis.memory.cache(calc_example)

    assert calc_cached(8) == 64
    assert calc_cached.check_call_in_cache(8) == True


@unittest.mock.patch(
    "jungfrau_utils.file_adapter.locate_gain_file",
    lambda path, **kwargs: "tests/data/gains.h5",
)
@unittest.mock.patch(
    "jungfrau_utils.file_adapter.locate_pedestal_file",
    lambda path, **kwargs: "tests/data/JF16T03V01.res.h5",
)
def test_image_calculations():
    res = cristallina.analysis.perform_image_calculations(["tests/data/p20841/raw/run0185/data/acq0001.*.h5"])
    # these values are only correct when using the specific gain and pedestal files included in the test data
    # they do not correspond to the gain and pedestal files used in the actual analysis
    intensity = [
        1712858.6,
        693994.06,
        1766390.0,
        1055504.9,
        1516520.9,
        461969.06,
        3148285.5,
        934917.5,
        1866691.6,
        798191.2,
        2250207.0,
        453842.6,
    ]
    assert np.allclose(res["JF16T03V01_intensity"], intensity)


@unittest.mock.patch(
    "jungfrau_utils.file_adapter.locate_gain_file",
    lambda path, **kwargs: "tests/data/gains.h5",
)
@unittest.mock.patch(
    "jungfrau_utils.file_adapter.locate_pedestal_file",
    lambda path, **kwargs: "tests/data/JF16T03V01.res.h5",
)
def test_roi_calculation():
    roi = cristallina.utils.ROI(left=575, right=750, top=750, bottom=600)

    cutouts = cristallina.analysis.perform_image_roi_crop(["tests/data/p20841/raw/run0205/data/acq0001.*.h5"], roi=roi)

    sum_roi, max_roi, min_roi = (
        np.sum(cutouts[11]),
        np.max(cutouts[11]),
        np.min(cutouts[11]),
    )

    # these values are only correct when using the specific gain and pedestal files included in the test data
    # they do not correspond to the gain and pedestal files used in the actual analysis
    assert np.allclose([3119071.8, 22381.547, -0.9425874], [sum_roi, max_roi, min_roi])


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)


def test_1d_gaussian():
    def gauss(x, H, A, x0, sigma):
        return H + A * np.exp(-((x - x0) ** 2) / (2 * sigma**2))

    x = np.linspace(0, 20)
    y = gauss(x, 0.5, 2, 10, 4)

    res = cristallina.analysis.fit_1d_gaussian(x, y)

    assert np.allclose(res.values["center"], 10)
    assert np.allclose(res.values["sigma"], 4)
    assert np.allclose(res.values["height"], 2)