#!/usr/bin/env python3

import unittest

import gi
gi.require_version('v_sim', '3.9')
from gi.repository import GLib, v_sim

import signals

class TestNodeArray(unittest.TestCase):
    def setUp(self):
        super(TestNodeArray, self).setUp()
        self.addTypeEqualityFunc(list, self.fuzzyList)
        self.addTypeEqualityFunc(float, self.fuzzyFloat)

    def fuzzyList(self, a, b, msg = None):
        if len(a) > 0 and isinstance(a[0], float) and len(a) == len(b):
            try:
                for i, j in zip(a, b):
                    self.fuzzyFloat(i, j, msg)
            except (self.failureException):
              raise self.failureException(msg if msg is not None else "" + "delta %s" % ([abs(i - j) for i, j in zip(a, b)]))
            return
        return self.assertListEqual(a, b, msg)
        
    def fuzzyFloat(self, a, b, msg = None):
        if abs(b-a) > 6e-8:
            raise self.failureException(msg)
        
    def test_population(self):
        # Empty structure.
        array = v_sim.Data.new()
        self.assertEqual(array.getNElements(True), 0)
        self.assertEqual(array.getNElements(False), 0)
        self.assertEqual(array.getNNodes(), 0)

        # Population allocation
        array.allocateByNames((1, 2, 45), ("O", "H", "g"))
        self.assertTrue(v_sim.Element.lookup("O") is not None)
        self.assertTrue(v_sim.Element.lookup("H") is not None)
        self.assertTrue(v_sim.Element.lookup("g") is not None)
        self.assertEqual(array.getNElements(True), 0)
        self.assertEqual(array.getNElements(False), 0)
        self.assertEqual(array.getNNodes(), 0)
        self.assertTrue(array.containsElement(v_sim.Element.lookup("O")))
        self.assertTrue(array.containsElement(v_sim.Element.lookup("H")))
        self.assertTrue(array.containsElement(v_sim.Element.lookup("g")))
        self.assertEqual(array.getElementId(v_sim.Element.lookup("O")), 0)
        self.assertEqual(array.getElementId(v_sim.Element.lookup("H")), 1)
        self.assertEqual(array.getElementId(v_sim.Element.lookup("g")), 2)

        # Population partial re-allocation
        array.allocateByNames((1, 4), ("C", "H"))
        self.assertTrue(v_sim.Element.lookup("C") is not None)
        self.assertEqual(array.getNElements(True), 0)
        self.assertEqual(array.getNElements(False), 0)
        self.assertEqual(array.getNNodes(), 0)
        self.assertTrue(array.containsElement(v_sim.Element.lookup("O")))
        self.assertTrue(array.containsElement(v_sim.Element.lookup("H")))
        self.assertTrue(array.containsElement(v_sim.Element.lookup("g")))
        self.assertTrue(array.containsElement(v_sim.Element.lookup("C")))
        self.assertEqual(array.getElementId(v_sim.Element.lookup("O")), 0)
        self.assertEqual(array.getElementId(v_sim.Element.lookup("H")), 1)
        self.assertEqual(array.getElementId(v_sim.Element.lookup("g")), 2)
        self.assertEqual(array.getElementId(v_sim.Element.lookup("C")), 3)

        # Node adding
        node = array.getNewNode(v_sim.Element.lookup("O"))
        self.assertTrue(node is not None)
        self.assertEqual(node.xyz, [0., 0., 0.])
        self.assertEqual(node.translation, [0., 0., 0.])
        self.assertTrue(not node.getVisibility())
        node.newValues((1,2,3))
        self.assertEqual(node.number, 0)
        self.assertEqual(node.xyz, [1., 2., 3.])
        self.assertEqual(node.translation, [0., 0., 0.])
        self.assertTrue(node.getVisibility())
        self.assertEqual(array.getNElements(True), 1)
        self.assertEqual(array.getNElements(False), 1)
        self.assertEqual(array.getNNodes(), 1)
        self.assertEqual(array.getNOriginalNodes(), 1)
        self.assertTrue(array.getElement(node), v_sim.Element.lookup("O"))
        self.assertEqual(array.getFromId(0), node)
        self.assertTrue(array.getFromId(1) is None)
        self.assertEqual(array.getOriginal(0), -1)
        self.assertEqual(array.get_property("n-nodes"), 1)

        node = array.getNewNodeForId(1)
        self.assertTrue(node is not None)
        node.newValues((2,0,0))
        self.assertEqual(node.number, 1)
        self.assertEqual(node.xyz, [2., 0., 0.])
        self.assertEqual(node.translation, [0., 0., 0.])
        self.assertTrue(node.getVisibility())
        self.assertEqual(array.getNElements(True), 2)
        self.assertEqual(array.getNElements(False), 2)
        self.assertEqual(array.getNNodes(), 2)
        self.assertEqual(array.getNOriginalNodes(), 2)
        self.assertTrue(array.getElement(node), v_sim.Element.lookup("H"))
        self.assertEqual(array.getFromId(1), node)
        self.assertTrue(array.getFromId(0) is not None)
        self.assertTrue(array.getFromId(2) is None)
        self.assertEqual(array.getOriginal(1), -1)

        node = array.copyNode(node)
        self.assertTrue(node is not None)
        node.newValues((-2,0,0))
        self.assertEqual(node.number, 2)
        self.assertEqual(node.xyz, [-2., 0., 0.])
        self.assertEqual(node.translation, [0., 0., 0.])
        self.assertTrue(node.getVisibility())
        self.assertEqual(array.getNElements(True), 2)
        self.assertEqual(array.getNElements(False), 2)
        self.assertEqual(array.getNNodes(), 3)
        self.assertEqual(array.getNOriginalNodes(), 2)
        self.assertTrue(array.getElement(node), v_sim.Element.lookup("H"))
        self.assertEqual(array.getFromId(2), node)
        self.assertTrue(array.getFromId(1) is not None)
        self.assertTrue(array.getFromId(3) is None)
        self.assertEqual(array.getOriginal(2), 1)
        self.assertTrue(array.setOriginal(2))
        self.assertEqual(array.getOriginal(2), -1)
        self.assertEqual(array.getNOriginalNodes(), 3)

        # Node removal
        array.removeNodes((0, 1))
        self.assertEqual(array.getNElements(True), 1)
        self.assertEqual(array.getNNodes(), 1)
        self.assertTrue(array.getFromId(0) is None)
        self.assertTrue(array.getFromId(1) is None)
        self.assertTrue(array.getFromId(2) is not None)

        node = array.copyNode(node)
        self.assertTrue(array.getFromId(3) is not None)
        array.removeAllDuplicateNodes()
        self.assertTrue(array.getFromId(3) is None)
        self.assertEqual(array.getNElements(True), 1)
        self.assertEqual(array.getNNodes(), 1)

        array.removeNodesOfElement(v_sim.Element.lookup("H"))
        self.assertEqual(array.getNElements(True), 0)
        self.assertEqual(array.getNNodes(), 0)

    def test_populationSignals(self):
      array = v_sim.Data.new()
      with signals.Listener(array, "notify::n-nodes") as nodes:
        with signals.Listener(array, "notify::elements") as elems:
          with signals.Listener(array, "PopulationIncrease") as popInc:
            with signals.Listener(array, "PopulationDecrease") as popDec:
              self.assertEqual(nodes.triggered(), 0)
              self.assertEqual(elems.triggered(), 0)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 0)
      
              array.allocateByNames((1, 4), ("C", "H"))
              self.assertEqual(nodes.triggered(), 0)
              self.assertEqual(elems.triggered(), 1)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 0)
              elems.reset()

              array.getNewNodeForId(0)
              self.assertEqual(nodes.triggered(), 1)
              self.assertEqual(elems.triggered(), 0)
              self.assertEqual(popInc.triggered(), 1)
              self.assertEqual(popDec.triggered(), 0)
              nodes.reset()
              popInc.reset()

              array.startAdding()
              array.getNewNodeForId(1)
              array.getNewNodeForId(1)
              array.getNewNodeForId(1)
              array.getNewNodeForId(1)
              self.assertEqual(nodes.triggered(), 0)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 0)
              array.completeAdding()
              self.assertEqual(nodes.triggered(), 1)
              self.assertEqual(popInc.triggered(), 1)
              self.assertEqual(popDec.triggered(), 0)
              nodes.reset()
              popInc.reset()

              array.removeNodes((1, 2))
              self.assertEqual(nodes.triggered(), 1)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 1)

              array.removeNodesOfElement(v_sim.Element.lookup("C"))
              self.assertEqual(nodes.triggered(), 2)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 2)

              array.freeNodes()
              self.assertEqual(nodes.triggered(), 3)
              self.assertEqual(elems.triggered(), 1)
              self.assertEqual(popInc.triggered(), 0)
              self.assertEqual(popDec.triggered(), 2) # freeNodes() does not emit PopulationDecrease

    def test_nodeMovements(self):
      array = v_sim.Data.new()
      array.allocateByNames((1, 4), ("C", "H"))
      array.getNewNodeForId(0).newValues(( 0, 0, 0))
      array.getNewNodeForId(1).newValues((-1,-1,-1))
      array.getNewNodeForId(1).newValues(( 1, 1,-1))
      array.getNewNodeForId(1).newValues(( 1,-1, 1))
      array.getNewNodeForId(1).newValues((-1, 1, 1))

      with signals.Listener(array, "position-changed") as moves:
        array.getFromId(0).setCoordinates((2.,2.,2.))
        self.assertEqual(array.getFromId(0).xyz, [2.,2.,2.])
        self.assertEqual(moves.triggered(), 0)

        array.moveNode(0, (0.,0.,0.))
        self.assertEqual(array.getFromId(0).xyz, [0.,0.,0.])
        self.assertEqual(moves.triggered(), 1)
        array.moveNodes((3,4), (4.,5.,6.,7.,8.,9.))
        self.assertEqual(array.getFromId(3).xyz, [4.,5.,6.])
        self.assertEqual(array.getFromId(4).xyz, [7.,8.,9.])
        self.assertEqual(moves.triggered(), 2)
        moves.reset()

        array.shiftNode(0, (1,2,3))
        self.assertEqual(array.getFromId(0).xyz, [1,2,3])
        self.assertEqual(moves.triggered(), 1)
        array.shiftNodes((3,4), (-3.,-2.,-1.))
        self.assertEqual(array.getFromId(3).xyz, [1.,3.,5.])
        self.assertEqual(array.getFromId(4).xyz, [4.,6.,8.])
        self.assertEqual(moves.triggered(), 2)
        moves.reset()

        array.startMoving()
        array.shiftNode(0, (-1,-2,-3))
        array.moveNodes((3, 4), (1,-1,1,-1,1,1))
        array.rotateNodes((1,2), (0,0,1), (0,0,0), 45);
        self.assertEqual(moves.triggered(), 0)
        array.completeMoving()
        self.assertEqual(moves.triggered(), 1)
        moves.reset()

        array.rotateNodes((1,2), (0,0,1), (0,0,0), 45);
        self.assertEqual(array.getFromId(1).xyz, [1.,-1.,-1.])
        self.assertEqual(array.getFromId(2).xyz, [-1.,1.,-1.])
        self.assertEqual(moves.triggered(), 1)

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