import unittest

import moderngl
import numpy as np

from common import get_context


class TestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.ctx = get_context()

    def test_1(self):
        prog = self.ctx.program(
            vertex_shader='''
                #version 330

                in vec4 in_vert;
                out vec4 out_vert;

                void main() {
                    out_vert = in_vert;
                }
            ''',
            varyings=['out_vert']
        )

        vbo1 = self.ctx.buffer(np.array([4.0, 2.0, 7.5, 1.8], dtype='f4').tobytes())
        vbo2 = self.ctx.buffer(reserve=vbo1.size)
        vao = self.ctx.simple_vertex_array(prog, vbo1, 'in_vert')
        vao.transform(vbo2, moderngl.POINTS)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 2.0, 7.5, 1.8])

    def test_2(self):
        prog = self.ctx.program(
            vertex_shader='''
                #version 330

                in vec4 in_vert;
                out vec4 out_vert;

                void main() {
                    out_vert = in_vert;
                }
            ''',
            varyings=['out_vert']
        )

        vbo1 = self.ctx.buffer(np.array([4.0, 2.0, 7.5, 1.8], dtype='f4').tobytes())
        vbo2 = self.ctx.buffer(reserve=16)

        vao = self.ctx.vertex_array(prog, [(vbo1, '4f', 'in_vert')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 2.0, 7.5, 1.8])

        vao = self.ctx.vertex_array(prog, [(vbo1, '3f', 'in_vert')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 2.0, 7.5, 1.0])

        vao = self.ctx.vertex_array(prog, [(vbo1, '2f', 'in_vert')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 2.0, 0.0, 1.0])

        vao = self.ctx.vertex_array(prog, [(vbo1, '1f', 'in_vert')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 0.0, 0.0, 1.0])

    def test_3(self):
        prog = self.ctx.program(
            vertex_shader='''
                #version 330

                in mat4 in_mat;
                out mat4 out_mat;

                void main() {
                    out_mat = in_mat;
                }
            ''',
            varyings=['out_mat']
        )

        vbo1 = self.ctx.buffer(np.arange(1, 17, dtype='f4').tobytes())
        vbo2 = self.ctx.buffer(reserve=64)

        vao = self.ctx.vertex_array(prog, [(vbo1, '16f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [
            1, 2, 3, 4,
            5, 6, 7, 8,
            9, 10, 11, 12,
            13, 14, 15, 16,
        ])

        vao = self.ctx.vertex_array(prog, [(vbo1, '12f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [
            1, 2, 3, 1,
            4, 5, 6, 1,
            7, 8, 9, 1,
            10, 11, 12, 1,
        ])

        vao = self.ctx.vertex_array(prog, [(vbo1, '8f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [
            1, 2, 0, 1,
            3, 4, 0, 1,
            5, 6, 0, 1,
            7, 8, 0, 1,
        ])

        vao = self.ctx.vertex_array(prog, [(vbo1, '4f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [
            1, 0, 0, 1,
            2, 0, 0, 1,
            3, 0, 0, 1,
            4, 0, 0, 1,
        ])

    def test_4(self):
        prog = self.ctx.program(
            vertex_shader='''
                #version 330

                in mat2 in_mat;
                out mat2 out_mat;

                void main() {
                    out_mat = in_mat;
                }
            ''',
            varyings=['out_mat']
        )

        vbo1 = self.ctx.buffer(np.array([4.0, 2.0, 7.5, 1.8], dtype='f4').tobytes())
        vbo2 = self.ctx.buffer(reserve=16)

        vao = self.ctx.vertex_array(prog, [(vbo1, '4f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 2.0, 7.5, 1.8])

        vao = self.ctx.vertex_array(prog, [(vbo1, '2f', 'in_mat')])
        vao.transform(vbo2, moderngl.POINTS, vertices=1)
        res = np.frombuffer(vbo2.read(), dtype='f4')
        np.testing.assert_almost_equal(res, [4.0, 0.0, 2.0, 0.0])


if __name__ == '__main__':
    unittest.main()
