import pytest
import numpy as np
from nabu.utils import median2 as nabu_median_filter
from nabu.testutils import get_data
from nabu.cuda.utils import get_cuda_context, __has_pycuda__
from nabu.preproc.ccd import CCDCorrection
if __has_pycuda__:
    import pycuda.gpuarray as garray
    from nabu.preproc.ccd_cuda import CudaCCDCorrection


@pytest.fixture(scope='class')
def bootstrap(request):
    cls = request.cls

    cls.data = get_data("mri_proj_astra.npz")["data"]
    cls.data /= cls.data.max()
    cls.put_hotspots_in_data()
    cls.data = np.tile(cls.data, (1, 1, 1))
    if __has_pycuda__:
        cls.ctx = get_cuda_context()


@pytest.mark.usefixtures('bootstrap')
class TestCCDCorrection:

    @classmethod
    def put_hotspots_in_data(cls):
        # Put 5 hot spots in the data
        # (row, column, deviation from median)
        cls.hotspots = [
            (50, 51, 0.04),
            (151, 150, 0.08),
            (202, 303, 0.12),
            (322, 203, 0.14)
        ]
        cls.threshold = 0.1 # parameterize ?
        data_medfilt = nabu_median_filter(cls.data)
        for r, c, deviation_from_median in cls.hotspots:
            cls.data[r, c] = data_medfilt[r, c] + deviation_from_median


    def check_detected_hotspots_locations(self, res):
        diff = self.data[0] - res[0]
        rows, cols = np.where(diff > 0)
        hotspots_arr = np.array(self.hotspots)
        M = hotspots_arr[:, -1] > self.threshold
        hotspots_rows = hotspots_arr[M, 0]
        hotspots_cols = hotspots_arr[M, 1]
        assert np.allclose(hotspots_rows, rows)
        assert np.allclose(hotspots_cols, cols)


    def test_median_clip(self):
        ccd_correction = CCDCorrection(
            self.data.shape,
            median_clip_thresh=self.threshold
        )
        res = ccd_correction.median_clip_correction(self.data, inplace=False)
        self.check_detected_hotspots_locations(res)

    @pytest.mark.skipif(not(__has_pycuda__), reason="Need cuda/pycuda for this test")
    def test_cuda_median_clip(self):
        d_radios = garray.to_gpu(self.data)
        cuda_ccd_correction = CudaCCDCorrection(
            d_radios.shape, median_clip_thresh=self.threshold
        )
        d_out = garray.zeros_like(d_radios)
        cuda_ccd_correction.median_clip_correction(d_radios, output=d_out)
        res = d_out.get()
        self.check_detected_hotspots_locations(res)

