#!/usr/bin/env python3
"""Test drawing functions from qspoc.utils"""

import unittest
from pathlib import Path
import pandas as pd

from qspoc.utils import smiles2png, smileslist2png


class TestDrawPictures(unittest.TestCase):
    """Test drawing functions for molecules"""

    def setUp(self):
        """Set up test environment"""
        self.test_output_dir = Path(__file__).parent / "test_file/pics"
        self.test_output_dir.mkdir(exist_ok=True)

        # Test SMILES strings
        self.test_smiles = "[H][C@@]12C[C@H](C)[C@](O)(C(=O)CO)[C@@]1(C)C[C@H](O)[C@@]1(F)[C@@]2([H])C[C@H](F)C2=CC(=O)C=C[C@]12C"  # Triflate anion
        self.test_smiles_list = ["O=S(=O)([O-])C(F)(F)F", "O=C([O-])c1ccccc1", "CC(=O)[O-]"]  # Triflate  # Benzoate  # Acetate

    def test_single_smiles_to_png(self):
        """Test generating PNG from a single SMILES string"""
        smiles2png(self.test_smiles, file_path=self.test_output_dir / "test_img.png")
        smiles2png(self.test_smiles, file_path=self.test_output_dir / "test_img_highlight.png", highlight_atoms=[0, 1])
        smiles2png(self.test_smiles, file_path=self.test_output_dir / "test_img_index.png", show_index=True)
        smiles2png(self.test_smiles, file_path=self.test_output_dir / "test_img_hydrogens.png", addHs=True)
        smiles2png(
            self.test_smiles,
            file_path=self.test_output_dir / "test_img_all.png",
            addHs=True,
            show_index=True,
            highlight_atoms=[[13, 21], [4, 5, 6]],
        )

    def test_smiles_list_to_png_from_csv(self):
        """Test generating PNG images from CSV input using smileslist2png"""
        csv_path = self.test_output_dir / "../input_test_simple.csv"
        self.assertTrue(csv_path.exists(), "input_test.csv should exist in test_file directory")
        df = pd.read_csv(csv_path)
        # Extract SMILES from different columns for testing
        smiles_list, name_list = df["smiles_list"].tolist(), df["name_list"].tolist()
        name_list = [f"{i}.png" for i in name_list]
        highlight_list = df["atom_lists"].tolist()

        # Test output directory for CSV results
        csv_output_dir = self.test_output_dir / "csv_output"

        # Generate images using smileslist2png
        smileslist2png(
            SMILES_list=smiles_list, savepath=str(csv_output_dir), name_list=name_list, highlight_atoms_list=highlight_list, show_index=True
        )

        # Verify all images were created
        for tag in name_list:
            img_path = csv_output_dir / tag
            self.assertTrue(img_path.exists(), f"Image {tag} should be created")
            self.assertGreater(img_path.stat().st_size, 0, f"Image {tag} should have content")

    def test_smiles_list_to_png_direct(self):
        """Test generating PNG images from direct SMILES list using smileslist2png"""
        name_list = [f"direct_test_{i}.png" for i in range(len(self.test_smiles_list))]
        direct_output_dir = self.test_output_dir / "direct_output"

        # Generate images from direct SMILES list
        smileslist2png(SMILES_list=self.test_smiles_list, savepath=str(direct_output_dir), name_list=name_list, show_index=False)

        # Verify all images were created
        for name in name_list:
            img_path = direct_output_dir / name
            self.assertTrue(img_path.exists(), f"Direct image {name} should be created")
            self.assertGreater(img_path.stat().st_size, 0, f"Direct image {name} should have content")

    def tearDown(self):
        """Clean up after tests if needed"""
        # Note: We keep the generated images for inspection
        # In a production environment, you might want to clean them up
        pass


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