import os

import os.path as osp

import re

import pickle

import numpy as np

import pandas as pd

from tqdm import tqdm

import torch

from torch_geometric.data import Data, InMemoryDataset

from rdkit import Chem, RDLogger

from src.data.utils import label2onehot

RDLogger.DisableLog('rdApp.*') 

class DruggenDataset(InMemoryDataset):

    def __init__(self, root, dataset_file, raw_files, max_atom, features, 

                 atom_encoder, atom_decoder, bond_encoder, bond_decoder,

                 transform=None, pre_transform=None, pre_filter=None):

        """

        Initialize the DruggenDataset with pre-loaded encoder/decoder dictionaries.

        Parameters:

            root (str): Root directory.

            dataset_file (str): Name of the processed dataset file.

            raw_files (str): Path to the raw SMILES file.

            max_atom (int): Maximum number of atoms allowed in a molecule.

            features (bool): Whether to include additional node features.

            atom_encoder (dict): Pre-loaded atom encoder dictionary.

            atom_decoder (dict): Pre-loaded atom decoder dictionary.

            bond_encoder (dict): Pre-loaded bond encoder dictionary.

            bond_decoder (dict): Pre-loaded bond decoder dictionary.

            transform, pre_transform, pre_filter: See PyG InMemoryDataset.

        """

        self.dataset_name = dataset_file.split(".")[0]

        self.dataset_file = dataset_file

        self.raw_files = raw_files

        self.max_atom = max_atom

        self.features = features

        # Use the provided encoder/decoder mappings.

        self.atom_encoder_m = atom_encoder

        self.atom_decoder_m = atom_decoder

        self.bond_encoder_m = bond_encoder

        self.bond_decoder_m = bond_decoder

        self.atom_num_types = len(atom_encoder)

        self.bond_num_types = len(bond_encoder)

        super().__init__(root, transform, pre_transform, pre_filter)

        path = osp.join(self.processed_dir, dataset_file)

        self.data, self.slices = torch.load(path)

        self.root = root

    @property

    def processed_dir(self):

        """

        Returns the directory where processed dataset files are stored.

        """

        return self.root

    @property

    def raw_file_names(self):

        """

        Returns the raw SMILES file name.

        """

        return self.raw_files

    @property

    def processed_file_names(self):

        """

        Returns the name of the processed dataset file.

        """

        return self.dataset_file

    def _filter_smiles(self, smiles_list):

        """

        Filters the input list of SMILES strings to keep only valid molecules that:

         - Can be successfully parsed,

         - Have a number of atoms less than or equal to the maximum allowed (max_atom),

         - Contain only atoms present in the atom_encoder,

         - Contain only bonds present in the bond_encoder.

        Parameters:

            smiles_list (list): List of SMILES strings.

        Returns:

            max_length (int): Maximum number of atoms found in the filtered molecules.

            filtered_smiles (list): List of valid SMILES strings.

        """

        max_length = 0

        filtered_smiles = []

        for smiles in tqdm(smiles_list, desc="Filtering SMILES"):

            mol = Chem.MolFromSmiles(smiles)

            if mol is None:

                continue

            # Check molecule size

            molecule_size = mol.GetNumAtoms()

            if molecule_size > self.max_atom:

                continue

            # Filter out molecules with atoms not in the atom_encoder

            if not all(atom.GetAtomicNum() in self.atom_encoder_m for atom in mol.GetAtoms()):

                continue

            # Filter out molecules with bonds not in the bond_encoder

            if not all(bond.GetBondType() in self.bond_encoder_m for bond in mol.GetBonds()):

                continue

            filtered_smiles.append(smiles)

            max_length = max(max_length, molecule_size)

        return max_length, filtered_smiles

    def _genA(self, mol, connected=True, max_length=None):

        """

        Generates the adjacency matrix for a molecule based on its bond structure.

        Parameters:

            mol (rdkit.Chem.Mol): The molecule.

            connected (bool): If True, ensures all atoms are connected.

            max_length (int, optional): The size of the matrix; if None, uses number of atoms in mol.

        Returns:

            np.array: Adjacency matrix with bond types as entries, or None if disconnected.

        """

        max_length = max_length if max_length is not None else mol.GetNumAtoms()

        A = np.zeros((max_length, max_length))

        begin = [b.GetBeginAtomIdx() for b in mol.GetBonds()]

        end = [b.GetEndAtomIdx() for b in mol.GetBonds()]

        bond_type = [self.bond_encoder_m[b.GetBondType()] for b in mol.GetBonds()]

        A[begin, end] = bond_type

        A[end, begin] = bond_type

        degree = np.sum(A[:mol.GetNumAtoms(), :mol.GetNumAtoms()], axis=-1)

        return A if connected and (degree > 0).all() else None

    def _genX(self, mol, max_length=None):

        """

        Generates the feature vector for each atom in a molecule by encoding their atomic numbers.

        Parameters:

            mol (rdkit.Chem.Mol): The molecule.

            max_length (int, optional): Length of the feature vector; if None, uses number of atoms in mol.

        Returns:

            np.array: Array of atom feature indices, padded with zeros if necessary, or None on error.

        """

        max_length = max_length if max_length is not None else mol.GetNumAtoms()

        try:

            return np.array([self.atom_encoder_m[atom.GetAtomicNum()] for atom in mol.GetAtoms()] +

                            [0] * (max_length - mol.GetNumAtoms()))

        except KeyError as e:

            print(f"Skipping molecule with unsupported atom: {e}")

            print(f"Skipped SMILES: {Chem.MolToSmiles(mol)}")

            return None

    def _genF(self, mol, max_length=None):

        """

        Generates additional node features for a molecule using various atomic properties.

        Parameters:

            mol (rdkit.Chem.Mol): The molecule.

            max_length (int, optional): Number of rows in the features matrix; if None, uses number of atoms.

        Returns:

            np.array: Array of additional features for each atom, padded with zeros if necessary.

        """

        max_length = max_length if max_length is not None else mol.GetNumAtoms()

        features = np.array([[*[a.GetDegree() == i for i in range(5)],

                               *[a.GetExplicitValence() == i for i in range(9)],

                               *[int(a.GetHybridization()) == i for i in range(1, 7)],

                               *[a.GetImplicitValence() == i for i in range(9)],

                               a.GetIsAromatic(),

                               a.GetNoImplicit(),

                               *[a.GetNumExplicitHs() == i for i in range(5)],

                               *[a.GetNumImplicitHs() == i for i in range(5)],

                               *[a.GetNumRadicalElectrons() == i for i in range(5)],

                               a.IsInRing(),

                               *[a.IsInRingSize(i) for i in range(2, 9)]]

                              for a in mol.GetAtoms()], dtype=np.int32)

        return np.vstack((features, np.zeros((max_length - features.shape[0], features.shape[1]))))

    def decoder_load(self, dictionary_name, file):

        """

        Returns the pre-loaded decoder dictionary based on the dictionary name.

        Parameters:

            dictionary_name (str): Name of the dictionary ("atom" or "bond").

            file: Placeholder parameter for compatibility.

        Returns:

            dict: The corresponding decoder dictionary.

        """

        if dictionary_name == "atom":

            return self.atom_decoder_m

        elif dictionary_name == "bond":

            return self.bond_decoder_m

        else:

            raise ValueError("Unknown dictionary name.")

    def matrices2mol(self, node_labels, edge_labels, strict=True, file_name=None):

        """

        Converts graph representations (node labels and edge labels) back to an RDKit molecule.

        Parameters:

            node_labels (iterable): Encoded atom labels.

            edge_labels (np.array): Adjacency matrix with encoded bond types.

            strict (bool): If True, sanitizes the molecule and returns None on failure.

            file_name: Placeholder parameter for compatibility.

        Returns:

            rdkit.Chem.Mol: The resulting molecule, or None if sanitization fails.

        """

        mol = Chem.RWMol()

        for node_label in node_labels:

            mol.AddAtom(Chem.Atom(self.atom_decoder_m[node_label]))

        for start, end in zip(*np.nonzero(edge_labels)):

            if start > end:

                mol.AddBond(int(start), int(end), self.bond_decoder_m[edge_labels[start, end]])

        if strict:

            try:

                Chem.SanitizeMol(mol)

            except Exception:

                mol = None

        return mol

    def check_valency(self, mol):

        """

        Checks that no atom in the molecule has exceeded its allowed valency.

        Parameters:

            mol (rdkit.Chem.Mol): The molecule.

        Returns:

            tuple: (True, None) if valid; (False, atomid_valence) if there is a valency issue.

        """

        try:

            Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_PROPERTIES)

            return True, None

        except ValueError as e:

            e = str(e)

            p = e.find('#')

            e_sub = e[p:]

            atomid_valence = list(map(int, re.findall(r'\d+', e_sub)))

            return False, atomid_valence

    def correct_mol(self, mol):

        """

        Corrects a molecule by removing bonds until all atoms satisfy their valency limits.

        Parameters:

            mol (rdkit.Chem.Mol): The molecule.

        Returns:

            rdkit.Chem.Mol: The corrected molecule.

        """

        while True:

            flag, atomid_valence = self.check_valency(mol)

            if flag:

                break

            else:

                # Expecting two numbers: atom index and its valence.

                assert len(atomid_valence) == 2

                idx = atomid_valence[0]

                queue = []

                for b in mol.GetAtomWithIdx(idx).GetBonds():

                    queue.append((b.GetIdx(), int(b.GetBondType()), b.GetBeginAtomIdx(), b.GetEndAtomIdx()))

                queue.sort(key=lambda tup: tup[1], reverse=True)

                if queue:

                    start = queue[0][2]

                    end = queue[0][3]

                    mol.RemoveBond(start, end)

        return mol

    def process(self, size=None):

        """

        Processes the raw SMILES file by filtering and converting each valid SMILES into a PyTorch Geometric Data object.

        The resulting dataset is saved to disk.

        Parameters:

            size (optional): Placeholder parameter for compatibility.

        Side Effects:

            Saves the processed dataset as a file in the processed directory.

        """

        # Read raw SMILES from file (assuming CSV with no header)

        smiles_list = pd.read_csv(self.raw_files, header=None)[0].tolist()

        max_length, filtered_smiles = self._filter_smiles(smiles_list)

        data_list = []

        self.m_dim = len(self.atom_decoder_m)

        for smiles in tqdm(filtered_smiles, desc='Processing dataset', total=len(filtered_smiles)):

            mol = Chem.MolFromSmiles(smiles)

            A = self._genA(mol, connected=True, max_length=max_length)

            if A is not None:

                x_array = self._genX(mol, max_length=max_length)

                if x_array is None:

                    continue

                x = torch.from_numpy(x_array).to(torch.long).view(1, -1)

                x = label2onehot(x, self.m_dim).squeeze()

                if self.features:

                    f = torch.from_numpy(self._genF(mol, max_length=max_length)).to(torch.long).view(x.shape[0], -1)

                    x = torch.concat((x, f), dim=-1)

                adjacency = torch.from_numpy(A)

                edge_index = adjacency.nonzero(as_tuple=False).t().contiguous()

                edge_attr = adjacency[edge_index[0], edge_index[1]].to(torch.long)

                data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles)

                if self.pre_filter is not None and not self.pre_filter(data):

                    continue

                if self.pre_transform is not None:

                    data = self.pre_transform(data)

                data_list.append(data)

        torch.save(self.collate(data_list), osp.join(self.processed_dir, self.dataset_file))