#!/usr/bin/python3
# copyright 2023 weber yann
#
# this file is part of rpnifs.
#
#        geneifs is free software: you can redistribute it and/or modify
#        it under the terms of the gnu general public license as published by
#        the free software foundation, either version 3 of the license, or
#        (at your option) any later version.
#
#        geneifs is distributed in the hope that it will be useful,
#        but without any warranty; without even the implied warranty of
#        merchantability or fitness for a particular purpose.  see the
#        gnu general public license for more details.
#
#        you should have received a copy of the gnu general public license
#        along with geneifs.  if not, see <http://www.gnu.org/licenses/>.
#

import copy
import pickle
import random
import sys

import unittest

try:
    import pyrpn
except (ImportError, NameError) as e:
    print("error importing pyrpn. try to run make.",
          file=sys.stderr)
    raise e

from utils import Progress


class TestRpnExprCopy(unittest.TestCase):
    """ Testing RPNExpr sequence method """


import unittest
import warnings

try:
    import numpy as np
except (ImportError, NameError) as e:
    np = None
    warnings.warn("Numpy not found, some tests skipped", warnings.ImportWarning)

try:
    import pyrpn
except (ImportError, NameError) as e:
    print("Error importing pyrpn. Try to run make.",
          file=sys.stderr)
    raise e

def skipIfNoNumpy():
    if np is not None:
        return lambda func: func
    return unittest.skip("Numpy not found")



class TestRPNIterInit(unittest.TestCase):

    def test_init(self):
        """ Test instanciation """
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                                pyrpn.const.RESULT_COUNT,
                                (640,480))
    
    def test_params_simple(self):
        """ Test parameters fetch from instanciation """
        args = [(pyrpn.const.POS_XY,
                 pyrpn.const.RESULT_COUNT,
                 (640,480)),
                (pyrpn.const.POS_LINEAR,
                 pyrpn.const.RESULT_BOOL,
                 (1024,))
        ]
        for arg in args:
            with self.subTest(args=arg):
                rif = pyrpn.RPNIterExpr(*arg)
                params = rif.get_params()
                self.assertEqual(params.pos_flag, arg[0])
                self.assertEqual(params.res_flag, arg[1])
                self.assertEqual(params.size_lim, arg[2])
                self.assertIsNone(params.const_values)
    
    def test_params_const(self):
        """ Test parameters from instanciation when const values used as result """
        args = [(pyrpn.const.POS_XY,
                 pyrpn.const.RESULT_CONST,
                 (640,480),
                 (42,)),
                (pyrpn.const.POS_LINEAR,
                 pyrpn.const.RESULT_CONST_RGBA,
                 (1024,), (13,37,13,12))
        ]
        for arg in args:
            with self.subTest(args=arg):
                rif = pyrpn.RPNIterExpr(*arg)
                params = rif.get_params()
                self.assertEqual(params.pos_flag, arg[0])
                self.assertEqual(params.res_flag, arg[1])
                self.assertEqual(params.size_lim, arg[2])
                self.assertEqual(params.const_values, arg[3])

    def test_params_xdim(self):
        """ Test parameters from instanciation when using xdim positionning """
        args = [(pyrpn.const.POS_XDIM,
                 pyrpn.const.RESULT_BOOL,
                 (2,640,480)),
                (pyrpn.const.POS_XDIM,
                 pyrpn.const.RESULT_BOOL,
                 (5,13,37,13,12,42)),
        ]
        for arg in args:
            with self.subTest(args=arg):
                rif = pyrpn.RPNIterExpr(*arg)
                params = rif.get_params()
                self.assertEqual(params.pos_flag, arg[0])
                self.assertEqual(params.res_flag, arg[1])
                self.assertEqual(params.size_lim, arg[2])
                self.assertIsNone(params.const_values)
    
    def test_shape(self):
        """ Test the shape method """
        tests = [((pyrpn.const.POS_XY,
                          pyrpn.const.RESULT_COUNT,
                          (640,480)), (640,480)),
                        ((pyrpn.const.POS_LINEAR,
                          pyrpn.const.RESULT_BOOL,
                          (1024,)), (1024,)),
                        ((pyrpn.const.POS_LINEAR,
                          pyrpn.const.RESULT_RGBA,
                          (1024,)), (1024,4)),
                        ((pyrpn.const.POS_XDIM,
                          pyrpn.const.RESULT_BOOL,
                          (2,640,480)),  (640,480)),
                        ((pyrpn.const.POS_XDIM,
                          pyrpn.const.RESULT_RGB,
                          (5,13,37,13,12,42)), (13,37,13,12,42,3)),
                ]
        for args, expt in tests:
            with self.subTest(args=args, expt_shape=expt):
                rif = pyrpn.RPNIterExpr(*args)
                self.assertEqual(rif.shape(), expt)

    @skipIfNoNumpy()
    def test_buffer(self):
        """ Test the len on buffer interface using numpy """
        args = [((pyrpn.const.POS_XY,
                  pyrpn.const.RESULT_COUNT,
                  (640,480)), 640*480),
                ((pyrpn.const.POS_LINEAR,
                  pyrpn.const.RESULT_BOOL,
                  (1024,)), 1024),
                ((pyrpn.const.POS_XDIM,
                  pyrpn.const.RESULT_BOOL,
                  (2,640,480)),  640*480),
                ((pyrpn.const.POS_XDIM,
                  pyrpn.const.RESULT_BOOL,
                  (5,13,37,13,12,42)), 13*37*13*12*42),
        ]
        
        for arg, expt_sz in args:
                with self.subTest(args=arg):
                    rif = pyrpn.RPNIterExpr(*arg)
                    arr = np.frombuffer(rif, dtype=np.uint64)
                    self.assertEqual(len(arr), expt_sz)
                    arr += 0x11111111 # check writing to all bytes

    def test_str(self):
        """ Test string representation of rif """

        tests = [[pyrpn.const.POS_XY, pyrpn.const.RESULT_RGB, (1024,1024)],
                 [pyrpn.const.POS_LINEAR, pyrpn.const.RESULT_CONST_RGBA, (1024,), (255,0,0,255)],
                 [pyrpn.const.POS_XY, pyrpn.const.RESULT_CONST_RGBA, (640,480), (255,0,0,255)],
                 [pyrpn.const.POS_LINEAR, pyrpn.const.RESULT_RGBA, (1024,)],
                 [pyrpn.const.POS_XDIM, pyrpn.const.RESULT_RGB, (4, 2, 2, 640, 480)],
                 [pyrpn.const.POS_XY, pyrpn.const.RESULT_CONST, (1024,512),
                     (2,)],
        ]
        for args in tests:
            rif = pyrpn.RPNIterExpr(*args)
            [expr.mutate(n_mutations=5) for expr in rif]
            codes = {spl[1]:spl[2].strip('"')
                     for spl in [s.split('=') for s in str(rif).split(';')
                                 if len(s.strip())]}
            argc = rif.get_params().argc
            for key, orig in rif.items():
                expr = pyrpn.RPNExpr(codes[key], argc)
                with self.subTest(rif=rif, key=key, orig=orig, clone=expr):
                    self.assertEqual(orig, expr)
            str_repr = repr(rif)


class TestRPNIterCoordinates(unittest.TestCase):
    """ Testing methods for coordinates <-> position convertions """

    def test_position_linear(self):
        """ Testing linear coordinate convertion methods """
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_LINEAR,
                    pyrpn.const.RESULT_CONST,
                    (256,), (42,))

        for pos in range(256):
            with self.subTest(rif=rif, position=pos, expt=(pos,)):
                res = rif.from_pos(pos)
                self.assertEqual(res, (pos,))
            with self.subTest(rif=rif, expt=pos, coord=(pos,)):
                res = rif.to_pos(pos)
                self.assertEqual(res, pos)

    def test_position_xy(self):
        """ Testing XY coordinate convertion methods """
        for _ in Progress(5):
            sz = (random.randint(32,128), random.randint(32,128))
            rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                    pyrpn.const.RESULT_CONST,
                    sz, (42,))
            for pos in range(sz[0]*sz[1]):
                coord = rif.from_pos(pos)
                expt = (pos % sz[0], pos // sz[0])
                with self.subTest(sz=sz, pos=pos, expt=expt, result=coord):
                    self.assertEqual(expt, coord)
                result = rif.to_pos(*coord)
                with self.subTest(sz=sz, coord=coord, expt=pos, result=result):
                    self.assertEqual(result, pos)

    def test_position_xy_neg(self):
        """ Testing XY coordinates with negative values """
        sz = (100,200)
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                pyrpn.const.RESULT_CONST,
                sz, (42,))
        
        for _ in range(1000):
            ny = random.randint(-0x1000000, sz[1])
            nx = random.randint(0,sz[0]-1)
            pos = rif.to_pos(nx, ny)
            coord = rif.from_pos(pos)
            self.assertEqual(coord, (nx, (ny%(1<<64))%sz[1]))

    def test_position_xy_random(self):
        """ Testing XY coordinate convertion methods (random samples)"""
        for _ in Progress(256):
            sz = (random.randint(32,4096), random.randint(32,4096))
            rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                    pyrpn.const.RESULT_CONST,
                    sz, (42,))
            for _ in range(500):
                pos = random.randint(0, (sz[0]*sz[1])-1)
                coord = rif.from_pos(pos)
                expt = (pos % sz[0], pos // sz[0])
                with self.subTest(sz=sz, pos=pos, expt=expt, result=coord):
                    self.assertEqual(expt, coord)
                result = rif.to_pos(*coord)
                with self.subTest(sz=sz, coord=coord, expt=pos, result=result):
                    self.assertEqual(result, pos)

    def test_position_xdim(self):
        """ Testing XDIM coordinate convertion methods """
        ndim = 5
        resol = [8 for _ in range(ndim)]
        sz = [ndim]+resol
        linear_size = 1
        for e in resol:
            linear_size *= e
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XDIM,
                pyrpn.const.RESULT_CONST,
                sz, (42,))
        for pos in range(linear_size):
            coord = rif.from_pos(pos)
            expt = []
            cur = pos
            for dim in resol:
                elt = cur % dim
                cur //= dim
                expt.append(elt)
            expt = tuple(expt)
            with self.subTest(sz=sz, pos=pos, expt=expt, result=coord):
                self.assertEqual(expt, coord)
            result = rif.to_pos(*coord)
            with self.subTest(sz=sz, coord=coord, expt=pos, result=result):
                self.assertEqual(result, pos)

    def test_position_xdim_neg(self):
        """ Testing XDIM negative position convertion """
        # TODO test negativ for other coordinate systems
        ndim = 4
        szlim = (ndim, 10,20,30,40)
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XDIM,
                pyrpn.const.RESULT_CONST,
                szlim, (42,))

        pos = rif.to_pos(0,0,-5,0)
        self.assertEqual(rif.from_pos(pos), (0,0,(-5%(1<<64))%30, 0))

    def test_position_xdim(self):
        """ Testing XDIM coordinate convertion methods (random sampling)"""
        for _ in Progress(16):
            ndim = random.randint(3,8)
            resol = [random.randint(2, 16) for _ in range(ndim)]
            sz = [ndim]+resol
            linear_size = 1
            for e in resol:
                linear_size *= e
            rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XDIM,
                    pyrpn.const.RESULT_CONST,
                    sz, (42,))
            for _ in range(2000):
                pos = random.randint(0, linear_size-1)
                coord = rif.from_pos(pos)
                expt = []
                cur = pos
                for dim in resol:
                    elt = cur % dim
                    cur //= dim
                    expt.append(elt)
                expt = tuple(expt)
                with self.subTest(sz=sz, pos=pos, expt=expt, result=coord):
                    self.assertEqual(expt, coord)
                result = rif.to_pos(*coord)
                with self.subTest(sz=sz, coord=coord, expt=pos, result=result):
                    self.assertEqual(result, pos)


class TestRPNIterConst(unittest.TestCase):
    """ Testing various coordinate systems with constant result value """

    @skipIfNoNumpy()
    def test_simple_linear_const(self):
        """ Testing a simple constant result on linear coord  """
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_LINEAR,
                                pyrpn.const.RESULT_CONST,
                                (512,), (42,))

        data = np.frombuffer(rif, dtype=np.uint64)
        for elt in data:
            self.assertEqual(elt, 0)

        rif['X'] = 'A0 A1 +'
        pos = rif.step(0)
        self.assertEqual(pos, 0)
        self.assertEqual(data[0], 42)

        data = np.frombuffer(rif, dtype=np.uint64)
        self.assertEqual(data[0], 42)
        for elt in data[1:]:
            self.assertEqual(elt, 0)



        pos = rif.step(0)
        self.assertEqual(pos, 42)

        data = np.frombuffer(rif, dtype=np.uint64)
        for i, elt in enumerate(data):
            with self.subTest(cur_pos=i):
                if i in (0, 42):
                    self.assertEqual(elt, 42)
                else:
                    self.assertEqual(elt, 0)


        pos = rif.step(1)
        self.assertEqual(pos, 1)
        pos = rif.step(pos)
        self.assertEqual(pos, 43)

        newpos = 512 - 40
        pos = rif.step(newpos)
        self.assertEqual(pos, newpos)
        pos = rif.step(pos)
        self.assertEqual(pos, 2)

    def test_random_linear_const(self):
        """ Testing linear coord with const with random expressions """
        for _ in range(200):
            const_val = random.randint(0,0xFFFF)
            rif = pyrpn.RPNIterExpr(pyrpn.const.POS_LINEAR,
                                    pyrpn.const.RESULT_CONST,
                                    (512,), (const_val,))
            rif['X'].mutate(n_mutations=random.randint(10,100))
            pos=0
            all_pos=[]
            for  _ in range(100):
                pos = rif.step(pos)
                all_pos.append(pos)
                data = np.frombuffer(rif, dtype=np.uint64)
                self.assertEqual(data[pos], const_val)
            for i, elt in enumerate(np.frombuffer(rif, dtype=np.uint64)):
                if i in all_pos:
                    self.assertEqual(elt, const_val)
                else:
                    self.assertEqual(elt, 0)

    def test_simple_xy_const(self):
        """ Testing xy coord with const value """
        sz = (1024,256)

        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                                pyrpn.const.RESULT_CONST,
                                sz, (42,))
        
        rif['X'] = 'A0 A1 +'
        rif['Y'] = 'A2'

        pos = rif.step(0)
        self.assertEqual(pos, 0)

        pos = rif.step(0)
        self.assertEqual(rif.from_pos(pos), (0,42))

        pos = rif.step(rif.to_pos(1,1))
        self.assertEqual(rif.from_pos(pos), (2,0))

        pos = rif.step(rif.to_pos(1,1))
        self.assertEqual(rif.from_pos(pos), (2,0))

        pos = rif.step(rif.to_pos(2,0))
        self.assertEqual(rif.from_pos(pos), (2,42))

    def test_random_xy_const(self):
        """ Testing xy coord with const with random expressions """
        for _ in Progress(200):
            const_val = random.randint(0,0xFFFF)
            rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XY,
                                    pyrpn.const.RESULT_CONST,
                                    (1024,1024,), (const_val,))
            rif['X'].mutate(n_mutations=random.randint(10,100))
            rif['Y'].mutate(n_mutations=random.randint(10,100))
            pos=0
            all_pos=[]
            for  _ in range(100):
                sys.stdout.flush()
                pos = rif.step(pos)
                all_pos.append(pos)
                data = np.frombuffer(rif, dtype=np.uint64)
                self.assertEqual(data[pos], const_val)
            expt = np.zeros(len(data), dtype=np.uint64)
            for p in all_pos:
                expt[p] = const_val
            self.assertTrue((expt == data).all())

    def test_simple_xdim_const(self):
        """ Testing xdim coord with const """
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_XDIM,
                                pyrpn.const.RESULT_CONST,
                                (4,100,200,300,400), (42,))
        rif['D0'] = 'A1'
        rif['D1'] = 'A0 A2 +'
        rif['D2'] = 'A0 A1 -'
        rif['D3'] = 'A3 A4 +'

        pos = rif.step(0)
        self.assertEqual(pos, 0)
        pos = rif.step(0)
        self.assertEqual(rif.from_pos(pos), (0,0,0,42))

        pos = rif.step(5)
        self.assertEqual(rif.from_pos(pos), (0,5,5,0))
        pos = rif.step(rif.to_pos(0,5,5,0))
        self.assertEqual(rif.from_pos(pos), (5,5,(-5%(1<<64))%300,42))
        expt = (5,5,-5,42)
        expt_pos = rif.to_pos(*expt)
        self.assertEqual(pos, expt_pos)

    def test_linear_rgb(self):
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_LINEAR,
                                pyrpn.const.RESULT_RGB,
                                (512,))
        rif['X'] = 'A0 2 +'
        rif['R'] = 'A0 1 +'
        rif['G'] = 'A1 A2 + 255 %'
        rif['B'] = 'A2 2 * A3 +'
        
        pos = rif.step(0)
        self.assertEqual(pos, 2)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[2])
        self.assertEqual(colors, [1,0,0])

        pos = rif.step(2)
        self.assertEqual(pos, 4)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[4])
        self.assertEqual(colors, [3, 1, 0])

    def test_linear_rgba(self):
        rif = pyrpn.RPNIterExpr(pyrpn.const.POS_LINEAR,
                                pyrpn.const.RESULT_RGBA,
                                (512,))
        rif['X'] = 'A0 2 +'
        rif['R'] = 'A0 1 +'
        rif['G'] = 'A1 A2 + 255 %'
        rif['B'] = 'A2 2 * A3 +'
        rif['A'] = 'A1 A2 A3 + +'
        
        pos = rif.step(0)
        self.assertEqual(pos, 2)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[2])
        self.assertEqual(colors, [1,0,0,0])

        pos = rif.step(2)
        self.assertEqual(pos, 4)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[4])
        self.assertEqual(colors, [3, 1, 0, 1])

        pos = rif.step(4)
        self.assertEqual(pos, 6)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[6])
        self.assertEqual(colors, [5, 4, 2, 4])

        pos = rif.step(6)
        self.assertEqual(pos, 8)
        data = np.frombuffer(rif, dtype=np.uint64).reshape(rif.shape())
        colors = list(data[8])
        self.assertEqual(colors, [7, 9, 10, 11])


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