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

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

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

    m_GeneralOptions = static_cast<OEBioisostere::OEBroodGeneralOptions*>(AddOption(OEBioisostere::OEBroodGeneralOptions()));
    m_ScoreOptions = static_cast<OEBioisostere::OEBroodScoreOptions*>(AddOption(OEBioisostere::OEBroodScoreOptions()));
    m_HitlistOptions = static_cast<OEBioisostere::OEBroodHitlistOptions*>(AddOption(OEBioisostere::OEBroodHitlistOptions()));
  }

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

  std::string GetDataBase() const { return m_dbParam->GetStringValue(); }
  OEBioisostere::OEBroodGeneralOptions* GetGenOpts() const { return m_GeneralOptions; }
  OEBioisostere::OEBroodScoreOptions* GetScoreOpts() const { return m_ScoreOptions; }
  OEBioisostere::OEBroodHitlistOptions* GetHitlistOpts() const { return m_HitlistOptions; }

private:
  OESystem::OEParameter* m_dbParam;
  OEBioisostere::OEBroodGeneralOptions* m_GeneralOptions;
  OEBioisostere::OEBroodScoreOptions* m_ScoreOptions;
  OEBioisostere::OEBroodHitlistOptions* m_HitlistOptions;
};

int main(int argc, char* argv[])
{
  BroodOptions broodOpts;
  OEChem::OESimpleAppOptions opts(broodOpts, "BroodCluster", OEChem::OEFileStringType::Mol3D, OEChem::OEFileStringType::Mol3D);

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

  OEChem::oemolistream ifs;
  if (!ifs.open(opts.GetInFile()))
    OESystem::OEThrow.Fatal("Unable to open %s for reading", opts.GetInFile().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 query;
  unsigned retCode = OEBioisostere::OEReadBroodQuery(ifs, query);
  if (retCode != OEBioisostere::OEBroodStatusCode::Success)
    OESystem::OEThrow.Fatal("Failed to read query: %s", OEBioisostere::OEGetBroodStatus(retCode).c_str());

  OEBioisostere::OEDBReader reader;
  unsigned retValue = reader.Init(broodOpts.GetDataBase(), query, *broodOpts.GetGenOpts());
  if (retValue != OEBioisostere::OEBroodStatusCode::Success)
    OESystem::OEThrow.Fatal("Unable to load Brood database");

  OEBioisostere::OEBroodOverlay overlay(*broodOpts.GetGenOpts(), *broodOpts.GetScoreOpts());
  overlay.SetupRef(query);

  OEBioisostere::OEHitlistBuilder hitlist(query, *broodOpts.GetGenOpts(), *broodOpts.GetHitlistOpts());

  unsigned packetCount = 0;
  OEBioisostere::OEBroodDBPacket packet;
  while (reader.GetNextPacket(packet))
  {
    ++packetCount;
    printf("Processing packet %u with %u fragments\n", packetCount, packet.GetFragCount());
    const std::vector<OEBioisostere::OEBroodScore>& vecScores = overlay.Overlay(packet);
    hitlist.AddScores(vecScores);
  }

  hitlist.Build();
  std::vector<OEBioisostere::OEBroodHit> vecHits = hitlist.GetHits();
  printf("Initial hit count: %u\n", static_cast<unsigned>(vecHits.size()));

  OEBioisostere::OEBroodClusterBuilder clBuilder(query);
  if (!clBuilder.Add(vecHits) || !clBuilder.Rank())
    OESystem::OEThrow.Fatal("Unable to build hit clusters");

  const std::vector<OEBioisostere::OEBroodCluster>& vecClusters = clBuilder.GetClusters();
  printf("Number of clusters: %u\n", static_cast<unsigned>(vecClusters.size()));

  unsigned written = 0;
  for (const auto& cluster : vecClusters)
  {
    const unsigned rank = cluster.GetRank();
    printf("Cluster %u: head=%s members=%u interest=%.3f\n",
           rank,
           cluster.GetHead().GetFragSmiles().c_str(),
           cluster.Count(),
           cluster.GetInterest());
  }

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

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

  printf("Wrote %u clustered molecules to %s\n", written, opts.GetOutFile().c_str());
  return 0;
}