import unittest
import numpy as np
import os
import shutil

# import sys
# sys.path.append('../src')

from saclay_format import saclay_parser as parser

class TestParser(unittest.TestCase):

    def setUp(self):
        """Create a temporary working directory."""
        self.tmpdir = "tmp_test_parser"
        os.makedirs(self.tmpdir, exist_ok=True)

    def tearDown(self):
        """Remove temporary directory after each test."""
        shutil.rmtree(self.tmpdir)

    # ----------------------------------------------------------
    def create_dummy_rho(self, shape):
        ''' Create a field array (xright convention) '''
        nx, ny, nz = shape
        x = np.arange(nx)
        y = np.arange(ny)
        z = np.arange(nz)
        x, y, z = np.meshgrid(x, y, z, indexing="ij")

        field = (x + 2.*y + 3.*z)

        return field

    # ----------------------------------------------------------
    def create_dummy_data(self):
        nx, ny, nz = 6, 7, 8
        nw_n, nw_p = 5, 4

        data = {
            "metadata": {
                "nx": nx, "ny": ny, "nz": nz,
                "dx": 0.9, "dy": 0.8, "dz": 0.7,
                "prefix": "dummy",
                "description": "dummy data",
                "Z": 10, "N": 12,
                "frame": 1,
                "t0": 0., "dt": 0.,
            }
        }

        # -------------------- fields --------------------
        rho_n = {
            "name": "rho_n",
            "n_components": 1,
            "type": "real",
            "unit": "fm^-3",
            "suffix": "_rho_n.bin"
        }
        j = {
            "name": "j",
            "n_components": 3,
            "type": "real",
            "unit": "arb",
            "suffix": "_j.bin"
        }

        data["metadata"]["field"] = [rho_n, j]
        data["rho_n"] = self.create_dummy_rho((nx, ny, nz))
        data["j"] = np.random.rand(nx, ny, nz, 3).astype(np.float64, order="F")

        # ---------------- wavefunctions ----------------
        wf = {
            "representation": "canonical",
            "n_neutron_states": nw_n,
            "n_proton_states": nw_p,
            "suffix": "_state"
        }
        data["metadata"]["wavefunction"] = wf
     
        statesn = np.random.rand(nx, ny, nz, 2, nw_n).astype(np.complex128, order='F')
        statesn[1,0,0,0,0] = 1.j
        data["statesn"] = statesn
        data["un"] = np.random.rand(nw_n).astype(np.complex128)
        data["vn"] = np.random.rand(nw_n).astype(np.complex128)

        data["statesp"] = np.random.rand(nx, ny, nz, 2, nw_p).astype(np.complex128, order="F")
        data["up"] = np.random.rand(nw_p).astype(np.complex128)
        data["vp"] = np.random.rand(nw_p).astype(np.complex128)

        return data

    
    # ----------------------------------------------------------
    def test_xleft_to_xright(self):

        # Fortran index order with column-major storage
        a = np.zeros((2,3), order='F')
        a[0,0] = 0.
        a[1,0] = 1.
        a[0,1] = 2.
        a[1,1] = 3.
        a[0,2] = 4.
        a[1,2] = 5.

        b = parser.ixleft_to_ixright(a)
        self.assertTrue(b.flags['C_CONTIGUOUS'])
        self.assertEqual(b.shape, (3,2))

        diff = np.ravel(b, order='A') - np.array([0,1,2,3,4,5])
        self.assertLess(np.linalg.norm(diff), 1.e-15)

        # Fortran index order with row-major storage
        a = np.zeros((2,3), order='C')
        a[0,0] = 0.
        a[1,0] = 1.
        a[0,1] = 2.
        a[1,1] = 3.
        a[0,2] = 4.
        a[1,2] = 5.

        b = parser.ixleft_to_ixright(a)
        diff = np.ravel(b, order='A') - np.array([0,1,2,3,4,5])

        self.assertTrue(b.flags['C_CONTIGUOUS'])
        self.assertEqual(b.shape, (3,2))
        self.assertLess(np.linalg.norm(diff), 1.e-15)

    # ----------------------------------------------------------
    def test_xright_to_xleft(self):

        # C index order with row-major storage
        a = np.zeros((3,2), order='C')
        a[0,0] = 0.
        a[0,1] = 1.
        a[1,0] = 2.
        a[1,1] = 3.
        a[2,0] = 4.
        a[2,1] = 5.

        b = parser.ixright_to_ixleft(a)
        diff = np.ravel(b, order='A') - np.array([0,1,2,3,4,5])

        self.assertTrue(b.flags['F_CONTIGUOUS'])
        self.assertEqual(b.shape, (2,3))
        self.assertLess(np.linalg.norm(diff), 1.e-15)

    # ----------------------------------------------------------
    def test_expand_field_dims(self):
        a = np.arange(24)
        f = a.reshape((2,3,4), order='F')
        self.assertEqual(f[1,0,0], 1.)

        g = parser.expand_field_dims(f)
        self.assertEqual(g.shape, (2,3,4,1,1))
        self.assertEqual(g[1,0,0,0,0], 1.)
        self.assertEqual(g[0,1,0,0,0], 2.)
        self.assertTrue(g.flags['F_CONTIGUOUS'])

    # ----------------------------------------------------------
    def test_reduce_field_dims(self):
        a = np.arange(24)
        f = a.reshape((2,3,4,1,1), order='F')
        g = parser.reduce_last_dims(f, n_dims=2)
        self.assertEqual(g.shape, (2,3,4))
        self.assertEqual(g[1,0,0], 1.)
        self.assertTrue(g.flags['F_CONTIGUOUS'])

        a = np.arange(48)
        f = a.reshape((2,3,4,2,1), order='C')
        g = parser.reduce_last_dims(f, n_dims=2)
        self.assertEqual(g.shape, (2,3,4,2))
        self.assertEqual(g[1,0,0,0], 24.)
        self.assertTrue(g.flags['C_CONTIGUOUS'])

        a = np.arange(48)
        f = a.reshape((2,3,4,1,2), order='F')
        g = parser.reduce_last_dims(f, n_dims=2)
        self.assertEqual(g.shape, (2,3,4,1,2))
        self.assertEqual(g[1,0,0,0,0], 1.)
        self.assertTrue(g.flags['F_CONTIGUOUS'])

    # ----------------------------------------------------------
    def test_yaml_write_read(self):
        data = self.create_dummy_data()

        yamlfile = os.path.join(self.tmpdir, "dummy.yaml")
        parser.write(data, yamlfile)

        read_data = parser.read(yamlfile)

        # version
        self.assertEqual(read_data["metadata"]["version"], parser.PARSER_VERSION)

        # fields
        self.assertLess(np.linalg.norm(read_data["rho_n"] - data["rho_n"]), 1e-15)
        self.assertLess(np.linalg.norm(read_data["j"] - data["j"]), 1e-15)

        serialized_rho = np.ravel(read_data["rho_n"], order='A')
        self.assertLess(serialized_rho[1] - 1., 1e-15)
        self.assertLess(serialized_rho[2] - 2., 1e-15)

        # wavefunctions
        for key in ["statesn", "statesp", "un", "vn", "up", "vp"]:
            self.assertLess(np.linalg.norm(read_data[key] - data[key]), 1e-15)

        serialized_statesn = np.ravel(data['statesn'], order='A')
        self.assertLess(serialized_statesn[1] - 1.j, 1e-15)

    # ----------------------------------------------------------
    def test_hdf5_write_read(self):
        data = self.create_dummy_data()

        h5file = os.path.join(self.tmpdir, "dummy.h5")
        parser.write(data, h5file)

        read_data = parser.read(h5file)
        
        # version
        self.assertEqual(read_data["metadata"]["version"], parser.PARSER_VERSION)
        
        # fields
        self.assertLess(np.linalg.norm(read_data["rho_n"] - data["rho_n"]), 1e-15)
        self.assertLess(np.linalg.norm(read_data["j"] - data["j"]), 1e-15)
        
        serialized_rho = np.ravel(read_data["rho_n"], order='A')
        self.assertLess(serialized_rho[1] - 1., 1e-15)
        self.assertLess(serialized_rho[2] - 2., 1e-15)
        
        # wavefunctions
        for key in ["statesn", "statesp", "un", "vn", "up", "vp"]:
            self.assertLess(np.linalg.norm(read_data[key] - data[key]), 1e-15)

        serialized_statesn = np.ravel(data['statesn'], order='A')
        self.assertLess(serialized_statesn[1] - 1.j, 1e-15)

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