/*
 * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "luci/CircleOptimizer.h"

#include "luci/Pass/FuseInstanceNormPass.h"
// TODO add more passes

#include "luci/Pass/ShapeInferencePass.h"
#include "luci/Pass/TypeInferencePass.h"

// logo passes
#include <logo/RemoveDeadNodePass.h>

#include "ProgressReporter.h"

#include <logo/Phase.h>

#include <memory>

namespace
{

using namespace luci;

class OptimizeOptionsImpl : public luci::CircleOptimizer::Options
{
public:
  void enable(Algorithm) final;
  bool query(Algorithm) final;

private:
  std::vector<Algorithm> _algorithms;
};

void OptimizeOptionsImpl::enable(Algorithm algo) { _algorithms.push_back(algo); }

bool OptimizeOptionsImpl::query(Algorithm algo)
{
  std::vector<Algorithm>::iterator it = std::find(_algorithms.begin(), _algorithms.end(), algo);
  if (it == _algorithms.end())
    return false;

  return true;
}

} // namespace

namespace luci
{

CircleOptimizer::Options *CircleOptimizer::options(void)
{
  if (_options == nullptr)
  {
    _options = std::make_unique<OptimizeOptionsImpl>();
  }

  return _options.get();
}

void CircleOptimizer::optimize(loco::Graph *g) const
{
  logo::Phase phase;

  /* TRANSFORM DECLARATION BEGIN */
  if (_options->query(Options::Algorithm::FuseInstanceNorm))
  {
    phase.emplace_back(std::make_unique<FuseInstanceNormPass>());
  }
  // Shape inference is needed for added nodes doing above transformations
  phase.emplace_back(std::make_unique<luci::ShapeInferencePass>());
  phase.emplace_back(std::make_unique<luci::TypeInferencePass>());
  phase.emplace_back(std::make_unique<logo::RemoveDeadNodePass>());
  /* TRANSFORM DECLARATION END */

  ProgressReporter prog(g, logo::PhaseStrategy::Saturate);
  logo::PhaseRunner<logo::PhaseStrategy::Saturate> phase_runner{g};
  phase_runner.attach(&prog);
  phase_runner.run(phase);
}

} // namespace luci