import io
from typing import List, Dict
import os
import requests
import pandas as pd
import gzip
import shutil

from rdkit import Chem
from pycomex.functional.experiment import Experiment
from pycomex.utils import folder_path, file_namespace


# :param DATASET_NAME:
#       This string determines the name of the message pack dataset file that is then 
#       stored into the "results" folder of the experiment as the result of the 
#       processing process. The corresponding file extensions will be added 
#       automatically.
DATASET_NAME: str = 'compas_1x'
# :param METADATA:
#       A dictionary which will be used as the basis for the metadata that will be added 
#       as additional information to the file share server.
METADATA: dict = {
    'verbose': 'Cata-Condensed Polybenzenoid Hydrocarbons (COMPAS-1x)',
    'description': (
        'The COMPAS-1x dataset is part of the COMPAS Project, a computational database of polycyclic '
        'aromatic systems. This dataset contains optimized ground-state structures and molecular properties '
        'for approximately 34,000 cata-condensed polybenzenoid hydrocarbons (also known as polycyclic aromatic '
        'hydrocarbons or PAHs), computed at the GFN2-xTB level of theory. The dataset covers molecules ranging '
        'from 1 to 11 rings and includes key quantum chemical properties such as HOMO/LUMO energies, energy gaps, '
        'dipole moments, and electron affinities. This resource supports data-driven investigation of '
        'structure-property relationships in aromatic systems and facilitates machine learning applications '
        'in organic electronics and materials design.'
    ),
    'target_type': ['regression'],
    'tags': ['SMILES', 'Molecules', 'Quantum Chemistry', 'Molecular Properties', 'Polycyclic Aromatic Hydrocarbons', 'PAH', 'GFN2-xTB'],
    'sources': [
        'https://pubs.acs.org/doi/full/10.1021/acs.jcim.2c00503',
        'https://pubmed.ncbi.nlm.nih.gov/35881922/',
        'https://chemrxiv.org/engage/chemrxiv/article-details/626836b8368ab655418e3c7f',
        'https://acs.figshare.com/collections/The_COMPAS_Project_A_Computational_Database_of_Polycyclic_Aromatic_Systems_Phase_1_i_cata_i_-Condensed_Polybenzenoid_Hydrocarbons/6118490',
        'https://gitlab.com/porannegroup/compas/-/tree/main/COMPAS-1?ref_type=heads',
    ],
    'target_descriptions': {
        0: 'HOMO_eV - Energy of the highest occupied molecular orbital (HOMO) in electron volts (eV)',
        1: 'LUMO_eV - Energy of the lowest unoccupied molecular orbital (LUMO) in electron volts (eV)',
        2: 'GAP_eV - HOMO-LUMO energy gap in electron volts (eV)',
        3: 'Dipmom_Debye - Dipole moment in Debye units',
        4: 'Etot_eV - Total energy of the molecule in electron volts (eV)',
        5: 'aEA_eV - Adiabatic electron affinity in electron volts (eV)',
        6: 'aIP_eV - Adiabatic ionization potential in electron volts (eV)',
        7: 'NFOD - Fractional occupation density (indicator of polyradical character)',
        8: 'n_rings - Number of aromatic rings in the molecule',
    }
}

__TESTING__ = False

experiment = Experiment.extend(
    'create_graph_datasets.py',
    base_path=folder_path(__file__),
    namespace=file_namespace(__file__),
    glob=globals(),
)


@experiment.hook('add_graph_metadata', default=False, replace=True)
def add_graph_metadata(e: Experiment, data: dict, graph: dict) -> dict:
    """
    We add the compound id for identification and the molecular weight
    """
    pass #graph['graph_id'] = data['ID']


@experiment.hook('load_dataset', default=False, replace=True)
def load_dataset(e: Experiment) -> Dict[int, dict]:
    
    e.log('downloading raw csv file...')
    raw_url = 'https://gitlab.com/porannegroup/compas/-/raw/main/COMPAS-1/compas-1x.csv?ref_type=heads&inline=false'
    response = requests.get(raw_url, verify=False)
    csv_file = io.StringIO(response.text)
    df = pd.read_csv(csv_file)
    
    print(df.head())
    
    compressed_path = os.path.join(e.path, f'{e.DATASET_NAME}.csv.gz')
    with gzip.open(compressed_path, mode='wb') as compressed_file:
        shutil.copyfileobj(csv_file, compressed_file)

    columns = [
        'HOMO_eV',
        'LUMO_eV',
        'GAP_eV',
        'Dipmom_Debye',
        'Etot_eV',
        'aEA_eV',
        'aIP_eV',
        'NFOD',
        'n_rings',
    ]

    dataset: Dict[int, dict] = {}
    index: int = 0
    for data in df.to_dict('records'):
        data['smiles'] = data['smiles']
        
        # == MOLECULE FILTERS ==
        # We don't want to use compounds with '.' in the smiles (separate molecules)
        if '.' in data['smiles']:
            continue
        
        # We don't want to use compounds that only consist of a single atom
        mol = Chem.MolFromSmiles(data['smiles'])
        if not mol:
            continue
        
        # We also don't want to accept "molecules" that are essentially just individual atoms
        if len(mol.GetAtoms()) < 2:
            continue
        
        # == TARGETS ==
        # In this dataset we only have one target
        data['targets'] = [data[key] for key in columns]
        dataset[index] = data
        
        index += 1

    return dataset

experiment.run_if_main()
