#!/usr/bin/env python
# (C) 2022 Cadence Design Systems, Inc. (Cadence) 
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.

import os
import sys
from openeye import oechem, oedocking

COMBINED_DESIGNUNIT_CONSTANTS = {
    "All":oechem.OEDesignUnitComponents_All,
    "TargetComplex":oechem.OEDesignUnitComponents_TargetComplex,
    "TargetComplexNoSolvent":oechem.OEDesignUnitComponents_TargetComplexNoSolvent,
    "Default":oechem.OEDesignUnitComponents_Default,
    "ListComponents":oechem.OEDesignUnitComponents_ListComponents,
    "MacroMolComponents":oechem.OEDesignUnitComponents_MacroMolComponents,
    "MolComponents":oechem.OEDesignUnitComponents_MolComponents,
}

def make_receptor_with_custom_mask_and_predicate(du, receptor_out, mask, pred_str):
    """
    Inputs:
        du (str): The file containing OEDesignUnit (.oedu)
        receptor_out (str): Path for the output receptor file
        pred (str): Pipe separated string of residues to create subset. 
        If no predicate is provided, the entirety of an included 
        component will be part of the receptor mask.
        mask(OEDesignUnitComponent): Specifies the component of the design unit
        that is to be used as as receptor mask. 

    """

    opts = oedocking.OEMakeReceptorOptions()

    # Add target predicate to create OESubSetDesignUnit
    opts.SetTargetPred(pred_str)
    target_mask = mask
    opts.SetTargetMask(target_mask)

    if oedocking.OEMakeReceptor(du, opts):
        oechem.OEThrow.Info(f"Successfully created receptor: {receptor_out}")
        return True
    
    oechem.OEThrow.Warning(f"{du.GetTitle()} failed to make receptor")
    return False

def get_target_component_mask(component_string):

    if component_string in COMBINED_DESIGNUNIT_CONSTANTS:
        return COMBINED_DESIGNUNIT_CONSTANTS[component_string]
    
    component_id = oechem.OEGetDesignUnitComponentID(component_string)

    if component_id == 0:
        oechem.Throw.Fatal(f"{component_string} is not a OEDesignUnit Component")
    
    return component_id

def main(argv=[__name__]):
    """
    Add solvent or cofactor atoms/molecules to a receptor docking grid.
    """
    itf = oechem.OEInterface(InterfaceData, argv)
    infile = itf.GetString("-du")
    receptor_out = itf.GetString("-out")
    target_mask = itf.GetString("-mask")
    pred_str = itf.GetString("-pred")

    if not os.path.exists(infile):
        oechem.OEThrow.Fatal(f"{infile} does not exist.")

    du = oechem.OEDesignUnit()
    if not oechem.OEReadDesignUnit(infile, du):
        oechem.OEThrow.Fatal("Cannot read design unit!")

    if not oechem.OEIsWriteableDesignUnit(receptor_out):
        oechem.OEThrow.Fatal(f"Can not write design unit to {receptor_out}")

    target_mask = target_mask.split(",")
    if len(target_mask) == 0:
        processed_mask = oechem.OEDesignUnitComponents_TargetComplexNoSolvent
    elif len(target_mask) == 1:
        processed_mask = get_target_component_mask(target_mask[0].strip())
    else:
        processed_mask = get_target_component_mask(target_mask[0].strip())
        for component_string in target_mask[1:]:
            print("hit it")
            processed_mask = processed_mask | get_target_component_mask(component_string.strip())

    success = make_receptor_with_custom_mask_and_predicate(du, receptor_out, processed_mask, pred_str)

    if success:
        oechem.OEWriteDesignUnit(receptor_out, du)

InterfaceData = """
!BRIEF [To run] python add_receptor_obj_to_du.py -du <oedu> 

!PARAMETER -du
    !TYPE string
    !REQUIRED true
    !BRIEF Input OEDesignUnit file.
    !DETAIL
        Input OEDesignUnit (.oedu) file containing the prepared protein structure.
!END

!PARAMETER -out
    !TYPE string
    !REQUIRED false
    !DEFAULT receptor_out.oedu
    !BRIEF Output receptor file.
    !DETAIL
        Output receptor file (.oedu).
!END

!PARAMETER -mask
    !TYPE string
    !REQUIRED false
    !DEFAULT TargetComplexNoSolvent
    !BRIEF Mask string containing the component names for the receptor mask. Components should be separated by string.
    !DETAIL
        Mask string containing the component names for the receptor mask. Components should be separated by string.
        Example: "protein,ligand,cofactors". Options to choose from: protein, nucleic, ligand, cofactors, solvent, 
        metals, counter_ions, lipids, packing_residues, excipients, suagrs, polymers, post_translational, other_proteins, 
        other_nucleics, other_ligands, other_cofactors, undefined, All, TargetComplex, TargetComplexNoSolvent, Default,
        ListComponents, MacroMolComponents, and MolComponents.
!END

!PARAMETER -pred
    !TYPE string
    !REQUIRED false
    !BRIEF Predicate string containing residue IDs, solvent or cofactor molecules to create a subset
    !DETAIL
        Predicate string containing solvent and cofactor molecules with their residue IDs. Accepted format is, 
        "MN:508: :D:1:|HOH:435: :D:1:"
        where, if needed, empty space is filled with alt loc.
!END

"""

if __name__ == "__main__":
    sys.exit(main(sys.argv))