app.py

import base64
import csv
import io
import json
import logging
import os
import time
import zipfile

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# Timing log configuration
TIMING_LOG_FILE = 'prediction_times.json'
MAX_LOG_ENTRIES = 100

def log_prediction_time(is_ifp_model, num_molecules, elapsed_time):
    """Log timing data and maintain rolling window of MAX_LOG_ENTRIES entries."""
    try:
        # Load existing log
        if os.path.exists(TIMING_LOG_FILE):
            with open(TIMING_LOG_FILE, 'r') as f:
                log_data = json.load(f)
        else:
            log_data = []

        # Add new entry
        from datetime import datetime
        log_data.append({
            "timestamp": datetime.now().isoformat(),
            "is_ifp": is_ifp_model,
            "num_molecules": num_molecules,
            "elapsed_time": elapsed_time
        })

        # Keep only the last MAX_LOG_ENTRIES entries
        if len(log_data) > MAX_LOG_ENTRIES:
            log_data = log_data[-MAX_LOG_ENTRIES:]

        # Save log
        with open(TIMING_LOG_FILE, 'w') as f:
            json.dump(log_data, f, indent=2)
    except Exception as e:
        logging.error(f"Error logging prediction time: {e}")

def get_timing_stats():
    """Calculate average time-per-molecule for IFP and standard models."""
    # Default values (fallback if no historical data)
    default_stats = {"ifp_time_per_mol": 45.0, "standard_time_per_mol": 2.0}

    try:
        if not os.path.exists(TIMING_LOG_FILE):
            return default_stats

        with open(TIMING_LOG_FILE, 'r') as f:
            log_data = json.load(f)

        if not log_data:
            return default_stats

        # Separate IFP and standard model timings
        ifp_times = []
        standard_times = []

        for entry in log_data:
            if entry.get("num_molecules", 0) > 0:
                time_per_mol = entry["elapsed_time"] / entry["num_molecules"]
                if entry.get("is_ifp", False):
                    ifp_times.append(time_per_mol)
                else:
                    standard_times.append(time_per_mol)

        # Calculate averages, use defaults if no data for that type
        stats = {
            "ifp_time_per_mol": sum(ifp_times) / len(ifp_times) if ifp_times else default_stats["ifp_time_per_mol"],
            "standard_time_per_mol": sum(standard_times) / len(standard_times) if standard_times else default_stats["standard_time_per_mol"]
        }

        return stats
    except Exception as e:
        logging.error(f"Error reading timing stats: {e}")
        return default_stats

# IMPORTANT: ifp (which imports vina) must be imported BEFORE rdkit to avoid segfault
# due to conflicting native library dependencies
try:
    import ifp
    IFP_AVAILABLE = True
except ImportError:
    IFP_AVAILABLE = False
    logger.warning("IFP module not available (vina not installed). IFP-based models will not work.")

from flask import Flask, Response, jsonify, render_template, request, send_file
from flask_cors import CORS
from werkzeug.utils import secure_filename
from qsprpred.models import SklearnModel
from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem, Draw
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer, Table, TableStyle

from qprf.qprf_utils import render_qprf
from spock.utils.standardizers.papyrus import PapyrusStandardizer

app = Flask(__name__)
app.jinja_env.filters['zip'] = zip
CORS(app)  # Allow all origins by default

# Configure logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

# Define the models directory
MODELS_DIR = 'models'

# Maximum number of SMILES accepted in a single prediction request
MAX_SMILES = 1000

# Folder where uploaded CSVs are stored so they persist across re-runs
UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Initialize SMILES standardizer for handling stereochemistry
smiles_standardizer = PapyrusStandardizer()

# define RDKit image implementer
def smiles_to_image(smiles):
    try:
        mol = Chem.MolFromSmiles(smiles) # attempt conversion to RDKit molecule
        if mol is None: # return None if not possible
            return None
        img = Draw.MolToImage(mol) # Image generation
        buffer = io.BytesIO()
        img.save(buffer, format="PNG")
        buffer.seek(0)
        img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
        buffer.close()
        return f"data:image/png;base64,{img_base64}"
    except Exception as e:
        logging.error(f"Error generating image for SMILES {smiles}: {e}") # Log the error message if any exception occurs during the process.
        return None

# define invalid SMILES scrubber
def validate_smiles(smiles_list):
    """
    Validates a list of SMILES strings. Returns a list of invalid SMILES.
    """
    invalid_smiles = []
    for smile in smiles_list:
        if Chem.MolFromSmiles(smile) is None:
            invalid_smiles.append(smile)
    return invalid_smiles

def standardize_smiles(smiles):
    """Standardize a SMILES string, handling stereochemistry."""
    try:
        standardized, _ = smiles_standardizer.convert_smiles(smiles)
        return standardized
    except Exception:
        return None

class SimilaritySearcher():
    def __init__(self):
        self.descgen = AllChem.GetMorganGenerator(radius=3)
        self.scorer = DataStructs.BulkTanimotoSimilarity

    def get_nearest_neighbors(self, smile, ms):
        """Find the most similar molecule in a reference set.

        Args:
            smile (str): smiles string
            ms (list): list of rdkit molecules from reference set

        Returns:
            tuple: (index of most similar molecule, similarity score)
                   Returns (None, 0.0) if query SMILES is invalid or no valid molecules in reference set
        """
        m1 = Chem.MolFromSmiles(smile)
        if m1 is None:
            logging.warning(f"Could not parse query SMILES: {smile}")
            return None, 0.0

        query_fp = self.descgen.GetSparseCountFingerprint(m1)

        # Filter out None molecules from reference set, keeping track of original indices
        valid_molecules = [(i, mol) for i, mol in enumerate(ms) if mol is not None]
        if not valid_molecules:
            logging.warning("No valid molecules in reference set")
            return None, 0.0

        original_indices, valid_mols = zip(*valid_molecules)
        target_fingerprints = [self.descgen.GetSparseCountFingerprint(x) for x in valid_mols]
        scores = self.scorer(query_fp, target_fingerprints)

        id_top = np.argmax(np.array(scores))

        # Return the original index (not the filtered index)
        return original_indices[id_top], max(scores)

def extract_model_info(directory):
    models_info = []
    for d in os.listdir(directory):
        meta_path = os.path.join(directory, d, f"{d}_meta.json")
        if os.path.isfile(meta_path):
            with open(meta_path, 'r') as meta_file:
                meta_data = json.load(meta_file)
                state = meta_data['py/state']
                curr_dir = os.path.join(directory, d)
                model_info = {
                    'name': state['name'],
                    'pref_name': state['pref_name'],
                    'case_study': state['case_study'],
                    'target_property_name': state['targetProperties'][0]['py/state']['name'],
                    'target_property_task': state['targetProperties'][0]['py/state']['task']['py/reduce'][1]['py/tuple'][0],
                    'feature_calculator': state['featureCalculators'][0]['py/object'].split('.')[-1],
                    'algorithm': state['alg'].split('.')[-1],
                    'currDir': curr_dir,
                    # Only some models ship a QMRF document; used to hide a dead download link
                    'has_qmrf': os.path.isfile(os.path.join(curr_dir, f"qmrf_{state['name']}.docx")),
                }
                models_info.append(model_info)
                logging.info(f"Loaded model metadata: {model_info}")
    return models_info

@app.route('/download')
def download_qmrf():
    path = request.args.get('path')
    target = request.args.get('target')
    qmrf_path = f'{path}/qmrf_{target}.docx'
    if not os.path.isfile(qmrf_path):
        return jsonify({'error': f'No QMRF document is available for {target}.'}), 404
    return send_file(qmrf_path, as_attachment=True)

@app.route('/downloadqprf')
def download_qprf():
    model = request.args.get('model')
    smile = request.args.get('smile')
    return send_file(f"qprf/output/{model}/{smile}.docx", as_attachment=True)

@app.route('/downloadposes')
def download_poses():
    smile = request.args.get('smile')
    protein_id = request.args.get('protein_id')
    return send_file(f"{smile}_{protein_id}.sdf", as_attachment=True)

# def zipdir(path, ziph):
#     # ziph is zipfile handle
#     for root, dirs, files in os.walk(path):
#         for file in files:
#             ziph.write(os.path.join(root, file),
#                        os.path.relpath(os.path.join(root, file),
#                                        os.path.join(path, '..')))



# @app.route('/downloadposes')
# def download_poses():
#     logging.basicConfig(filename='myapp.log', level=logging.INFO)
#     logger.info('Started')
#     store_names = request.args.get('store_names')
#     logger.info(store_names)
#     for store_name in store_names:
#         with zipfile.ZipFile('Python.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
#             zipdir(f'{store_name}/', zipf)
#     return send_file('Python.zip', as_attachment=True)

@app.route('/')
@app.route('/predict')
def home():
    available_models = extract_model_info(MODELS_DIR)
    timing_stats = get_timing_stats()
    return render_template('index.html', models=available_models, timing_stats=timing_stats)

@app.route('/predict', methods=['POST'])
def predict():
    logging.info("Handling prediction request.")
    start_time = time.time()
    available_models = extract_model_info(MODELS_DIR)
    try:
        smiles_input = request.form.get('smiles')
        uploaded_file = request.files.get('file')
        model_names = request.form.getlist('model')
        file_name = request.form.get('uploaded_file_name')

        logging.debug(f"Received SMILES input: {smiles_input}")
        logging.debug(f"Uploaded file: {uploaded_file}")
        logging.debug(f"Selected models: {model_names}")
        logging.debug(f"Previous uploaded file name: {file_name}")

        if not model_names:
            logging.error("No model selected.")
            timing_stats = get_timing_stats()
            return render_template('index.html', models=available_models, timing_stats=timing_stats, error="No model selected.")

        smiles_list = []
        invalid_smiles = []

        # Handle SMILES string input
        if smiles_input:
            input_smiles = [smile.strip() for smile in smiles_input.split(',')]

            # Check if only one SMILES string is entered
            if len(input_smiles) == 1:
                standardized = standardize_smiles(input_smiles[0])
                if standardized is None:  # Check for invalid single SMILES
                    logging.error(f"Invalid SMILES string: {input_smiles[0]}")  # Log the invalid SMILES
                    timing_stats = get_timing_stats()
                    return render_template('index.html', models=available_models, timing_stats=timing_stats, error="Invalid SMILES string")  # Display error for single invalid SMILES
                else:
                    smiles_list.append(standardized)  # Add valid standardized SMILES to processing list
            else:
                for smile in input_smiles:
                    standardized = standardize_smiles(smile)
                    if standardized is None:
                        invalid_smiles.append(smile)  # Collect invalid SMILES
                    else:
                        smiles_list.append(standardized)  # Collect valid standardized SMILES

        # Handle uploaded file
        if uploaded_file and uploaded_file.filename != '':
            file_name = secure_filename(uploaded_file.filename)
            logging.debug("Processing uploaded file.")
            # Persist the upload so a subsequent run can reuse it without re-uploading
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_name)
            uploaded_file.save(file_path)
            uploaded_df = pd.read_csv(file_path)
            logging.debug(f"Uploaded file contents: {uploaded_df.head()}")
            if 'SMILES' in uploaded_df.columns:
                file_smiles = uploaded_df['SMILES'].tolist()
                for smile in file_smiles:
                    standardized = standardize_smiles(smile)
                    if standardized is None:
                        invalid_smiles.append(smile)  # Collect invalid SMILES from file
                    else:
                        smiles_list.append(standardized)  # Collect valid standardized SMILES from file
        elif file_name:
            file_name = secure_filename(file_name)
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], file_name)
            if os.path.exists(file_path):
                logging.debug("Reprocessing previous uploaded file.")
                uploaded_df = pd.read_csv(file_path)
                if 'SMILES' in uploaded_df.columns:
                    file_smiles = uploaded_df['SMILES'].tolist()
                    for smile in file_smiles:
                        standardized = standardize_smiles(smile)
                        if standardized is None:
                            invalid_smiles.append(smile)  # Collect invalid SMILES from previous file
                        else:
                            smiles_list.append(standardized)  # Collect valid standardized SMILES from previous file

        logging.debug(f"Final SMILES list: {smiles_list}")
        logging.debug(f"Invalid SMILES detected: {invalid_smiles}")  # Log invalid SMILES

        if not smiles_list and not invalid_smiles:
            error_message = "No SMILES strings provided"
            logging.error(error_message)
            timing_stats = get_timing_stats()
            return render_template('index.html', models=available_models, timing_stats=timing_stats, error=error_message)

        total_submitted = len(smiles_list) + len(invalid_smiles)
        if total_submitted > MAX_SMILES:
            error_message = (f"Too many SMILES submitted: {total_submitted} > {MAX_SMILES}. "
                             "Please split your input into smaller batches.")
            logging.error(error_message)
            timing_stats = get_timing_stats()
            return render_template('index.html', models=available_models, timing_stats=timing_stats, error=error_message)

        # Identical structures (e.g. the same molecule written two ways) standardize to
        # the same SMILES and collapse to a single entry in the docking store, which
        # breaks the per-molecule mapping. Reject duplicates up front with a clear,
        # fast message instead of failing deep in the (slow) docking pipeline.
        unique_smiles = list(dict.fromkeys(smiles_list))
        if len(unique_smiles) < len(smiles_list):
            error_message = (f"Duplicate structures detected: {len(smiles_list)} molecules were "
                             f"submitted but only {len(unique_smiles)} are unique after "
                             "standardization. Please submit each unique structure once.")
            logging.error(error_message)
            timing_stats = get_timing_stats()
            return render_template('index.html', models=available_models, timing_stats=timing_stats, error=error_message)

        store_names = []
        all_predictions = {}
        all_raw_predictions = {}
        all_ads = {}
        model_info_dict = {}
        used_ifp_model = False
        data_dict = {
            "Model": {},
            "Input": {'SMILES': {}, 'Input structure': {}},
            "Nearest Neighbor": {'NN structure': {}, 'Similarity': {}, 'Source': {}, 'Experimental value': {}},
            "Output": {'Predicted pChEMBL Value': {}, 'Within Applicability Domain': {}, 'Download QPRF': {}}
        }
        tooltip_dict = {"Model": 'Unique identifier of the model that made the prediction',
                        "Input structure": '2D depiction of molecule',
                        "NN structure": '2D depiction of molecule',
                        "SMILES": 'Line representation of molecule',
                        "Nearest Neighbor": 'Molecule from the model training set that is the most similar to input molecule',
                        "Similarity": 'Tanimoto similarity score based on same chemical descriptor as used for model',
                        "Source": 'Document(s) containing experimental data for nearest neighbor',
                        "Experimental value": 'Average of all reported experimental values',
                        "Predicted pChEMBL Value": 'Model prediction for input molecule. pChEMBL is defined as -log(response). More information available in QMRF & QPRF',
                        "Within Applicability Domain": 'AD is based on descriptors of training set. An input molecule is within AD if the distance to the training set is lower than a set threshold. More information available in QMRF & QPRF',
                        "Download QPRF": 'Automatically filled in report about the prediction',
                        "Download Poses": 'Get SDF with docked ligand poses'
                        }

        for model_name in model_names:
            logging.debug(f"Processing model: {model_name}")
            model_path = os.path.join(MODELS_DIR, model_name, f"{model_name}_meta.json")
            model = SklearnModel.fromFile(model_path)
            ad = []
            if model.featureCalculators[0].__class__.__name__ == "DataFrameDescriptorSet":
                if not IFP_AVAILABLE:
                    logging.error(f"Model {model_name} requires IFP/vina but vina is not installed")
                    return jsonify({"error": f"Model {model_name} requires docking (vina) which is not available"}), 500
                predictions, ad, lib_name = ifp.predict(smiles_list, model, model.name.replace("_final", ""), "dimorph")
                ad = list(ad)
                store_names.append(lib_name)
                used_ifp_model = True
            elif getattr(model, 'applicabilityDomain', None):
                predictions, ad = model.predictMols(smiles_list, use_applicability_domain=True)
                ad = list(ad)
            else:
                predictions = model.predictMols(smiles_list)

            # Check if the model is regression or classification
            if model.task.isRegression():
                formatted_predictions = [f"{pred[0]:.2f}" for pred in predictions]
            else:
                # Format classification output as Active/Inactive
                formatted_predictions = ["Active" if pred[0] == 1 else "Inactive" for pred in predictions]

            all_predictions[model_name] = formatted_predictions
            all_raw_predictions[model_name] = predictions
            all_ads[model_name] = ad

            # Extract model info for report
            with open(model_path, 'r') as meta_file:
                meta_data = json.load(meta_file)
                state = meta_data['py/state']
                model_info = {
                    'name': state['name'],
                    'pref_name': state['pref_name'],
                    'case_study': state['case_study'],
                    'target_property_name': state['targetProperties'][0]['py/state']['name'],
                    'target_property_task': state['targetProperties'][0]['py/state']['task']['py/reduce'][1]['py/tuple'][0],
                    'feature_calculator': state['featureCalculators'][0]['py/object'].split('.')[-1],
                    'algorithm': state['alg'].split('.')[-1]
                }
                model_info_dict[model_name] = model_info

        table_data = []

        for i, smile in enumerate(smiles_list):
            image_data = smiles_to_image(smile)
            # Check if any model has AD data
            if any(all_ads[m] for m in model_names):
                row = [image_data] + [smile] + [
                    all_predictions[m][i] + (f' ({format_ad(all_ads[m][i])})' if all_ads[m] else '')
                    for m in model_names
                ]
            else:
                row = [image_data] + [smile] + [all_predictions[m][i] for m in model_names]

            table_data.append(row)

        # Update headers
        headers = ['Structure', 'SMILES']
        tooltips = ['2D depiction of input molecule', 'Line representation of input molecule']

        searcher = SimilaritySearcher()
        for model_name in model_names:
            accession = model_name.split("_")[0]
            train_df = pd.read_csv(f'data/{accession}_Data/train_full_model_{accession}.csv').reset_index()
            train_smiles = train_df['SMILES'].to_list()
            ms = [Chem.MolFromSmiles(x) for x in train_smiles]
            model_path = os.path.join(MODELS_DIR, model_name, f"{model_name}_meta.json")
            model = SklearnModel.fromFile(model_path)

            if getattr(model, 'applicabilityDomain', None):
                if model.task.isRegression():
                    # Format regression table header
                    headers.append(f'Predicted pChEMBL Value ({model_name})')
                    tooltips.append('Model prediction for input molecule. pChEMBL is defined as -log(response). The value in brackets shows if input molecule is within AD')
                else:
                    # Format classification table header
                    headers.append(f'Predicted class label ({model_name})')
                    tooltips.append('Model prediction for input molecule. The value in brackets shows if input molecule is within AD')
            else:
                if model.task.isRegression():
                    # Format regression table header
                    headers.append(f'Predicted pChEMBL Value ({model_name})')
                    tooltips.append('Model prediction for input molecule. pChEMBL is defined as -log(response)')
                else:
                    # Format classification table header
                    headers.append(f'Predicted class label ({model_name})')
                    tooltips.append('Model prediction for input molecule')


            for i, smile in enumerate(smiles_list):
                image_data = smiles_to_image(smile)
                id_top, score = searcher.get_nearest_neighbors(smile, ms)

                # Handle case where nearest neighbor search failed
                if id_top is None:
                    logging.warning(f"Could not find nearest neighbor for SMILES: {smile}")
                    # Use placeholder values for nearest neighbor data
                    nearest_neighbor = {
                        "smiles": "N/A",
                        "reference": ["N/A"],
                        "value": 0.0,
                        "predicted_value": "N/A",
                        "similarity": "Could not compute similarity"
                    }
                    nn_smiles = None
                    doc_ids = ["N/A"]
                    doc_links = []
                    image_data_nn = None
                else:
                    nearest_neighbor = {}
                    nn_smiles = train_df.iloc[id_top]['SMILES']
                    nearest_neighbor["smiles"] = nn_smiles
                    doc_ids = train_df.iloc[id_top]['all_doc_ids']
                    doc_ids = str(doc_ids).split(';')
                    doc_links = []
                    for doc_id in doc_ids:
                        if doc_id.startswith('PMID:'):
                            doc_links.append(f'https://pubmed.ncbi.nlm.nih.gov/{doc_id.lstrip("PMID:")}/')
                        elif doc_id.startswith('DOI:'):
                            doc_links.append('https://doi.org/' + doc_id)
                        elif doc_id.startswith('PubChemAID:'):
                            doc_links.append(f'https://pubchem.ncbi.nlm.nih.gov/bioassay/{doc_id.lstrip("PubChemAID:")}/')

                    nearest_neighbor["reference"] = doc_ids
                    nearest_neighbor["value"] = train_df.iloc[id_top][model_info_dict[model_name]['target_property_name']]
                    if model.featureCalculators[0].__class__.__name__ == "DataFrameDescriptorSet":
                        if IFP_AVAILABLE:
                            nearest_neighbor["predicted_value"] = ifp.predict([nn_smiles], model, model.name.replace("_final", ""), "dimorph")[0][0]
                        else:
                            nearest_neighbor["predicted_value"] = "N/A (vina not available)"
                    else:
                        nearest_neighbor["predicted_value"] = model.predictMols([nn_smiles])[0][0]
                    nearest_neighbor["similarity"] = f"Nearest neighbor was found using {searcher.scorer.__name__} based on {searcher.descgen.__class__.__name__}"
                    image_data_nn = smiles_to_image(nn_smiles)

                # Use correct predictions and AD values for this model
                raw_pred = all_raw_predictions[model_name][i] if model_name in all_raw_predictions else None
                raw_ad = all_ads[model_name][i] if all_ads[model_name] else None
                render_qprf(smile, model, raw_pred, raw_ad, nearest_neighbor)
                data_dict["Model"].setdefault("value", []).append(model_name)
                data_dict["Input"]["Input structure"].setdefault("image", []).append(image_data)
                data_dict["Input"]["SMILES"].setdefault("value", []).append(smile)
                data_dict["Nearest Neighbor"]["NN structure"].setdefault("image", []).append(image_data_nn)
                data_dict["Nearest Neighbor"]["Similarity"].setdefault("value", []).append(f"{score:.2f}")
                data_dict["Nearest Neighbor"]["Source"].setdefault("value", []).append(doc_ids)
                data_dict["Nearest Neighbor"]["Source"].setdefault("links", []).append(doc_links)
                data_dict["Nearest Neighbor"]["Experimental value"].setdefault("value", []).append(f"{nearest_neighbor['value']:.2f}")
                data_dict["Output"]["Predicted pChEMBL Value"].setdefault("value", []).append(all_predictions[model_name][i])
                if len(all_ads[model_name]) > i:
                    data_dict["Output"]["Within Applicability Domain"].setdefault("value", []).append(all_ads[model_name][i])
                else:
                    data_dict["Output"]["Within Applicability Domain"].setdefault("value", []).append(None)

                data_dict["Output"]["Download QPRF"].setdefault("url", []).append("hi")
                if model.featureCalculators[0].__class__.__name__ == "DataFrameDescriptorSet":
                    data_dict["Output"].setdefault("Download Poses", {}).setdefault("pose", []).append("hi")
        error_message = None
        if invalid_smiles:
            error_message = f"Invalid SMILES, could not be processed: {', '.join(invalid_smiles)}"  # Mention invalid SMILES in error message

        elapsed_time = time.time() - start_time

        # Log timing data for self-learning estimates
        log_prediction_time(used_ifp_model, len(smiles_list), elapsed_time)
        timing_stats = get_timing_stats()

        # Generate a downloadable report bundle (PDF + CSV + QMRF + QPRF) instead of re-rendering
        if request.form.get('download_report'):
            return build_report_bundle(model_names, smiles_list, model_info_dict, headers, table_data)

        return render_template('index.html',
                               models=available_models,
                               data_dict=data_dict,
                               headers=headers,
                               tooltips=tooltips,
                               data=table_data,
                               tooltips_extensive = tooltip_dict,
                               smiles_input=smiles_input,
                               model_names=model_names,
                               file_name=file_name,
                               store_names=store_names,
                               elapsed_time=elapsed_time,
                               timing_stats=timing_stats,
                               error=error_message)
    except ifp.DockingError as e:
        logging.warning(f"Docking could not complete: {e}")
        timing_stats = get_timing_stats()
        return render_template('index.html', models=available_models, timing_stats=timing_stats, error=str(e))
    except Exception:
        logging.exception("An error occurred while processing the request.")
        timing_stats = get_timing_stats()
        return render_template('index.html', models=available_models, timing_stats=timing_stats, error="An error occurred while processing the request.")

def format_ad(ad_value):
    """Render an applicability-domain value as a readable label instead of a raw array."""
    try:
        flag = bool(ad_value[0]) if hasattr(ad_value, '__len__') else bool(ad_value)
    except (TypeError, IndexError):
        return str(ad_value)
    return 'within AD' if flag else 'outside AD'


def build_report_bundle(model_names, smiles_list, model_info_dict, headers, table_data):
    """Build a ZIP with the summary PDF, a predictions CSV, and per-model QMRF /
    per-prediction QPRF documents, and return it as a download."""
    # Summary PDF (create_report mutates the rows it receives, so pass a fresh copy)
    report_models = [model_info_dict[m] for m in model_names if m in model_info_dict]
    report_headers = headers[1:]  # drop the 'Structure' image column
    report_rows = [[row[1]] + [str(cell) for cell in row[2:]] for row in table_data]
    pdf_buffer = create_report(report_models, report_headers, report_rows)

    # Predictions CSV (rebuilt from the original table_data, text only)
    csv_buffer = io.StringIO()
    writer = csv.writer(csv_buffer)
    writer.writerow(headers[1:])
    for row in table_data:
        writer.writerow([row[1]] + [str(cell) for cell in row[2:]])

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf:
        zf.writestr('prediction_report.pdf', pdf_buffer.getvalue())
        zf.writestr('predictions.csv', csv_buffer.getvalue())
        for model_name in model_names:
            # QMRF: per-model methodology document shipped alongside the model
            qmrf_path = os.path.join(MODELS_DIR, model_name, f'qmrf_{model_name}.docx')
            if os.path.exists(qmrf_path):
                zf.write(qmrf_path, arcname=f'qmrf/{model_name}.docx')
            # QPRF: per-prediction report generated during the prediction loop
            for idx, smile in enumerate(smiles_list):
                qprf_path = os.path.join('qprf', 'output', model_name, f'{smile}.docx')
                if os.path.exists(qprf_path):
                    safe = secure_filename(smile) or f'molecule_{idx + 1}'
                    zf.write(qprf_path, arcname=f'qprf/{model_name}/{safe}.docx')
    zip_buffer.seek(0)
    response = send_file(zip_buffer, as_attachment=True,
                         download_name='qsprpred_report.zip', mimetype='application/zip')
    # Echo the client's download token back as a cookie so the UI knows the file was
    # delivered and can dismiss the loading overlay (the page itself never navigates).
    token = request.form.get('download_token')
    if token:
        response.set_cookie('fileDownloadToken', token, max_age=60)
    return response


def create_report(model_info_list, headers, table_data):
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    elements = []

    # Add title
    styles = getSampleStyleSheet()
    title_style = styles['Title']
    elements.append(Paragraph("Prediction Report", title_style))
    elements.append(Spacer(1, 12))

    # Add model metadata
    for model_info in model_info_list:
        elements.append(Paragraph(f"Model: {model_info['name']}", styles['Heading2']))
        elements.append(Paragraph(f"Model Name: {model_info['pref_name']}", styles['Heading2']))
        elements.append(Paragraph(f"Target Property Name: {model_info['target_property_name']}", styles['Normal']))
        elements.append(Paragraph(f"Target Property Task: {model_info['target_property_task']}", styles['Normal']))
        elements.append(Paragraph(f"Feature Calculator: {model_info['feature_calculator']}", styles['Normal']))
        elements.append(Paragraph(f"Algorithm: {model_info['algorithm']}", styles['Normal']))
        elements.append(Spacer(1, 12))

    # Convert headers to Paragraphs for wrapping
    header_paragraphs = [Paragraph(header, styles['Normal']) for header in headers]

    # Convert SMILES strings to Paragraphs for wrapping
    for i in range(len(table_data)):
        table_data[i][0] = Paragraph(table_data[i][0], styles['Normal'])

    # Add prediction table
    data = [header_paragraphs] + table_data
    table = Table(data, repeatRows=1)

    # Apply style to table
    table.setStyle(TableStyle([
#        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
#        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, 0), 12),
        ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
#        ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
        ('GRID', (0, 0), (-1, -1), 1, colors.black),
    ]))

    # Ensure SMILES strings can wrap
    table._argW[0] = 2.5 * inch  # width of SMILES column
    elements.append(table)

    doc.build(elements)
    buffer.seek(0)
    return buffer

@app.route('/get_models', methods=['GET'])
def get_models():
    """Return JSON metadata for all available models (name, target, task, algorithm, etc.)."""
    return jsonify(extract_model_info(MODELS_DIR))

@app.route('/api', methods=['POST'])
def apipredict():
    data = request.json
    smiles = data.get('smiles', [])
    models = data.get('models', [])
    output_format = data.get('format', 'json')  # Default to JSON format

    if not smiles or not isinstance(smiles, list):
        return jsonify({'error': 'Invalid input: please provide a list of SMILES strings.'}), 400

    if not models or not isinstance(models, list):
        return jsonify({'error': 'Invalid input: please provide a list of model names.'}), 400

    all_predictions = {}
    all_ads = {}

    for model_name in models:
        model_path = os.path.join(MODELS_DIR, model_name, f"{model_name}_meta.json")
        if not os.path.exists(model_path):
            return jsonify({'error': f"Model {model_name} does not exist."}), 400

        model = SklearnModel.fromFile(model_path)
        # Include the applicability domain in the output when the model supports it (#41)
        if getattr(model, 'applicabilityDomain', None):
            predictions, ad = model.predictMols(smiles, use_applicability_domain=True)
            all_ads[model_name] = [bool(ad[i][0]) for i in range(len(smiles))]
        else:
            predictions = model.predictMols(smiles)
            all_ads[model_name] = None
        all_predictions[model_name] = [f"{pred[0]:.4f}" for pred in predictions]

    # Format the result
    result = []
    for i, smile in enumerate(smiles):
        result_entry = {'smiles': smile}
        for model_name in models:
            result_entry[f'prediction ({model_name})'] = all_predictions[model_name][i]
            if all_ads[model_name] is not None:
                result_entry[f'within_applicability_domain ({model_name})'] = all_ads[model_name][i]
        result.append(result_entry)

    if output_format == 'text':
        result_text = '\n'.join([f"SMILES: {entry['smiles']} -> " + ", ".join([f"{model}: {pred}" for model, pred in entry.items() if model != 'smiles']) for entry in result])
        return Response(result_text, mimetype='text/plain')

    if output_format == 'csv':
        # Create an in-memory output file for CSV data
        output = io.StringIO()
        writer = csv.DictWriter(output, fieldnames=result[0].keys())
        writer.writeheader()
        writer.writerows(result)
        output.seek(0)
        return Response(output.getvalue(), mimetype='text/csv', headers={"Content-Disposition": "attachment;filename=predictions.csv"})


    return jsonify(result)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)