import chembl_downloader
import re

def _infer_nct_year(nct_id):
    """
    Infer the approximate registration year from a ClinicalTrials.gov NCT identifier.
    NCT IDs follow the format NCT########, where the numeric portion generally increases
    over time. This function uses approximate year ranges based on observed NCT ID
    allocation patterns to estimate when a trial was registered.

    NCT IDs are sequential and follow approximate ranges:
    - NCT00000000-NCT00999999: ~1999-2004
    - NCT01000000-NCT01999999: ~2005-2011
    - NCT02000000-NCT02999999: ~2012-2015
    - NCT03000000-NCT03999999: ~2016-2018
    - NCT04000000-NCT04999999: ~2019-2021
    - NCT05000000-NCT05999999: ~2022-2023
    - NCT06000000+: ~2024+

    :param nct_id: A ClinicalTrials.gov identifier (e.g., "NCT00000001")
    :type nct_id: str
    :return: Estimated year of trial registration, or None if the NCT ID is invalid
    :rtype: int or None
    :example:
        >>> _infer_nct_year("NCT00500000")
        2002
        >>> _infer_nct_year("NCT03000000")
        2016
        >>> _infer_nct_year("invalid")
        None
    .. note::
        This function provides an approximation based on historical NCT ID allocation
        patterns and may not be accurate for all trials. The actual registration date
        should be obtained from the official ClinicalTrials.gov database when precision
        is required.
    """
    if not isinstance(nct_id, str) or not nct_id.startswith('NCT'):
        return None
    
    # Extract the numeric part
    match = re.search(r'NCT(\d{8})', nct_id)
    if not match:
        return None
    
    num = int(match.group(1))
    
    # Approximate year ranges based on NCT ID ranges
    if num < 1000000:
        return 1999 + (num // 200000)  # ~1999-2004
    elif num < 2000000:
        return 2005 + ((num - 1000000) // 140000)  # ~2005-2011
    elif num < 3000000:
        return 2012 + ((num - 2000000) // 250000)  # ~2012-2015
    elif num < 4000000:
        return 2016 + ((num - 3000000) // 330000)  # ~2016-2018
    elif num < 5000000:
        return 2019 + ((num - 4000000) // 350000)  # ~2019-2021
    elif num < 6000000:
        return 2022 + ((num - 5000000) // 500000)  # ~2022-2023
    else:
        return 2024 + ((num - 6000000) // 500000)  # ~2024+

def molecules(version: str = '36'):
    """
    Query ChEMBL database for all molecules with clinical trial data and drug indications.
    This function retrieves comprehensive information about drug molecules from the ChEMBL database,
    including their indications, clinical trial references, target information, and mechanism of action.
    The results are filtered to include only molecules with ClinicalTrials.gov references targeting
    human proteins.
    :param version: ChEMBL database version to query, defaults to '36'
    :type version: str, optional
    :return: DataFrame containing molecule information with the following key columns:
    - chembl_id: ChEMBL identifier for the molecule
    - pref_name: Preferred name of the molecule
    - mesh_heading: MeSH term for the indication
    - mesh_id: MeSH identifier
    - phase: Clinical trial phase (20 for approved drugs, 0-3 for others)
    - reference_type: Type of reference (filtered to 'ClinicalTrials')
    - clinical_trial_id: ClinicalTrials.gov identifier(s), exploded if multiple
    - target_chembl_id: ChEMBL identifier for the target
    - target_organism: Target organism (filtered to 'Homo sapiens')
    - target_type: Type of target
    - target_uniprot_id: UniProt accession for the target
    - target_gene_name: Gene symbol for the target
    - mechanism_of_action: Description of the mechanism of action
    - action_type: Type of action on the target
    - molregno: Internal molecule registry number
    - clinical_score: Numeric score based on phase (20 for approved)
    - approved: Boolean indicating if drug is approved (phase 4)
    - trial_year: Inferred year from clinical trial ID (nullable integer)
    :rtype: pandas.DataFrame
    .. note::
       Clinical trial IDs containing multiple comma-separated values are exploded into separate rows.
    .. note::
       Phase 4 (post-marketing) trials are coded as 20 to distinguish approved drugs.
    """
    sql_all_drugs = """
    SELECT 
        MD.chembl_id,
        MD.pref_name,
        DI.mesh_heading,
        DI.mesh_id,
        DI.max_phase_for_ind AS phase,
        IREF.ref_type AS reference_type,
        IREF.ref_id AS clinical_trial_id,
        TD.chembl_id AS target_chembl_id,
        TD.organism AS target_organism,
        TD.target_type,
        CS_TARGET.accession AS target_uniprot_id,
        COMP_SYN.component_synonym AS target_gene_name,
        DM.mechanism_of_action,
        DM.action_type,
        MD.molregno
    FROM MOLECULE_DICTIONARY MD
    INNER JOIN DRUG_INDICATION DI ON MD.molregno = DI.molregno
    INNER JOIN INDICATION_REFS IREF ON DI.drugind_id = IREF.drugind_id
    INNER JOIN DRUG_MECHANISM DM ON MD.molregno = DM.molregno
    INNER JOIN TARGET_DICTIONARY TD ON DM.tid = TD.tid
    LEFT JOIN TARGET_COMPONENTS TC ON TD.tid = TC.tid
    LEFT JOIN COMPONENT_SEQUENCES CS_TARGET ON TC.component_id = CS_TARGET.component_id
    LEFT JOIN COMPONENT_SYNONYMS COMP_SYN ON CS_TARGET.component_id = COMP_SYN.component_id AND COMP_SYN.syn_type = 'GENE_SYMBOL'
    WHERE IREF.ref_id IS NOT NULL 
        AND IREF.ref_type = 'ClinicalTrials'
        AND TD.organism = 'Homo sapiens'
    """
    df = chembl_downloader.query(sql_all_drugs, version=version)

    # Explode clinical_trial_id column if it contains multiple comma-separated IDs
    if 'clinical_trial_id' in df.columns and df['clinical_trial_id'].notna().any():
        df['clinical_trial_id'] = df['clinical_trial_id'].apply(
            lambda x: x.split(',') if isinstance(x, str) and ',' in x else [x] if x else []
        )
        df = df.explode('clinical_trial_id')
        df['clinical_trial_id'] = df['clinical_trial_id'].apply(lambda x: x.strip() if isinstance(x, str) else x)
        
        # Add inferred year column as integer
        df['trial_year'] = df['clinical_trial_id'].apply(_infer_nct_year)
        df['trial_year'] = df['trial_year'].astype('Int64')  # Use nullable integer type

    return df

if __name__ == '__main__':    
    # Query for all drugs
    print("\nQuerying all molecules with clinical trial data from ChEMBL...")
    df = molecules()    
    df.drop_duplicates().to_csv("molecules.csv", index=False)