/* 
(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.
*/

#include <openeye.h>

#include <oeplatform.h>
#include <oesystem.h>
#include <oechem.h>
#include <oebio.h>
#include <oespruce.h>

#include <vector>
#include <fstream>

using namespace OEPlatform;
using namespace OESystem;
using namespace OEChem;
using namespace OESpruce;
using namespace OEBio;
using namespace std;

//@ <SNIPPET-READPDB-CORRECTLY>
bool ReadProteinMol(const string& pdbFile, OEChem::OEMolBase& mol)
{
  const unsigned int pdbFlavor = OEIFlavor::PDB::SpruceDefault;
  const unsigned int mmcifFlavor = OEIFlavor::MMCIF::SpruceDefault;
  oemolistream ifs;
  ifs.SetFlavor(OEFormat::PDB, pdbFlavor);
  ifs.SetFlavor(OEFormat::MMCIF, mmcifFlavor);

  if (!ifs.open(pdbFile))
  {
    OEThrow.Warning("Unable to open %s for reading.", pdbFile.c_str());
    return false;
  }

  if (!OEReadMolecule(ifs, mol))
  {
    OEThrow.Warning("Unable to read molecule from %s.", pdbFile.c_str());
    return false;
  }
  ifs.close();

  return (mol);
}
//@ </SNIPPET-READPDB-CORRECTLY>

bool ReadDesignUnit(const string& duFile, OEBio::OEDesignUnit& du) {

  if (!OEReadDesignUnit(duFile, du)) {
    OEThrow.Warning("Unable to read molecule from %s.", duFile.c_str());
    return false;
  }

  return (du);
}

vector<string> ReadTextFile(const string& filename)
{
  ifstream file(filename);
  string str;
  vector<string> file_contents;
  while (std::getline(file, str))
  {
    file_contents.push_back(str);
  }

  return file_contents;
}

void ShowUsage()
{
  OEThrow.Usage("superpose <reference protein PDB> <fit protein PDB> [global|ddm|weighted|site|sse] [site-residue file]");
}

int main(int argc, char *argv[])
{
  if (argc < 3)
  {
    ShowUsage();
    return 1;
  }

  string inpMethod = "global"; 
  string siteFile;

  if (argc >= 4)
  {
    vector<string> allowedMethods;
    allowedMethods.push_back("global");
    allowedMethods.push_back("ddm");
    allowedMethods.push_back("weighted");
    allowedMethods.push_back("site");
    allowedMethods.push_back("sse");

    inpMethod = string(argv[3]);
    if(find(allowedMethods.begin(), allowedMethods.end(), inpMethod) == allowedMethods.end())
    {
      OEThrow.Warning("%s superposition method is not supported.\n", inpMethod.c_str());
      ShowUsage();
      return 1;
    }
  }

  if (argc >= 5)
    siteFile = string(argv[4]);

  string refProtFile(argv[1]);
  string fitProtFile(argv[2]);

  unsigned method = OEGetSuperposeMethodFromName(inpMethod);
  OESuperposeOptions spOpts(method);
  OESuperposeResults results;
  OESuperpose superposition(spOpts);

  cout << "Superposing " << fitProtFile << " to " << refProtFile 
       << " using " << inpMethod << endl;

  if (strcmp(OEGetFileExtension(refProtFile.c_str()), "oedu") == 0) 
  {
    OEDesignUnit refDU;
    OEDesignUnit fitDU;
    if (!ReadDesignUnit(refProtFile, refDU) || !ReadDesignUnit(fitProtFile, fitDU))
      OEThrow.Fatal("Unable to read protein(s) from OEDU file.");

    superposition.SetupRef(refDU);
    superposition.Superpose(results, fitDU);

    string oeduExt(".oedu");
    size_t strPos = fitProtFile.find(oeduExt);
    string baseName = fitProtFile.substr(0, strPos);
    string outputFitFile = baseName + string("_sp.oedu");
    cout << "Writing superimposed fit protein to " << outputFitFile << endl;

    results.Transform(fitDU);

    oeofstream ofs(outputFitFile);
    OEWriteDesignUnit(ofs, fitDU);
  } 
  else 
  {
    OEGraphMol refProt;
    OEGraphMol fitProt;
    if (!ReadProteinMol(refProtFile, refProt) || !ReadProteinMol(fitProtFile, fitProt))
      OEThrow.Fatal("Unable to read protein(s) from PDB file.");

    if (method == OESuperposeMethod::Site)
    {
      superposition.SetupRef(refProt);
    }
    else
    {
      vector<string> sites = ReadTextFile(siteFile);
      superposition.SetupRef(refProt, sites);
    }
    superposition.Superpose(results, fitProt);

    string pdbExt(".pdb");
    size_t strPos = fitProtFile.find(pdbExt);
    string baseName = fitProtFile.substr(0, strPos);
    string outputFitFile = baseName + string("_sp.oeb.gz");
    cout << "Writing superimposed fit protein to " << outputFitFile << endl;

    results.Transform(fitProt);

    oemolostream ofs(outputFitFile);
    OEWriteMolecule(ofs, fitProt);
  }

  double rmsd = results.GetRMSD();
  int seqscore = results.GetSeqScore();
  double tanimoto = results.GetTanimoto();

  if (method == OESuperposeMethod::SSE)
  {
    cout << "Tanimoto: " << tanimoto << endl;
  }
  else
  {
    cout << "RMSD: " << rmsd << " Angstroms" << endl;
    cout << "SeqScore: " << seqscore << endl;
  }

  return 0;
}