from crystalpol.shared.config import Config
from crystalpol.shared.system.atom import Atom
from crystalpol.shared.system.crystal import Crystal
from crystalpol.shared.system.molecule import Molecule
from crystalpol.shared.utils.log import Log

import logging
from io import StringIO
from unittest import TestCase, mock


class TestLog(TestCase):

    def setUp(self):

        self.log_stream = StringIO()
        logging.basicConfig(
            stream=self.log_stream, format="%(message)s", level=logging.INFO
        )
        self.logger = logging.getLogger()

    def tearDown(self):
        logging.getLogger().removeHandler(logging.getLogger().handlers[0])
        del self.logger
        del self.log_stream

    @mock.patch("crystalpol.shared.utils.log.sys")
    @mock.patch("crystalpol.shared.utils.log.weekday_date_time")
    def test_make_header(self, weekday_date_time_mock, sys_mock):
        weekday_date_time_mock.return_value = "Test Day"
        sys_mock.version = "Test Version"

        config = Config(mem=1, level="b3lyp/aug-cc-pVDZ", n_atoms=10)
        Log.make_header("test", config.to_dict(), logger=self.logger)

        expected_log_stream = [
            "##########################################################################################\n",
            "##############             Welcome to CRYSTALPOL version test             ##############\n",
            "##########################################################################################\n",
            "\n",
            "Your python version is Test Version\n",
            "\n",
            "Program started on Test Day\n",
            "\n",
            "------------------------------------------------------------------------------------------\n",
            "                   CRYSTALPOL variables being used in this run:                           \n",
            "------------------------------------------------------------------------------------------\n",
            "\n",
            "\tmem = 1\n",
            "\tlevel = b3lyp/aug-cc-pVDZ\n",
            "\tn_atoms = 10\n",
            "\tn_procs = 1\n",
            "\tpop = chelpg\n",
            "\tcomment = crystalpol\n",
            "\tcharge_tolerance = 0.02\n",
            "\tsimulation_dir = simfiles\n",
            "\tmult = [0, 1]\n",
            "------------------------------------------------------------------------------------------\n",
            "                                      RUN Results:                                       \n",
            "------------------------------------------------------------------------------------------\n",
            "\n",
        ]

        self.log_stream.seek(0)

        self.assertEqual(self.log_stream.readlines(), expected_log_stream)

    def test_make_run(self):

        Log.make_run(
            1, 0.00001, 0.000001, [], self.create_crystal(), logger=self.logger
        )

        expected_log_stream = [
            "cycle: 1\n",
            "\n",
            "Max charge diff: 0.00001\n",
            "Min charge diff: 0.00000\n",
            "Charge Diff: []\n",
            "\n",
            "------------------------------------------------------------------------------------------\n",
            "             S             rx             ry             rz             chg               \n",
            "------------------------------------------------------------------------------------------\n",
            "             H          0.000000       0.000000       0.000000       0.000000             \n",
            "\n",
            "------------------------------------------------------------------------------------------\n",
            "\n",
        ]

        self.log_stream.seek(0)

        self.assertEqual(self.log_stream.readlines(), expected_log_stream)

    @staticmethod
    def create_crystal():
        crystal_structure = [["H "]]

        crystal = Crystal(crystal_structure)

        molecule = Molecule("TESTE")
        molecule.add_atom(
            Atom(na=1, rx=0.000000, ry=0.000000, rz=0.000000, chg=0.000000)
        )

        crystal.add_cell([molecule])
        crystal.add_cell([molecule])

        return crystal