# -*- coding: utf-8 -*-
#
# Copyright (c) 2023, the cclib development team
#
# This file is part of cclib (http://cclib.github.io) and is distributed under
# the terms of the BSD 3-Clause License.

"""Unit tests for Molden writer."""

import os
import unittest

import cclib
from cclib.io.filewriter import MissingAttributeError
from cclib.io.moldenwriter import MoldenReformatter
from cclib.io.moldenwriter import round_molden
import numpy
import pytest

__filedir__ = os.path.dirname(__file__)
__filepath__ = os.path.realpath(__filedir__)
__datadir__ = os.path.join(__filepath__, "..", "..")
__testdir__ = __filedir__

class MOLDENTest(unittest.TestCase):

    def test_missing_attribute_error(self):
        """Check if MissingAttributeError is raised as expected."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/dvb_un_sp.out")
        required_attrs = ['atomcoords', 'atomnos', 'natom']
        for attr in required_attrs:
            data = cclib.io.ccread(fpath)
            delattr(data, attr)

            # Molden files cannot be written if required attrs are missing.
            with pytest.raises(MissingAttributeError):
                cclib.io.moldenwriter.MOLDEN(data)

    def test_atoms_section_size(self):
        """Check if size of Atoms section is equal to expected."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/dvb_un_sp.out")
        data = cclib.io.ccread(fpath)
        writer = cclib.io.moldenwriter.MOLDEN(data)
        # Check size of Atoms section.
        assert len(writer._coords_from_ccdata(-1)) == data.natom

    def test_atoms_section_size_with_ghost(self):
        """Check if size of Atoms section is equal to expected with a ghost atom present."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/dvb_un_sp_ghost.out")
        data = cclib.io.ccread(fpath)
        writer = cclib.io.moldenwriter.MOLDEN(data)
        writer.ghost = "GH"
        # Check for the ghost atom, and make sure it has the right label.
        assert "GH" in writer._coords_from_ccdata(-1)[0]
        # Check size of Atoms section.
        assert len(writer._coords_from_ccdata(-1)) == data.natom

    def test_gto_section_size(self):
        """Check if size of GTO section is equal to expected."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/dvb_un_sp.out")
        data = cclib.io.ccread(fpath)
        writer = cclib.io.moldenwriter.MOLDEN(data)
        # Check size of GTO section.
        size_gto_ccdata = 0
        for atom in data.gbasis:
            size_gto_ccdata += 1
            for prims in atom:
                size_gto_ccdata += len(prims[1]) + 1
        # Filter blank lines.
        size_gto_writer = len(list(filter(None, writer._gto_from_ccdata())))
        assert size_gto_writer == size_gto_ccdata

    def test_mo_section_size(self):
        """Check if size of MO section is equal to expected."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/dvb_un_sp.out")
        data = cclib.io.ccread(fpath)
        writer = cclib.io.moldenwriter.MOLDEN(data)
        # Check size of MO section.
        size_mo_ccdata = 0
        extra = 4 if hasattr(data, 'mosyms') else 3
        for i in range(data.mult):
            size_mo_ccdata += len(data.moenergies[i]) *\
                                (len(data.mocoeffs[i][0]) + extra)
        # Filter blank lines.
        mosyms, moenergies, mooccs, mocoeffs = writer._syms_energies_occs_coeffs_from_ccdata_for_moldenwriter()
        size_mo_writer = len(list(filter(None, writer._mo_from_ccdata(mosyms, moenergies, mooccs, mocoeffs))))
        assert size_mo_writer == size_mo_ccdata

    def test_no_section_size(self):
        """Check if size of NO section is equal to expected."""
        fpath = os.path.join(__datadir__,
                             "data/GAMESS/basicGAMESS-US2018/water_cis_dets.out")
        data = cclib.io.ccread(fpath)
        writer = cclib.io.moldenwriter.MOLDEN(data)
        # Check size of NO section.
        size_no_ccdata = 0
        extra = 4
        size_no_ccdata += len(data.nooccnos) *\
                                (len(data.nocoeffs[0]) + extra)
        # Filter blank lines.
        nosyms, noenergies, nooccs, nocoeffs = writer._syms_energies_occs_coeffs_from_ccdata_for_moldenwriter()
        size_no_writer = len(list(filter(None, writer._mo_from_ccdata(nosyms, noenergies, nooccs, nocoeffs))))
        assert size_no_writer == size_no_ccdata

    def test_round_molden(self):
        """Check if Molden Style number rounding works as expected."""
        # If the 6th digit after dot is greater than 5, but is not 7,
        # round the number upto 6th place.
        # Else truncate at 6th digit after dot.
        assert round_molden(1) == 1
        assert round_molden(-1) == -1
        assert round_molden(0.999995789) == 0.999995
        assert round_molden(-0.999995789) == -0.999995
        assert round_molden(0.999996789) == 0.999997
        assert round_molden(-0.999997789) == -0.999997
        assert round_molden(0.999997789) == 0.999997
        assert round_molden(-0.999998789) == -0.999999
        assert round_molden(-0.999999999) == -1.0

    def test_molden_cclib_diff(self):
        """Check if file written by cclib matched file written by Molden."""
        filenames = ['dvb_un_sp', 'C_bigbasis', 'water_mp2', 'dvb_ir']
        for fn in filenames:
            fpath = os.path.join(__datadir__,
                                 f"data/GAMESS/basicGAMESS-US2018/{fn}.out")
            data = cclib.io.ccread(fpath)
            cclib_out = cclib.io.moldenwriter.MOLDEN(data).generate_repr()
            # Reformat cclib's output to remove extra spaces.
            cclib_out_formatted = MoldenReformatter(cclib_out).reformat()
            fpath = os.path.join(__testdir__, f"data/molden5.7_{fn}.molden")
            with open(fpath) as handle:
                molden_out = handle.read()
            # Reformat Molden's output to remove extra spaces,
            # and fix number formatting.
            molden_out_formatted = MoldenReformatter(molden_out).reformat()
            # Assert if reformatted files from both writers are same.
            assert molden_out_formatted == cclib_out_formatted


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