import unittest

from pdgrapher import Trainer
import torch
import os
torch.set_num_threads(5)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"


class TestTrainer(unittest.TestCase):

    def test_trainer(self):
        trainer = Trainer()

    def test_trainer_args(self):
        trainer = Trainer(
            fabric_kwargs={"accelerator": "cpu"},
            log_train=False, use_backward_data=True, supervision_multiplier=1.3
        )

        self.assertFalse(trainer.log_train)
        self.assertTrue(trainer.use_backward_data)
        self.assertEqual(trainer.supervision_multiplier, 1.3)

        self.assertEqual(trainer.fabric.device.type, "cpu")

    def test_trainer_unknown_kwargs(self):
        with self.assertWarns(Warning):
            trainer = Trainer(
                unknown_arg=1, another_unknown=True
            )


if __name__ == "__main__":
    unittest.TestLoader.sortTestMethodsUsing = None
    unittest.main()