Source code for schrodinger.application.desmond.license

"""
Module for license related code in Desmond workflows.

Copyright Schrodinger LLC, All Rights Reserved.
"""
from collections import Counter
from pathlib import Path
from typing import List

from schrodinger.application.desmond import bld_def
from schrodinger.application.desmond import bld_ver
from schrodinger.application.desmond.constants import CUSTOM_CHARGE_MODE
from schrodinger.infra import licensing
from schrodinger.utils import license

DESMOND_GPU_TOKEN_MULTIPLIER = 16
JAGUAR_TOKEN_MULTIPLIER = 2


def add_fep_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the FEP GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.
    :param use_custom_charges: Set to True when using custom charges (F16).

    :return: The updated command.
    """
    lic_args = ["-lic", fep_lic(_get_num_procs(args))]
    return _deduplicate_lics(args + lic_args)


def add_md_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    MD GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(args + ["-lic", md_lic(_get_num_procs(args))])


def add_md_cpu_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    MD CPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(args + ["-lic", md_cpu_lic(_get_num_procs(args))])


def add_watermap_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    Watermap GPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(
        args + ["-lic", watermap_lic(_get_num_procs(args))])


def add_watermap_cpu_lic(args: List[str]) -> List[str]:
    """
    Add the command line arguments to properly handle the
    Watermap CPU license.

    NOTE: If the same license is present with a different
    number of tokens, this will override the current value.

    :param args: List of command line arguments. This function will look into
        `args` for the -HOST and -cpu options to determine the number of processors.
        If these are not found, this function will assume one processor.

    :return: The updated command.
    """
    return _deduplicate_lics(
        args + ["-lic", watermap_cpu_lic(_get_num_procs(args))])


def protein_mutation_lic() -> str:
    """
    Return the license string needed for Plop, used in the
    protein mutation generator stage.
    """
    return "%s:8" % licensing.getFeatureName(licensing.PSP_PLOP)


def fep_lic(ngpus: int) -> str:
    """
    Return the license string for FEP GPU, given the number of GPUs
    to be used.
    """
    license_string = licensing.getFeatureName(licensing.FEP_GPGPU)
    return f"{license_string}:{ngpus * DESMOND_GPU_TOKEN_MULTIPLIER}"


def md_lic(ngpus: int) -> str:
    """
    Return the license string for MD GPU, given the number of
    GPUs to be used.
    """
    license_string = licensing.getFeatureName(
        licensing.DESMOND_ACADEMIC) if _is_academic(
        ) else licensing.getFeatureName(licensing.DESMOND_GPGPU)
    return f"{license_string}:{ngpus * DESMOND_GPU_TOKEN_MULTIPLIER}"


def md_cpu_lic(ncpus: int) -> str:
    """
    Return the license string for MD CPU, given the number of
    CPUs to be used.
    """
    license_string = licensing.getFeatureName(
        licensing.DESMOND_ACADEMIC) if _is_academic(
        ) else licensing.getFeatureName(licensing.DESMOND_MAIN)
    return f"{license_string}:{ncpus}"


def watermap_lic(ngpus: int) -> str:
    """
    Return the license string for Watermap GPU, given the number
    of GPUs to be used.
    """
    license_string = licensing.getFeatureName(licensing.DESMOND_WATERMAP_GPGPU)
    return f"{license_string}:{ngpus * DESMOND_GPU_TOKEN_MULTIPLIER}"


def watermap_cpu_lic(ncpus: int) -> str:
    """
    Return the license string for Watermap CPU, given the number
    of CPUs to be used.
    """
    license_string = licensing.getFeatureName(licensing.DESMOND_WATERMAP)
    return f"{license_string}:{ncpus * 2}"


def _get_num_procs(args: List[str]) -> int:
    """
    Return the number of processors (GPUs or CPUs) from command line
    arguments.
    """
    if '-HOST' in args:
        host = args[args.index('-HOST') + 1].split(":")
        if len(host) > 1:
            return int(host[1])
    if '-cpu' in args:
        num_procs = int(args[args.index('-cpu') + 1])
    else:
        num_procs = 1
    return num_procs


def _is_academic() -> bool:
    """
    Return True if this is an academic build. False otherwise.
    """
    return bld_ver.desmond_build_version() == bld_def.bld_types[
        bld_def.DESMOND_ACADEMIC]


def _deduplicate_lics(args: List[str]) -> List[str]:
    """
    Go through the command line arguments `args`
    and return the arguments without duplicate licenses.
    NOTE: If the same license is present with a different
    number of tokens, the one that appears last in the list
    will be used.
    """
    lic_map = {}
    updated_args = []
    i = 0
    while i < len(args):
        if args[i] == '-lic':
            lic_value = args[i + 1]
            lic_map[lic_value.split(':')[0]] = ['-lic', lic_value]
            # Skip -lic value
            i += 1
        else:
            updated_args.append(args[i])
        i += 1
    return sum(sorted(lic_map.values()), updated_args)