import pandas as pd
from rdkit import Chem

# ---------- SETTINGS ----------
INPUT_FILE = 'S2_PMT_vPvM_classification.xlsx'
OUTPUT_FILE = 'S3_Chemical_features_grouping.xlsx'

#----------------------------
# Load the Excel file
df = pd.read_excel(INPUT_FILE)

# Halagonated present
def num_Br(mol):
    if mol is None:
        return 0
    count = sum(1 for atom in mol.GetAtoms() if atom.GetSymbol() == 'Br')
    return count

def num_Cl(mol):
    if mol is None:
        return 0
    count = sum(1 for atom in mol.GetAtoms() if atom.GetSymbol() == 'Cl')
    return count

def num_F(mol):
    if mol is None:
        return 0
    count = sum(1 for atom in mol.GetAtoms() if atom.GetSymbol() == 'F')
    return count


# function to count the number of -CF2- in a molecule OECD guideline
def check_CF2(structure):
    if structure is None:
        return 0
    cf_2 = Chem.MolFromSmarts("[CX4](F)(F)([!F])([!F])") # OECD guideline
    check = structure.GetSubstructMatches(cf_2)
    return len(check)

# function to count  the number of -CF3 in a molecule
def check_CF3(structure):
    if structure is None:
        return 0
    cf_3 = Chem.MolFromSmarts('[CX4](F)(F)(F)([!F])') # OECD guideline
    check = structure.GetSubstructMatches(cf_3)
    return len(check)

#function to count the number of potentially forming TFA in a molecule CF3-C
def check_TFA(structure):
    if structure is None:
        return 0
    tfa = Chem.MolFromSmarts('[CX4](F)(F)(F)[CX4;!$(*-F)]')
    check = structure.GetSubstructMatches(tfa)
    return len(check)

def has_TFA_precursor(mol):
    if mol is None:
        return 0

    patterns = [
        Chem.MolFromSmarts('[CX4](F)(F)(F)C(=O)'),   # strong
        Chem.MolFromSmarts('[CX4](F)(F)(F)[CH1,CH2]'), # oxidizable
        Chem.MolFromSmarts('c[CX4](F)(F)(F)')        # aromatic CF3
    ]

    for p in patterns:
        if mol.HasSubstructMatch(p):
            return 1

    return 0



#function to count triazine in a molecule SMILE C1=CN=NN=C1
def count_triazine(structure):
    if structure is None:
        return 0

    triazines = [
        Chem.MolFromSmarts("c1ncncn1"),  # 1,3,5-triazine (aromatic)
        Chem.MolFromSmarts("n1ncncc1"),  # 1,2,4-triazine
        Chem.MolFromSmarts("n1nnccc1")   # 1,2,3-triazine
    ]

    total = 0
    for patt in triazines:
        total += len(structure.GetSubstructMatches(patt))

    return total

#function for aromatic amine SMILE C1=CC=C(C=C1)N
def count_aromatic_amine(structure):
    if structure is None:
        return 0
    aromatic_amine = Chem.MolFromSmarts("c1ccc(cc1)N")
    return len (structure.GetSubstructMatches(aromatic_amine))

#function to count triazole in a molecule SMILE C1=NNN=C1
def count_triazole(structure):
    if structure is None:
        return 0
    triazole = Chem.MolFromSmarts("n1nncc1")  # aromatic 1,2,3-triazole
    return len(structure.GetSubstructMatches(triazole))

#function to count melamine SMILE C1(=NC(=NC(=N1)N)N)N
def count_melamine(structure):
    if structure is None:
        return 0
    melamine = Chem.MolFromSmarts("n1c(nc(nc1[NH2,NH,NX3])[NH2,NH,NX3])[NH2,NH,NX3]") # Matches melamine and substituted melamines (where the amino groups may be substituted):
    return len(structure.GetSubstructMatches(melamine))

# function to count benzotriazole SMILE C1=CC2=NNN=C2C=C1
def count_benzotriazole(structure):
    if structure is None:
        return 0
    benzotriazole = Chem.MolFromSmarts("[c]1[c][c][c]2[n][n][n][c]2[c]1")  # general benzotriazole core

    return len(structure.GetSubstructMatches(benzotriazole))

def count_benzidine(mol):
    if mol is None:
        return 0

    biphenyl = Chem.MolFromSmarts('c1ccccc1-c2ccccc2')
    aniline = Chem.MolFromSmarts('cN')

    if not mol.HasSubstructMatch(biphenyl):
        return 0

    return len(mol.GetSubstructMatches(aniline))


# function to count phthalate ester
def count_phthalate_esters(structure):
    if structure is None:
        return 0
    phthalate_ester = Chem.MolFromSmarts('c1cc(C(=O)O[#6])c(C(=O)O[#6])cc1')
    matches = structure.GetSubstructMatches(phthalate_ester)
    return len(matches)

#Function to count PCB
def count_pcb(mol):
    if mol is None:
        return 0

    biphenyl = Chem.MolFromSmarts('c1ccccc1-c2ccccc2')

    for match in mol.GetSubstructMatches(biphenyl):
        ring1 = match[:6]
        ring2 = match[6:]

        def has_cl(ring):
            return any(
                any(nbr.GetSymbol() == 'Cl'
                    for nbr in mol.GetAtomWithIdx(idx).GetNeighbors())
                for idx in ring
            )

        if has_cl(ring1) and has_cl(ring2):
            return 1

    return 0


df['Structure'] = df['SMILES'].apply(
    lambda x: Chem.MolFromSmiles(str(x)) if pd.notnull(x) else None
)

df['Br_groups'] = df['Structure'].apply(lambda x: num_Br(x))
df['Cl_groups'] = df['Structure'].apply(lambda x: num_Cl(x))
df['F_groups'] = df['Structure'].apply(lambda x: num_F(x))

df['CF2_groups'] = df['Structure'].apply(lambda x: check_CF2(x))
df['CF3_groups'] = df['Structure'].apply(lambda x: check_CF3(x))
df['TFA_groups'] = df['Structure'].apply(lambda x: check_TFA(x))
df['triazine_groups'] = df['Structure'].apply(lambda x: count_triazine(x))
df['triazole_groups'] = df['Structure'].apply(lambda x: count_triazole(x))
df['aromatic_amine_groups'] =df['Structure'].apply(lambda x:count_aromatic_amine(x))
df['melamine_groups'] = df['Structure'].apply(lambda x: count_melamine(x))
df['benzotriazole_groups'] = df['Structure'].apply(lambda x: count_benzotriazole(x))

df['phthalate_groups'] = df['Structure'].apply(lambda x: count_phthalate_esters(x))
df['benzidine_groups'] = df['Structure'].apply(count_benzidine)
df['pcb_groups']= df['Structure'].apply(count_pcb)
df['TFA_precursor'] = df['Structure'].apply(has_TFA_precursor)
total = df['TFA_precursor'].sum()

# Save
df.to_excel(OUTPUT_FILE, index=False)

# List of SMARTS count columns
smart_columns = [
    'Br_groups',
    'Cl_groups',
    'F_groups',
    'CF2_groups',
    'CF3_groups',
    'TFA_groups',
    'triazine_groups',
    'triazole_groups',
    'aromatic_amine_groups',
    'melamine_groups',
    'benzotriazole_groups',
    'benzidine_groups',
    'phthalate_groups',
    'TFA_precursor',
    'pcb_groups'
]

# Count how many compounds contain at least 1 match for each SMART
summary = {col: (df[col] > 0).sum() for col in smart_columns}

# Convert to a DataFrame for easy display
df_summary = pd.DataFrame.from_dict(summary, orient='index', columns=['Num_compounds_with_match'])
print(df_summary)

df_expanded = df.copy()

df_expanded['Use'] = df_expanded['Use'].str.split(',')
df_expanded = df_expanded.explode('Use')
df_expanded['Use'] = df_expanded['Use'].str.strip()

# Count compounds with at least one match for each SMARTS, grouped by use
df_summary_by_use = (
    df_expanded.groupby("Use")[smart_columns]
      .apply(lambda x: (x > 0).sum())
)

print(df_summary_by_use)


df_summary_by_use.to_excel('pesticide_biocides/Summary_smarts_by_use.xlsx')