Source code for qp.structure.setup

"""Fetch PDB and perform quality checks"""

import os
import csv
import sys
import yaml
import requests


def read_config(config_file):
    """Read and parse a YAML configuration file.

    Parameters
    ----------
    config_file : str
        Path to the configuration YAML file.

    Returns
    -------
    dict
        Parsed configuration data with all pipeline parameters.
    """
    with open(config_file, 'r') as file:
        return yaml.safe_load(file)


def parse_input(input, output, center_yaml_residues):
    """Parse input sources and determine center residues for each structure.

    Processes the user-provided input (PDB codes, file paths, or CSV) and
    determines the center residue definitions from either the CSV ``center``
    column or the YAML ``center_residues`` parameter.

    Parameters
    ----------
    input : str or list
        Input specification: PDB code(s), path to PDB file(s), or path to CSV.
    output : str
        Path to the output directory.
    center_yaml_residues : list
        Center residue definitions from the YAML config file.

    Returns
    -------
    tuple of (list, list)
        ``(pdb_all, center_residues)`` where ``pdb_all`` is a list of
        ``(pdb_id, pdb_path)`` tuples and ``center_residues`` is a list
        of center definition strings.

    Raises
    ------
    SystemExit
        If no center residues are provided in either the CSV or YAML.
    """

    # If the input was a path or a single PDB, convert it to a list
    if isinstance(input, str):
        input = [input]
    output = os.path.abspath(output)
    center_csv_residues = get_centers(input)
    pdb_all = get_pdbs(input, output)

    # Determine how the user has decided to provide the center residues
    if center_csv_residues:
        use_csv_centers = True
        print("> Using residues from the input csv as centers\n")
    elif center_yaml_residues:
        use_csv_centers = False
        print("> Using residues from the input yaml as centers.\n")
    else:
        print("> No center residues were provided in the input csv or the yaml.\n")
        exit()

    # Determine the center residues for the current PDB
    if use_csv_centers:
        center_residues = center_csv_residues
    else:
        if not center_yaml_residues:
            sys.exit("> No more center residues available.")
        center_residues = center_yaml_residues

    return pdb_all, center_residues


def fetch_pdb(pdb, out):
    """
    Fetches the PDB file for a given PDB code.

    Parameters
    ----------
    pdb: str
        PDB code
    out: str
        Path to output PDB file

    Raises
    ------
    ValueError
        If the PDB ID returns a 404 error (invalid user input).
    IOError
        If there is a network or server-side error.
    """
    url = f"https://files.rcsb.org/view/{pdb}.pdb"
    try:
        r = requests.get(url, timeout=15)
    except requests.exceptions.RequestException as e:
        # Raise an IOError for network-level problems
        raise IOError(f"Could not connect to the server. Details: {e}")

    # Check the status code from the server's response
    if r.status_code == 200:
        os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True)
        with open(out, "w") as f:
            f.write(r.text)
    elif r.status_code == 404:
        # Raise a ValueError for an invalid PDB ID
        raise ValueError(f"PDB ID '{pdb}' is not valid or does not exist.")
    else:
        # Raise an IOError for other server-side problems
        raise IOError(f"Server returned an error with status code {r.status_code}.")


def get_pdbs(input_path, output_path):
    """
    Parses the input PDBs and returns a list of tuples

    Parameters
    ----------
    input_path: list
        List of PDB codes, paths to PDB files, or path to CSV file
    output_path: str
        Path to output directory

    Returns
    -------
    pdb_all
        List of tuples containing parsed PDB ID and path to PDB file


    Notes
    -----
    Store input PDBs as a tuple of
    Parsed ID (PDB code or filename)
    Path to PDB file (existing or to download)
    """
    pdb_all = []
    for pdb_id in input_path:
        if os.path.isfile(pdb_id):
            pdb, ext = os.path.splitext(os.path.basename(pdb_id))
            pdb = pdb.replace(".", "_")
            if ext == ".pdb":
                pdb_all.append((pdb, pdb_id))
            elif ext == ".csv":
                with open(pdb_id, "r") as csvfile:
                    reader = csv.DictReader(csvfile)
                    for row in reader:
                        pdb = row['pdb_id']
                        pdb_all.append((pdb, os.path.join(output_path, pdb, f"{pdb}.pdb")))
            else:
                with open(pdb_id, "r") as f:
                    pdb_all.extend(
                        [(pdb, os.path.join(output_path, pdb, f"{pdb}.pdb")) for pdb in f.read().splitlines()]
                    )
        else:
            pdb_all.append((pdb_id, os.path.join(output_path, pdb_id, f"{pdb_id}.pdb")))

    return pdb_all


def get_centers(input_path):
    """
    Parses the input centers and returns a list of the center residue PDB IDs.

    Parameters
    ----------
    input_path: list
        List of pdbs or the input csv file.

    Returns
    -------
    centers: list
        List PDB IDs for user-defined center residues.
    """
    
    centers = []
    for pdb_id in input_path:
        if os.path.isfile(pdb_id):
            pdb, ext = os.path.splitext(os.path.basename(pdb_id))
            pdb = pdb.replace(".", "_")
            if ext == ".csv":
                input_csv = pdb_id
                with open(input_csv, "r") as csvfile:
                    reader = csv.DictReader(csvfile)
                    # Check if 'center' column exists
                    if 'center' not in reader.fieldnames:
                        print(f"> WARNING: The 'center' option is not being used in {input_csv}. Returning empty list.")
                        return []
                    # If the 'center' column exists, proceed to collect centers
                    for row in reader:
                        center = row.get('center', None)
                        if center:
                            centers.append(center)
    return centers