"""Main driver for module."""
import argparse
import logging
from sys import version_info
from osmolytes.pqr import parse_pqr_file, count_residues
from osmolytes.sasa import SolventAccessibleSurface
from osmolytes.energy import transfer_energy


assert version_info > (3, 5)
_LOGGER = logging.getLogger("main")


def build_parser():
    """Build argument parser.

    :returns:  argument parser
    :rtype:  argparse.ArgumentParser
    """
    parser = argparse.ArgumentParser(
        description=(
            "This code uses various models to attempt to predict the "
            "influence of osmolytes on protein stability"
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--log-level",
        help="set verbosity of code",
        nargs=1,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
    )
    parser.add_argument(
        "--model",
        help=(
            "model to use for stability prediction:\n"
            "  'auton' -- Auton M, Bolen DW. Predicting the energetics of "
            "osmolyte-induced protein folding/unfolding. Proc Natl Acad Sci "
            "102:15065 (2005) https://doi.org/10.1073/pnas.0507053102"
        ),
        nargs=1,
        default="auton",
        choices=["auton"],
    )
    parser.add_argument(
        "--surface-output",
        nargs=1,
        default=None,
        help="path for output of surface in XYZ format",
    )
    parser.add_argument(
        "--solvent-radius",
        nargs=1,
        default=1.4,
        type=float,
        help="radius of solvent molecule (in Å)",
    )
    parser.add_argument(
        "--surface-points",
        nargs=1,
        default=2000,
        type=int,
        help="number of points per atom for constructing surface",
    )
    parser.add_argument(
        "pqr_path",
        help=(
            "input PQR file for protein, as generated by PDB2PQR or "
            "similar program (see documentation) "
        ),
    )
    return parser


def main():
    """Main driver."""
    parser = build_parser()
    args = parser.parse_args()
    logging.basicConfig(level=getattr(logging, args.log_level))
    _LOGGER.debug(f"Got command-line arguments: {args}")
    _LOGGER.info(f"Reading PQR input from {args.pqr_path}...")
    with open(args.pqr_path, "rt") as pqr_file:
        atoms = parse_pqr_file(pqr_file)
    count_df = count_residues(atoms)
    _LOGGER.info(f"Protein composition:\n{count_df}")
    _LOGGER.info(
        f"Constructing protein solvent-accessible surface with "
        f"solvent radius {args.solvent_radius} and {args.surface_points} "
        f"points per atom."
    )
    sas = SolventAccessibleSurface(
        atoms,
        probe_radius=args.solvent_radius,
        num_points=args.surface_points,
        xyz_path=args.surface_output,
    )
    _LOGGER.info("Calculating transfer energies.")
    energy_df = transfer_energy(atoms, sas)
    _LOGGER.info(f"Detailed energies (kcal/mol/M):\n{energy_df.to_string()}")
    energies = energy_df.sum(axis=0).sort_values()
    _LOGGER.info(f"Summary energies (kcal/mol/M):\n{energies}")
