#!/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.

#############################################################################
# Simple superimposition of a fit protein on to a reference protein
#############################################################################
import sys
import os
from openeye import oechem
from openeye import oespruce
import tempfile


def ReadProteinMol(ifilename):
    if oechem.OEIsReadableDesignUnit(ifilename):
        du = oechem.OEDesignUnit()
        if not oechem.OEReadDesignUnit(ifilename, du):
            oechem.OEThrow.Fatal("Unable to open %s for reading OEDesignUnit" % ifilename)
        return du
    else:
        # @ <SNIPPET-READPDB-CORRECTLY>
        ifs = oechem.oemolistream()
        ifs.SetFlavor(oechem.OEFormat_PDB, oechem.OEIFlavor_PDB_Default | oechem.OEIFlavor_PDB_DATA | oechem.OEIFlavor_PDB_ALTLOC)  # noqa        
        if not ifs.open(ifilename):
            oechem.OEThrow.Fatal("Unable to open %s for reading" % ifilename)
        mol = oechem.OEGraphMol()
        if not oechem.OEReadMolecule(ifs, mol):
            oechem.OEThrow.Fatal("Unable to read molecule from %s" % ifilename)
        # @ </SNIPPET-READPDB-CORRECTLY>
        return mol


def ReadSiteResidues(in_file):
    site_residues = []
    with open(in_file, "r") as f:
        lines = f.read().splitlines()

        for line in lines:
            if line.strip() == "":
                continue
            site_residues.append(line)
    return site_residues


def main(argv=[__name__]):
    if len(argv) < 3:
        oechem.OEThrow.Usage(f"{argv[0]} <reference protein PDB> <fit protein PDB> [global|ddm|weighted|sse|site] [site-residue file] [nowrite]")  # noqa

    inp_method = "global"
    if len(argv) > 3:
        inp_method = argv[3]

    site_file = None
    do_write = True
    if inp_method == "site":
        if len(argv) > 4:
            site_file = argv[4]
        else:
            oechem.OEThrow.Warning(f"A text file containing site residues must be provided for using the SiteSequence method\n")
            sys.exit(1)

        if not os.path.isfile(site_file):
            oechem.OEThrow.Warning(f"File not found: {site_file}\n")
            sys.exit(1)
        nowrite_argidx = 5
    else:
        nowrite_argidx = 4

    if len(argv) > nowrite_argidx:
        if argv[nowrite_argidx] != "nowrite":
            oechem.OEThrow.Warning(f"{argv[nowrite_argidx]} is not a valid option.\n")
            sys.exit(1)
        else:
            do_write = False

    ref_prot_file = argv[1]
    fit_prot_file = argv[2]

    ref_prot = ReadProteinMol(ref_prot_file)
    fit_prot = ReadProteinMol(fit_prot_file)

    method = oespruce.OEGetSuperposeMethodFromName(inp_method)
    if method == oespruce.OESuperposeMethod_Undefined:
        oechem.OEThrow.Warning(f"{inp_method} superposition method is not supported.\n")
        sys.exit(1)

    opts = oespruce.OESuperposeOptions(method)
    print(f"Superposing {fit_prot_file} to {ref_prot_file} using {oespruce.OEGetSuperposeMethodName(method)}.\n")

    results = oespruce.OESuperposeResults()
    superposition = oespruce.OESuperpose(opts)
    if opts.GetMethod() == oespruce.OESuperposeMethod_Site and not isinstance(ref_prot, oechem.OEDesignUnit):
        site_residues = ReadSiteResidues(site_file)
        superposition.SetupRef(ref_prot, site_residues)
    else:
        superposition.SetupRef(ref_prot)
    superposition.Superpose(results, fit_prot)

    rmsd = results.GetRMSD()
    seqscore = results.GetSeqScore()
    tanimoto = results.GetTanimoto()

    results.Transform(fit_prot)

    if opts.GetMethod() == oespruce.OESuperposeMethod_SSE:
        print(f"Tanimoto: {tanimoto:4.2f}\n")
    else:
        print(f"RMSD: {rmsd:4.2f} Angstroms.")
        print(f"SeqScore: {seqscore:d}.\n")

    if do_write:
        temp_dir = tempfile.mkdtemp()
        if isinstance(fit_prot, oechem.OEDesignUnit):
            str_pos = fit_prot_file.find(".oedu")
            base_name = fit_prot_file[0:str_pos]
            output_fit_file = os.path.join(temp_dir, base_name + "_sp.oedu")
            ofs = oechem.oeofstream(output_fit_file)
            oechem.OEWriteDesignUnit(ofs, fit_prot)
        else:
            str_pos = fit_prot_file.find(".pdb")
            base_name = fit_prot_file[0:str_pos]
            output_fit_file = os.path.join(temp_dir, base_name + "_sp.oeb.gz")
            ofs = oechem.oemolostream(output_fit_file)
            oechem.OEWriteMolecule(ofs, fit_prot)

        print(f"Superimposed fit protein was written to {output_fit_file}.\n")


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