#! /opt/local/bin/python
# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2023-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

"""
This file includes tests for the holotomo python API.
"""

import os
import unittest
import tempfile
import shutil
import sys
import io
import warnings
import gc
import numpy as np
import h5py as h5
from skimage.data import shepp_logan_phantom, coins, immunohistochemistry, microaneurysms, retina
from pynx.processing_unit import has_cuda, has_opencl
from pynx.holotomo.utils import simulate_probe
from pynx.holotomo import *
from pynx.wavefront import Wavefront, PropagateNearField

if has_cuda:
    import pynx.holotomo.cu_operator as cuop

import pynx.holotomo.cpu_operator as cpuop

exclude_cuda = False
exclude_opencl = False
if 'PYNX_PU' in os.environ:
    if 'opencl' in os.environ['PYNX_PU'].lower():
        exclude_cuda = True
    elif 'cuda' in os.environ['PYNX_PU'].lower():
        exclude_opencl = True

if 'opencl' in sys.argv or '--opencl' in sys.argv or not has_cuda:
    exclude_cuda = True
if 'cuda' in sys.argv or '--cuda' in sys.argv or not has_opencl:
    exclude_opencl = True


def make_obj_probe(nproj=20, nz=4, ny=512, nx=400, wavelength=0.5e-10,
                   delta=1e-5, beta=1e-7, thickness=1e-6, pixel_size=50e-9, dz0=1e-1):
    # TODO: actually create projections from a tomography acquisition
    # Propagation distances
    vz = np.linspace(dz0, dz0 * 1.5, nz)
    # Create probe
    probe = simulate_probe((ny, nx), vz, nb_line_h=10, nb_line_v=10,
                           pixel_size=pixel_size, wavelength=wavelength, amplitude=2)

    # Create object projections from skimage examples
    k = 2 * np.pi / wavelength

    img_src = [shepp_logan_phantom, immunohistochemistry, retina, coins]

    obj = np.ones((nproj, ny, nx), dtype=np.complex64)
    i, ct = 0, 0
    while ct < nproj:
        img = img_src[i]()
        if img.ndim == 2:
            img = img[..., np.newaxis]
        for j in range(img.shape[-1]):
            d = img[:, :, j]
            if d.shape[0] > ny:
                d = d[:ny]
            if d.shape[1] > nx:
                d = d[:, :nx]
            ny1, nx1 = d.shape
            obj[ct, :ny1, :nx1] = d / d.max()

            if ny1 < ny:
                obj[ct] = np.roll(obj[ct], (ny - ny1) // 2, axis=0)
            if nx1 < nx:
                obj[ct] = np.roll(obj[ct], (nx - nx1) // 2, axis=1)
            ct += 1
            if ct == nproj:
                break
        i = (i + 1) % len(img_src)

    for i in range(0, nproj):
        obj[i] = np.exp(-k * (-1j * delta + beta) * thickness * obj[i])

    return obj, probe, vz


def make_data(obj, probe, vz, wavelength=0.5e-10, nb_photons=1e4, pixel_size=50e-9):
    # Create propagated projections
    nproj, ny, nx = obj.shape
    nz = len(vz)
    data = np.empty((nproj, nz, ny, nx), dtype=np.float32)
    w = Wavefront(d=data[0], pixel_size=pixel_size, wavelength=wavelength)
    for i in range(nproj):
        for iz in range(nz):
            w.set(obj[i] * probe[iz], shift=True)
            w = PropagateNearField(dz=vz[iz]) * w
            data[i, iz] = abs(w.get(shift=True).reshape(ny, nx)) ** 2
            data[i, iz] *= nb_photons / data[i, iz].mean()
    # Empty beam images
    ref = np.empty((nz, ny, nx), dtype=np.float32)
    for iz in range(nz):
        w.set(probe[iz], shift=True)
        w = PropagateNearField(dz=vz[iz]) * w
        ref[iz] = abs(w.get(shift=True).reshape(ny, nx)) ** 2
        ref[iz] *= nb_photons / data[i, iz].mean()
    return data, ref


class TestHoloTomo(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        # Directory contents will automatically get cleaned up on deletion
        cls.tmp_dir_obj = tempfile.TemporaryDirectory()
        cls.tmp_dir = cls.tmp_dir_obj.name
        # Holotomo data and main object, created in make_data and make_obj
        cls.pixel_size = 50e-9
        cls.wavelength = 0.5e-10
        cls.data = None
        cls.ref = None
        cls.p = None
        cls.obj0 = None
        cls.probe0 = None
        cls.vz = None
        cls.ops = []
        if 'cuda' not in sys.argv and 'opencl' not in sys.argv:
            cls.ops.append(cpuop)
        if has_cuda and 'opencl' not in sys.argv:
            cls.ops.append(cuop)
        # if has_opencl and 'cuda' not in sys.argv:
        #     cls.ops.append(clop)

    @classmethod
    def tearDownClass(cls):
        if cls.data is not None:
            del cls.data
        if cls.p is not None:
            del cls.p
        gc.collect()

    def make_data(self):
        """
        Make ptycho data, if it does not already exist.
        :return: Nothing
        """
        if self.data is None:
            self.make_obj_probe()
            self.data, self.ref = make_data(self.obj0, self.probe0, self.vz)

    def make_obj_probe(self):
        """
        Make ptycho obj, if it does not already exist.
        :return: Nothing
        """
        if self.p is None:
            self.obj0, self.probe0, self.vz = make_obj_probe()

    def test_00_make_obj_probe(self):
        self.make_obj_probe()

    def test_01_make_data(self):
        self.make_data()

    def test_make_holotomo_data(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)

    def test_make_holotomo(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h)

    def test_make_holotomo_o(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h, obj=self.obj0.copy())

    def test_make_holotomo_p(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h, probe=self.probe0.copy())

    def test_make_holotomo_op(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h, obj=self.obj0.copy(), probe=self.probe0.copy())

    def test_save_chunk_hdf5(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h, obj=self.obj0.copy(), probe=self.probe0.copy())
        p.save_obj_probe_chunk(chunk_prefix=os.path.join(self.tmp_dir, "chunk_%02d"), nxtomo=False)

    def test_save_nxtomo(self):
        self.make_data()
        h = HoloTomoData(self.data, self.ref, pixel_size_detector=self.pixel_size, wavelength=self.wavelength,
                         detector_distance=self.vz)
        p = HoloTomo(data=h, obj=self.obj0.copy(), probe=self.probe0.copy())
        fname = os.path.join(self.tmp_dir, "data.nx")
        p.save_obj_probe_chunk(chunk_prefix=fname, nxtomo=True,
                               sample_name="test", angles=np.linspace(0, 2 * np.pi, h.nproj))
        self.assertTrue(os.path.exists(fname))
        with h5.File(fname) as nx:
            self.assertTrue(nx['/entry_1/data/data'].shape == (h.nproj, h.ny, h.nx))
            self.assertTrue(len(nx['/entry_1/data/rotation_angle']) == h.nproj)
            self.assertTrue(len(nx['/entry_1/data/image_key']) == h.nproj)

        p.save_obj_probe_chunk(chunk_prefix=os.path.join(self.tmp_dir, "data.nx"), nxtomo=True,
                               sample_name="test", angles=np.linspace(0, 2 * np.pi, h.nproj), append=False)
        with h5.File(fname) as nx:
            self.assertTrue(nx['/entry_1/data/data'].shape == (h.nproj, h.ny, h.nx))
            self.assertTrue(len(nx['/entry_1/data/rotation_angle']) == h.nproj)
            self.assertTrue(len(nx['/entry_1/data/image_key']) == h.nproj)
        h.idx += 1
        p.save_obj_probe_chunk(chunk_prefix=os.path.join(self.tmp_dir, "data.nx"), nxtomo=True,
                               sample_name="test", angles=np.linspace(0, 2 * np.pi, h.nproj), append=True)
        with h5.File(fname) as nx:
            self.assertTrue('/entry_1/data/data' in nx)
            self.assertTrue('/entry_1/data/rotation_angle' in nx)
            self.assertTrue('/entry_1/data/image_key' in nx)
            self.assertTrue(nx['/entry_1/data/data'].shape == (2 * h.nproj, h.ny, h.nx))
            self.assertTrue(len(nx['/entry_1/data/rotation_angle']) == 2 * h.nproj)
            self.assertTrue(len(nx['/entry_1/data/image_key']) == 2 * h.nproj)


def suite():
    test_suite = unittest.TestSuite()
    loadTests = unittest.defaultTestLoader.loadTestsFromTestCase
    test_suite.addTest(loadTests(TestHoloTomo))
    return test_suite


if __name__ == '__main__':
    sys.stdout = io.StringIO()
    warnings.simplefilter('ignore')
    res = unittest.TextTestRunner(verbosity=2, descriptions=False).run(suite())
