import os
import shutil
import tempfile
from unittest import TestCase

import numpy as np

from yt.testing import assert_almost_equal, fake_random_ds
from yt.visualization.volume_rendering.api import (
    OpaqueSource,
    Scene,
    ZBuffer,
    create_volume_source,
)


class FakeOpaqueSource(OpaqueSource):
    # A minimal (mock) concrete implementation of OpaqueSource
    def render(self, camera, zbuffer=None):
        pass

    def _validate(self):
        pass


def setup():
    """Test specific setup."""
    from yt.config import ytcfg

    ytcfg["yt", "internals", "within_testing"] = True


class ZBufferTest(TestCase):
    # This toggles using a temporary directory. Turn off to examine images.
    use_tmpdir = True

    def setUp(self):
        np.random.seed(0)
        if self.use_tmpdir:
            self.curdir = os.getcwd()
            # Perform I/O in safe place instead of yt main dir
            self.tmpdir = tempfile.mkdtemp()
            os.chdir(self.tmpdir)
        else:
            self.curdir, self.tmpdir = None, None

    def tearDown(self):
        if self.use_tmpdir:
            os.chdir(self.curdir)
            shutil.rmtree(self.tmpdir)

    def test_composite_vr(self):
        ds = fake_random_ds(64)
        dd = ds.sphere(ds.domain_center, 0.45 * ds.domain_width[0])
        ds.field_info[ds.field_list[0]].take_log = False

        sc = Scene()
        cam = sc.add_camera(ds)
        cam.resolution = (512, 512)
        vr = create_volume_source(dd, field=ds.field_list[0])
        vr.transfer_function.clear()
        vr.transfer_function.grey_opacity = True
        vr.transfer_function.map_to_colormap(0.0, 1.0, scale=10.0, colormap="Reds")
        sc.add_source(vr)

        cam.set_width(1.8 * ds.domain_width)
        cam.lens.setup_box_properties(cam)

        # Create Arbitrary Z-buffer
        empty = cam.lens.new_image(cam)
        z = np.empty(empty.shape[:2], dtype="float64")
        # Let's put a blue plane right through the center
        z[:] = cam.width[2] / 2.0
        empty[:, :, 2] = 1.0  # Set blue to 1's
        empty[:, :, 3] = 1.0  # Set alpha to 1's
        zbuffer = ZBuffer(empty, z)
        zsource = FakeOpaqueSource()
        zsource.set_zbuffer(zbuffer)
        sc.add_source(zsource)

        im = sc.render()
        im.write_png("composite.png")

    def test_nonrectangular_add(self):
        rgba1 = np.ones((64, 1, 4))
        z1 = np.expand_dims(np.arange(64.0), 1)

        rgba2 = np.zeros((64, 1, 4))
        z2 = np.expand_dims(np.arange(63.0, -1.0, -1.0), 1)

        exact_rgba = np.concatenate((np.ones(32), np.zeros(32)))
        exact_rgba = np.expand_dims(exact_rgba, 1)
        exact_rgba = np.dstack((exact_rgba, exact_rgba, exact_rgba, exact_rgba))

        exact_z = np.concatenate((np.arange(32.0), np.arange(31.0, -1.0, -1.0)))
        exact_z = np.expand_dims(exact_z, 1)

        buff1 = ZBuffer(rgba1, z1)
        buff2 = ZBuffer(rgba2, z2)

        buff = buff1 + buff2

        assert_almost_equal(buff.rgba, exact_rgba)
        assert_almost_equal(buff.z, exact_z)

    def test_rectangular_add(self):
        rgba1 = np.ones((8, 8, 4))
        z1 = np.arange(64.0)
        z1 = z1.reshape((8, 8))
        buff1 = ZBuffer(rgba1, z1)

        rgba2 = np.zeros((8, 8, 4))
        z2 = np.arange(63.0, -1.0, -1.0)
        z2 = z2.reshape((8, 8))
        buff2 = ZBuffer(rgba2, z2)

        buff = buff1 + buff2

        exact_rgba = np.empty((8, 8, 4), dtype=np.float64)
        exact_rgba[0:4, 0:8, :] = 1.0
        exact_rgba[4:8, 0:8, :] = 0.0

        exact_z = np.concatenate((np.arange(32.0), np.arange(31.0, -1.0, -1.0)))
        exact_z = np.expand_dims(exact_z, 1)
        exact_z = exact_z.reshape(8, 8)

        assert_almost_equal(buff.rgba, exact_rgba)
        assert_almost_equal(buff.z, exact_z)
