#!/usr/bin/env python3

import sys, os, math
import unittest

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

class TestDumper(unittest.TestCase):
    def _data(self, filename):
        # Create a bilayer graphene cell
        a = 1.42 # CC bond length
        nA = 1 # AC direction
        A = nA * 3. * a
        nB = 1 # ZZ direction
        B = nB * math.sqrt(3.) * a

        ab = 3.41 # CC interplane distance

        C = 12.
        box = v_sim.Box.new((B, 0., C, 0., 0., A),
                            v_sim.BoxBoundaries.SURFACE_ZX)
        box.setUnit(v_sim.Units.ANGSTROEM)

        data = v_sim.DataAtomic.new(filename, None)
        data.setBox(box)

        elements = (v_sim.Element.new("C"), )
        data.allocate(elements, (1, ))

        zz = math.sqrt(3.) / 2.
        c0 = (
            (0.,     0., 0.0 * a),
            (zz * a, 0., 0.5 * a),
            (0.,     0., 2.0 * a),
            (zz * a, 0., 1.5 * a))
        c1 = (
            (0.,     ab, 1.0 * a),
            (zz * a, ab, 0.5 * a),
            (0.,     ab, 2.0 * a),
            (zz * a, ab, 2.5 * a))
        
        for iz in range(nA):
            for ix in range(nB):
                for u, v, w in c0:
                    data.addNodeFromElementName("C", (u + ix * 2*zz * a, v, w + iz * 3 * a), False)
                for u, v, w in c1:
                    data.addNodeFromElementName("C", (u + ix * 2*zz * a, v, w + iz * 3 * a), False)

        data.setTightBox()

        return data

    def test_ascii(self):
        data = self._data("bilayer.ascii")
        self.assertEqual(data.save(), True)
        loaded = v_sim.DataAtomic.new("bilayer.ascii", None)
        self.assertEqual(loaded.load(0, None), True)
        self.assertEqual(loaded.getNNodes(), 8)
        self.assertEqual(loaded.getNElements(True), 1)

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