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


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=(
            f"Osmolytes v. {osmolytes.__version__}. "
            f"This code uses various models to attempt to predict the "
            f"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",
        default=None,
        help="path for output of surface in XYZ format",
    )
    parser.add_argument(
        "--solvent-radius",
        default=1.4,
        type=float,
        help="radius of solvent molecule (in Å)",
    )
    parser.add_argument(
        "--surface-points",
        default=2000,
        type=int,
        help="number of points per atom for constructing surface",
    )
    parser.add_argument(
        "--output",
        choices=["xlsx", "csv"],
        help="output m-value results in Excel (xlsx) or CSV (csv) format",
    )
    parser.add_argument(
        "--output-dir", default=".", help="directory for m-value output"
    )
    parser.add_argument(
        "pqr_path",
        help=(
            "input PQR file for protein, as generated by PDB2PQR or "
            "similar program (see documentation) "
        ),
    )
    return parser


def main(args=None):
    """Main entry point.

    :param list(str) args:  list of strings to parse
    :returns:  dictionary of results
    :rtype:  dict
    """
    parser = build_parser()
    args = parser.parse_args(args)
    logging.basicConfig(level=getattr(logging, args.log_level))
    _LOGGER.info(f"Osmolytes version {osmolytes.__version__}")
    _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)
    _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,
    )
    counts = count_residues(atoms)
    areas = get_folded_areas(atoms, sas).rename(
        {"sidechain": "sidechain areas", "backbone": "backbone areas"},
        axis="columns"
    )
    areas["residue counts"] = counts
    print(areas)
    areas = areas[["residue counts", "sidechain areas", "backbone areas"]]
    _LOGGER.info(f"Protein composition:\n{areas}")
    _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}")
    output_dir = Path(args.output_dir)
    if args.output == "xlsx":
        output_path = output_dir / f"{Path(args.pqr_path).stem}-mvalues.xlsx"
        _LOGGER.info(f"Writing energies (kcal/mol/M) to {output_path}")
        energy_df.to_excel(output_path)
    elif args.output == "csv":
        output_path = output_dir / f"{Path(args.pqr_path).stem}-mvalues.csv"
        _LOGGER.info(f"Writing energies (kcal/mol/M) to {output_path}")
        energy_df.to_csv(output_path)
    return {
        "args": args,
        "atoms": atoms,
        "sas": sas,
        "energy_df": energy_df,
        "energies": energies,
    }


def console():
    """Run code from the console."""
    main()
