/*
 (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 <cstdio>
#include <vector>

#include <openeye.h>

#include <oechem.h>
#include <oebioisostere.h>
#include <oeplatform.h>
#include <oesystem.h>

class BroodComboOptions : public OESystem::OEOptions
{
public:
  BroodComboOptions(std::string name = "BroodComboOptions")
    : OESystem::OEOptions(name)
  {
    OESystem::OEStringParameter dbParam("-db");
    dbParam.SetRequired(true);
    dbParam.SetVisibility(OESystem::OEParamVisibility::Simple);
    dbParam.SetBrief("Database folder");
    m_dbParam = AddParameter(dbParam);

    OESystem::OEUIntParameter maxHitsParam("-primaryQueryMaxHits", 15);
    maxHitsParam.SetVisibility(OESystem::OEParamVisibility::Simple);
    maxHitsParam.SetBrief("Maximum number of primary hits collected from the first query");
    m_maxHitsParam = AddParameter(maxHitsParam);

    OESystem::OEUIntParameter secondaryMaxHitsParam("-secondaryQueryMaxHits", 15);
    secondaryMaxHitsParam.SetVisibility(OESystem::OEParamVisibility::Simple);
    secondaryMaxHitsParam.SetBrief("Maximum number of secondary combo hits collected for each primary hit");
    m_secondaryMaxHitsParam = AddParameter(secondaryMaxHitsParam);

  }

  BroodComboOptions(const BroodComboOptions&) = default;
  BroodComboOptions& operator=(const BroodComboOptions&) = default;
  ~BroodComboOptions() override = default;
  BroodComboOptions* CreateCopy() const override { return new BroodComboOptions(*this); }

  std::string GetDataBase() const { return m_dbParam->GetStringValue();}
  unsigned int GetPrimaryQueryMaxHits() const
  {
    if (m_maxHitsParam->GetHasValue())
      return static_cast<unsigned int>(std::stoi(m_maxHitsParam->GetStringValue()));
    return static_cast<unsigned int>(std::stoi(m_maxHitsParam->GetStringDefault()));
  }
  unsigned int GetSecondaryQueryMaxHits() const
  {
    if (m_secondaryMaxHitsParam->GetHasValue())
      return static_cast<unsigned int>(std::stoi(m_secondaryMaxHitsParam->GetStringValue()));
    return static_cast<unsigned int>(std::stoi(m_secondaryMaxHitsParam->GetStringDefault()));
  }
private:
  OESystem::OEParameter* m_dbParam;
  OESystem::OEParameter* m_maxHitsParam;
  OESystem::OEParameter* m_secondaryMaxHitsParam;
};

int main(int argc, char* argv[])
{
  BroodComboOptions broodOpts;
  OEChem::OERefInputAppOptions opts(
      broodOpts,
      "BroodComboBuilder",
      OEChem::OEFileStringType::Mol3D,
      OEChem::OEFileStringType::Mol3D,
      OEChem::OEFileStringType::Mol3D,
      "-in2");

  if (OESystem::OEConfigureOpts(opts, argc, argv, false) == OESystem::OEOptsConfigureStatus::Help)
    return 0;
  broodOpts.UpdateValues(opts);

  OEChem::oemolistream ifs1;
  if (!ifs1.open(opts.GetInFile()))
    OESystem::OEThrow.Fatal("Unable to open %s for reading", opts.GetInFile().c_str());

  OEChem::oemolistream ifs2;
  if (!ifs2.open(opts.GetRefFile()))
    OESystem::OEThrow.Fatal("Unable to open %s for reading", opts.GetRefFile().c_str());

  OEChem::oemolostream ofs;
  if (!ofs.open(opts.GetOutFile()))
    OESystem::OEThrow.Fatal("Unable to open %s for writing", opts.GetOutFile().c_str());

  OEBioisostere::OEBroodQuery query1;
  unsigned retCode = OEBioisostere::OEReadBroodQuery(ifs1, query1);
  if (retCode != OEBioisostere::OEBroodStatusCode::Success)
    OESystem::OEThrow.Fatal("Failed to read -in1 query: %s", OEBioisostere::OEGetBroodStatus(retCode).c_str());

  OEBioisostere::OEBroodQuery query2;
  retCode = OEBioisostere::OEReadBroodQuery(ifs2, query2);
  if (retCode != OEBioisostere::OEBroodStatusCode::Success)
    OESystem::OEThrow.Fatal("Failed to read -in2 query: %s", OEBioisostere::OEGetBroodStatus(retCode).c_str());

  // Build primary hits from the first query.
  std::vector<OEBioisostere::OEBroodHit> vecPrimaryHits;
  OEBioisostere::OEDBReader primaryReader;
  if (primaryReader.Init(broodOpts.GetDataBase()) != OEBioisostere::OEBroodStatusCode::Success)
    OESystem::OEThrow.Fatal("Unable to open database for primary query hits");
  OEBioisostere::OEBroodOverlay primaryOverlay;
  primaryOverlay.SetupRef(query1);
  OEBioisostere::OEBroodMolBuilder primaryBuilder(query1);
  OEBioisostere::OEBroodDBPacket primaryPacket;
  while (primaryReader.GetNextPacket(primaryPacket) && vecPrimaryHits.size() < broodOpts.GetPrimaryQueryMaxHits())
  {
    const std::vector<OEBioisostere::OEBroodScore>& vecMatches = primaryOverlay.Overlay(primaryPacket);
    for (const auto& match : vecMatches)
    {
      OEBioisostere::OEBroodHit hit;
      if (primaryBuilder.Build(match, hit) == OEBioisostere::OEBroodStatusCode::Success)
      {
        vecPrimaryHits.push_back(hit);
        if (vecPrimaryHits.size() >= broodOpts.GetPrimaryQueryMaxHits())
          break;
      }
    }
  }
  if (vecPrimaryHits.empty())
    OESystem::OEThrow.Fatal("Unable to build primary hits from -in1");

  unsigned totalSecondaryHits = 0;
  unsigned written = 0;
  std::vector<OEBioisostere::OEBroodHit> allSecondaryHits;

  // The attempt budget below bounds Build() calls per primary hit
  // regardless of success rate.
  unsigned int maxBuildAttempts = broodOpts.GetSecondaryQueryMaxHits() * 20;

  // For each primary hit from query1, run query2 as the reference query in combo mode.
  // The per-primary cap keeps each combo expansion bounded by -secondaryQueryMaxHits.
  for (const auto& primaryHit : vecPrimaryHits)
  {
    OEBioisostere::OEDBReader secondaryReader;
    if (secondaryReader.Init(broodOpts.GetDataBase()) != OEBioisostere::OEBroodStatusCode::Success)
      OESystem::OEThrow.Fatal("Unable to build secondary combo hits from -in2");
    // Build an overlay pipeline for query2, then combine each score with the current
    // primary hit via OEBroodComboBuilder.
    OEBioisostere::OEBroodOverlay secondaryOverlay;
    secondaryOverlay.SetupRef(query2);
    OEBioisostere::OEBroodComboBuilder secondaryBuilder(query2, primaryHit);

    std::vector<OEBioisostere::OEBroodHit> vecSecondaryHits;
    OEBioisostere::OEBroodDBPacket secondaryPacket;
    unsigned int buildAttempts = 0;
    bool done = false;
    // Stream the database packets; stop once this primary hit reaches its cap
    // or exhausts the attempt budget.
    while (secondaryReader.GetNextPacket(secondaryPacket) && !done)
    {
      const std::vector<OEBioisostere::OEBroodScore>& vecMatches = secondaryOverlay.Overlay(secondaryPacket);
      for (const auto& match : vecMatches)
      {
        OEBioisostere::OEBroodHit hit;
        // Combo build succeeds only when the secondary replacement can be merged with
        // the primary context for the current query2 reference.
        if (secondaryBuilder.Build(match, hit) == OEBioisostere::OEBroodStatusCode::Success)
        {
          vecSecondaryHits.push_back(hit);
          if (vecSecondaryHits.size() >= broodOpts.GetSecondaryQueryMaxHits())
          {
            done = true;
            break;
          }
        }
        ++buildAttempts;
        if (buildAttempts >= maxBuildAttempts)
        {
          done = true;
          break;
        }
      }
    }
    allSecondaryHits.insert(allSecondaryHits.end(), vecSecondaryHits.begin(), vecSecondaryHits.end());
  }

  if (!allSecondaryHits.empty())
  {
    // Cluster and rank all secondary hits before writing output molecules.
    OEBioisostere::OEBroodClusterBuilder clBuilder(query1);
    if (!clBuilder.Add(allSecondaryHits))
      OESystem::OEThrow.Fatal("Unable to cluster secondary combo hits");
    if (!clBuilder.Rank())
      OESystem::OEThrow.Fatal("Unable to rank secondary combo clusters");

    const std::vector<OEBioisostere::OEBroodCluster>& vecClusters = clBuilder.GetClusters();

    for (const auto& cluster : vecClusters)
    {
      OEChem::OEWriteConstMolecule(ofs, cluster.GetHead().GetMol());
      ++written;
      ++totalSecondaryHits;
    }

    for (const auto& cluster : vecClusters)
    {
      for (const auto& hit : cluster.GetMembers())
      {
        OEChem::OEWriteConstMolecule(ofs, hit.GetMol());
        ++written;
        ++totalSecondaryHits;
      }
    }
  }

  printf("Primary hits: %u\n", static_cast<unsigned>(vecPrimaryHits.size()));
  printf("Secondary combo hits: %u\n", totalSecondaryHits);
  printf("Wrote %u molecules to %s\n", written, opts.GetOutFile().c_str());
  return 0;
}