import pandas as pd

def run_admet_prediction(df, admet_model):
    smiles_list = df['canonical_smiles'].tolist()

    predictions = []
    for smiles in smiles_list:
        preds = admet_model.predict(smiles)
        predictions.append(preds)

    df_admet = pd.DataFrame(predictions)

    # Align with base fields
    df_base = df[['molecule_chembl_id', 'canonical_smiles']].reset_index(drop=True)
    df_admet = pd.concat([df_base, df_admet], axis=1)

    return df_admet

def filter_admet_candidates(df, strict=True):
    admet_groups = {
        'Absorption': {
            'HIA_Hou': lambda x: x >= 0.7,
            'PAMPA_NCATS': lambda x: x >= 0.7,
            'Caco2_Wang': lambda x: x >= 70,
            'Solubility_AqSolDB': lambda x: x >= -4,
            'PPBR_AZ': lambda x: x < 90,
            'Pgp_Broccatelli': lambda x: x <= 0.5,
            'tpsa': lambda x: x <= 140,
            'Bioavailability_Ma': lambda x: x >= 0.7
        },
        'Distribution': {
            'VDss_Lombardo': lambda x: x >= 0.7,
            'Lipophilicity_AstraZeneca': lambda x: 1 <= x <= 3,
            'BBB_Martins': lambda x: x >= 0.5
        },
        'Metabolism': {
            'CYP1A2_Veith': lambda x: x <= 0.5,
            'CYP2C19_Veith': lambda x: x <= 0.5,
            'CYP2C9_Substrate_CarbonMangels': lambda x: x <= 0.5,
            'CYP2C9_Veith': lambda x: x <= 0.5,
            'CYP2D6_Substrate_CarbonMangels': lambda x: x <= 0.5,
            'CYP2D6_Veith': lambda x: x <= 0.5,
            'CYP3A4_Substrate_CarbonMangels': lambda x: x <= 0.5,
            'CYP3A4_Veith': lambda x: x <= 0.5,
            'Clearance_Hepatocyte_AZ': lambda x: x <= 20,
            'Clearance_Microsome_AZ': lambda x: x <= 20,
            'Half_Life_Obach': lambda x: x >= 240
        },
        'Excretion': {
            'HydrationFreeEnergy_FreeSolv': lambda x: x < 0
        },
        'Toxicity': {
            'AMES': lambda x: x <= 0.5,
            'Carcinogens_Lagunin': lambda x: x <= 0.5,
            'ClinTox': lambda x: x <= 0.5,
            'DILI': lambda x: x <= 0.5,
            'NR-AR-LBD': lambda x: x <= 0.5,
            'NR-AR': lambda x: x <= 0.5,
            'NR-AhR': lambda x: x <= 0.5,
            'NR-Aromatase': lambda x: x <= 0.5,
            'NR-ER-LBD': lambda x: x <= 0.5,
            'NR-ER': lambda x: x <= 0.5,
            'NR-PPAR-gamma': lambda x: x <= 0.5,
            'SR-ARE': lambda x: x <= 0.5,
            'SR-ATAD5': lambda x: x <= 0.5,
            'SR-HSE': lambda x: x <= 0.5,
            'SR-MMP': lambda x: x <= 0.5,
            'SR-p53': lambda x: x <= 0.5,
            'Skin_Reaction': lambda x: x <= 0.5,
            'hERG': lambda x: x <= 0.5,
            'LD50_Zhu': lambda x: x >= 2
        }
    }

    # Coerce to numeric
    for group in admet_groups.values():
        for col in group:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

    # Evaluate groups
    for group, rules in admet_groups.items():
        df[group] = 'Fail'
        threshold = {
            'Absorption': 6,
            'Distribution': 2,
            'Metabolism': 8,
            'Excretion': 1,
            'Toxicity': 13
        }.get(group, 1)

        for i, row in df.iterrows():
            pass_count = sum(
                rules[col](row[col]) if pd.notnull(row[col]) else False
                for col in rules if col in row
            )
            if pass_count >= threshold:
                df.at[i, group] = 'Pass'

    # Composite suitability logic
    group_cols = ['Absorption', 'Distribution', 'Metabolism', 'Excretion', 'Toxicity']
    required_passes = 5 if strict else 4

    df['Suitable?'] = df[group_cols].apply(lambda r: 'Yes' if list(r).count('Pass') >= required_passes else 'No', axis=1)

    df = df[df['Suitable?'] == 'Yes'].reset_index(drop=True)

    ordered_columns = [
        'molecule_chembl_id',
        'canonical_smiles',
        'Absorption',
        'Distribution',
        'Metabolism',
        'Excretion',
        'Toxicity',
        'Suitable?'
    ]

    return df[ordered_columns]