/**
 * This file is part of sally.
 * Copyright (C) 2015 SRI International.
 *
 * Sally is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Sally is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with sally.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "engine/pdkind/pdkind_engine.h"
#include "engine/pdkind/solvers.h"
#include "engine/factory.h"

#include "smt/factory.h"
#include "utils/statistics.h"
#include "utils/trace.h"
#include "expr/gc_relocator.h"

#include <stack>
#include <cassert>
#include <sstream>
#include <iostream>
#include <fstream>
#include <algorithm>

#include "system/trace_helper.h"

#define unused_var(x) { (void)x; }

namespace sally {
namespace pdkind {

pdkind_engine::pdkind_engine(const system::context& ctx)
: engine(ctx)
, d_transition_system(0)
, d_property(0)
, d_trace(0)
, d_invariant(expr::term_ref(), 0)
, d_smt(0)
, d_reachability(ctx, d_cex_manager)
, d_induction_frame_index(0)
, d_induction_frame_depth(0)
, d_induction_frame_depth_count(0)
, d_induction_frame_next_index(0)
, d_property_invalid(false)
, d_learning_type(LEARN_UNDEFINED)
{
  utils::statistics& stats = ctx.get_statistics();
  d_stats.frame_index = static_cast<utils::stat_int*>(stats.register_stat("pdkind::frame_index"));
  d_stats.induction_depth = static_cast<utils::stat_int*>(stats.register_stat("pdkind::induction_depth"));
  d_stats.frame_size = static_cast<utils::stat_int*>(stats.register_stat("pdkind::frame_size"));
  d_stats.frame_pushed = static_cast<utils::stat_int*>(stats.register_stat("pdkind::frame_pushed"));
  d_stats.queue_size = static_cast<utils::stat_int*>(stats.register_stat("pdkind::queue_size"));
  d_stats.max_cex_depth = static_cast<utils::stat_int*>(stats.register_stat("pdkind::max_cex_depth"));
}

pdkind_engine::~pdkind_engine() {
  delete d_smt;
}

void pdkind_engine::reset() {
  d_transition_system = 0;
  d_property = 0;
  d_trace = 0;
  d_invariant = engine::invariant(expr::term_ref(), 0);
  d_induction_frame.clear();
  d_induction_frame_index = 0;
  d_induction_frame_depth = 0;
  d_induction_frame_depth_count = 0;
  d_induction_obligations.clear();
  d_induction_obligations_next.clear();
  d_induction_obligations_count.clear();
  delete d_smt;
  d_smt = 0;
  d_properties.clear();
  d_property_invalid = false;
  d_reachability.clear();
  d_cex_manager.clear();
}

induction_obligation pdkind_engine::pop_induction_obligation() {
  assert(d_induction_obligations.size() > 0);
  induction_obligation ind = d_induction_obligations.top();
  d_induction_obligations.pop();
  d_induction_obligations_handles.erase(ind);
  d_stats.queue_size->set_value(d_induction_obligations.size());
  return ind;
}

void pdkind_engine::enqueue_induction_obligation(const induction_obligation& ind) {
  assert(d_induction_frame.find(ind) != d_induction_frame.end());
  induction_obligation_queue::handle_type h = d_induction_obligations.push(ind);
  d_induction_obligations_handles[ind] = h;
  d_stats.queue_size->set_value(d_induction_obligations.size());
}

pdkind_engine::induction_result pdkind_engine::push_obligation(induction_obligation& ind) {

  TRACE("pdkind") << "pdkind: Trying F_fwd at " << d_induction_frame_index << ": " << ind.F_fwd << std::endl;

  // Check if F_cex is inductive. If not then, if it can be reached, we can find a counter-example.
  solvers::query_result fwd_result = d_smt->check_inductive(ind.F_fwd);

  // If UNSAT we can push it
  if (fwd_result.result == smt::solver::UNSAT) {
    // It's pushed so add it to induction assumptions
    assert(d_induction_frame.find(ind) != d_induction_frame.end());
    TRACE("pdkind") << "pdkind: pushed " << ind.F_fwd << std::endl;
    // Add it to set of pushed facts
    d_induction_obligations_next.push_back(ind);
    d_stats.frame_pushed->set_value(d_induction_obligations_next.size());
    // We're done
    return INDUCTION_SUCCESS;
  }

  expr::term_ref F_fwd_not = tm().mk_not(ind.F_fwd);
  expr::term_ref F_cex_not = tm().mk_not(ind.F_cex);
  F_cex_not = tm().mk_or(F_cex_not);

  // We have a model for
  //
  //   F F ..... F !F
  //
  // If the model satisfies CEX in the last state, we try to extend the
  // counter-example, otherwise, we fail the push.
  solvers::query_result cex_result = d_smt->check_inductive_model(fwd_result.model, ind.F_cex);
  if (cex_result.result == smt::solver::UNSAT) {
    cex_result = d_smt->check_inductive(F_cex_not);
  }
  if (cex_result.result == smt::solver::SAT) {
    // We can actually reach the counterexample of induction from G, so we check if
    // it's reachable.
    TRACE("pdkind") << "pdkind: F_cex generalization " << cex_result.generalization << std::endl;

    // Check if G is reachable. We know that F_cex is not reachable up to induction frame index.
    // This means that G can be reached at index i, then F_cex is reachable at i + induction_depth.
    // Therefore i + induction_depth > frame_index, hence we check from i = frame_index - depth + 1 to frame_index
    size_t start = (d_induction_frame_index + 1) - d_induction_frame_depth;
    size_t end = d_induction_frame_index;
    reachability::status reachable = d_reachability.check_reachable(start, end, cex_result.generalization, 0);

    // If reachable, then G leads to F_cex, and F_cex leads to !P
    if (reachable.r == reachability::REACHABLE) {
      d_property_invalid = true;
      d_cex_manager.add_edge(cex_result.generalization, ind.F_cex, d_induction_frame_depth, 0);
      MSG(1) << "pdkind: found counter-example of length " << d_cex_manager.get_cex_length(0) + 1 << std::endl;
      return INDUCTION_FAIL;
    }

    // G is not reachable, so !G holds up to current frame index
    expr::term_ref F_cex = cex_result.generalization;
    TRACE("pdkind") << "pdkind: new F_cex: " << F_cex << std::endl;

    // Learn something forward that refutes G
    // We know that G_not is not satisfiable
    expr::term_ref F_fwd = d_smt->learn_forward(d_induction_frame_index, F_cex);
    TRACE("pdkind") << "pdkind: new F_fwd: " << F_fwd << std::endl;

    // Add to counter-example to induction frame
    induction_obligation new_ind(tm(), F_fwd, F_cex, d_induction_frame_depth + ind.d, 1);
    assert(d_induction_frame.find(new_ind) == d_induction_frame.end());
    d_induction_frame.insert(new_ind);
    d_stats.frame_size->set_value(d_induction_frame.size());
    d_smt->add_to_induction_solver(F_fwd, solvers::INDUCTION_FIRST);
    d_smt->add_to_induction_solver(F_fwd, solvers::INDUCTION_INTERMEDIATE);
    enqueue_induction_obligation(new_ind);

    // Remember the counter-example
    d_cex_manager.add_edge(F_cex, ind.F_cex, d_induction_frame_depth, 0);

    // We have to learn :(, decrease the score, let others go for a change
    ind.score *= 0.9;

    // Now, retry
    return INDUCTION_RETRY;
  }

  // So we have a model of !F_fwf && !F_cex
  solvers::query_result cti_result = d_smt->check_inductive_model(fwd_result.model, F_fwd_not);
  assert(cti_result.result == smt::solver::SAT);

  // Check if it's reachable
  size_t start = (d_induction_frame_index + 1) - d_induction_frame_depth;
  size_t end = d_induction_frame_index;
  reachability::status reachable = d_reachability.check_reachable(start, end, cti_result.generalization, cex_manager::null_property_id);

  if (reachable.r == reachability::REACHABLE) {
    // Current obligation is false at reach + depth. So we can move to at most
    // reach + k next time;
    //
    // This particular CTI is reachable in k + depth. There might be others that
    // are shorter than this.

    // Find out the smallest index where F_fwd is falsified
    start = d_induction_frame_index + 1;
    end = std::min(reachable.k + d_induction_frame_depth, d_induction_frame_next_index);
    if (start < end) {
      reachable = d_reachability.check_reachable(start, end, F_fwd_not, cex_manager::null_property_id);
      if (reachable.r == reachability::REACHABLE) {
        // We have to adapt the next index
        d_induction_frame_next_index = reachable.k;
      }
    } else {
      d_induction_frame_next_index = end;
    }

    // We know that CEX is not reachable (FULL CHECK ABOVE), otherwise we wouldn't be here.
    // Therefore !CEX is k-inductive modulo current assumptions and we can just push it
    induction_obligation new_ind(tm(), F_cex_not, ind.F_cex, ind.d, ind.d);
    assert(d_induction_frame.find(new_ind) == d_induction_frame.end());
    d_induction_frame.insert(new_ind);
    d_stats.frame_size->set_value(d_induction_frame.size());
    // No need to assert anything, we already have F_fwd => !F_cex
    // Also, just add to next
    d_induction_obligations_next.push_back(new_ind);
    d_stats.frame_pushed->set_value(d_induction_obligations_next.size());

    // Current obligation has failed, we know it will become invalid
    return INDUCTION_FAIL;
  }

  // Learn something forward that refutes reachability of CTI
  expr::term_ref F_fwd = d_smt->learn_forward(d_induction_frame_index, cti_result.generalization);
  d_smt->add_to_induction_solver(F_fwd, solvers::INDUCTION_FIRST);
  d_smt->add_to_induction_solver(F_fwd, solvers::INDUCTION_INTERMEDIATE);

  F_fwd = tm().mk_and(ind.F_fwd, F_fwd);
  TRACE("pdkind") << "pdkind: new F_fwd: " << F_fwd << std::endl;

  // We know what F_fwd removes the CTI, and F_fwd => !CEX, so we can add it
  d_induction_frame.erase(ind);
  ind = induction_obligation(tm(), F_fwd, ind.F_cex, ind.d, ind.score*0.9, ind.refined + 1);
  d_induction_frame.insert(ind);

  // Current obligation has failed, but we replaced it
  return INDUCTION_RETRY;
}

void pdkind_engine::push_current_frame() {

  // Search while we have something to do
  while (!d_induction_obligations.empty() && !d_property_invalid) {

    // Pick a formula to try and prove inductive, i.e. that F_k & P & T => P'
    induction_obligation ind = pop_induction_obligation();

    // Push the formula forward if it's inductive at the frame
    induction_result ind_result = push_obligation(ind);

    // See what happened
    switch (ind_result) {
    case INDUCTION_RETRY:
      // We'll retry the same formula (it's already added to the solver)
      enqueue_induction_obligation(ind);
      break;
    case INDUCTION_SUCCESS:
      // Boss, we're done with this one
      break;
    case INDUCTION_FAIL:
      // Failure, we didn't push, either counter-example found, or we couldn't push
      // and decided not to try again
      break;
    }
  }
}

engine::result pdkind_engine::search() {

  // Push frame by frame */
  for(;;) {

    // d_induction_cutoff = d_induction_frame_index + d_induction_frame_depth;

    // Set how far we can go
    // d_induction_frame_next_index = d_induction_frame_index + 1;
    d_induction_frame_next_index = d_induction_frame_index + d_induction_frame_depth;

    MSG(1) << "pdkind: working on induction frame " << d_induction_frame_index << " (" << d_induction_frame.size() << ") with induction depth " << d_induction_frame_depth << std::endl;

    // Push the current induction frame forward
    push_current_frame();

    // If we've disproved the property, we're done
    if (d_property_invalid) {
      return engine::INVALID;
    }

    MSG(1) << "pdkind: pushed " << d_induction_obligations_next.size() << " of " << d_induction_frame.size() << std::endl;

    // If we pushed everything, we're done
    if (d_induction_frame.size() == d_induction_obligations_next.size()) {
      std::set<expr::term_ref> invariant;
      induction_frame_type::const_iterator it = d_induction_frame.begin(), end = d_induction_frame.end();
      for (; it != end; ++ it) { invariant.insert(it->F_fwd); }
      d_invariant = engine::invariant(tm().mk_and(invariant), d_induction_frame_depth);
      return engine::VALID;
    }

    // Set depth of induction for next time
    d_induction_frame_depth ++;

    // Clear induction obligations queue and the frame
    d_induction_obligations.clear();
    d_induction_frame.clear();
    d_stats.frame_size->set_value(0);

    // If exceeded number of frames
    if (ctx().get_options().get_unsigned("pdkind-max") > 0 && d_induction_frame_index >= ctx().get_options().get_unsigned("pdkind-max")) {
      return engine::INTERRUPTED;
    }

    // Next frame position
    d_induction_frame_index = d_induction_frame_next_index;

    if (ctx().get_options().get_unsigned("pdkind-induction-max") != 0 && d_induction_frame_depth > ctx().get_options().get_unsigned("pdkind-induction-max")) {
      d_induction_frame_depth = ctx().get_options().get_unsigned("pdkind-induction-max");
    }
    d_smt->reset_induction_solver(d_induction_frame_depth);

    if (ctx().get_options().get_bool("pdkind-minimize-frames")) {
      d_smt->minimize_frame(d_induction_obligations_next);
    }

    // Add formulas to the new frame
    d_induction_frame.clear();
    std::vector<induction_obligation>::const_iterator next_it = d_induction_obligations_next.begin();
    for (; next_it != d_induction_obligations_next.end(); ++ next_it) {
      // The formula
      induction_obligation ind = *next_it;
      ind.score = ind.score/2 + 0.5; // Keep old score and add 1 for effort
      assert(d_induction_frame.find(ind) == d_induction_frame.end());
      d_smt->add_to_induction_solver(ind.F_fwd, solvers::INDUCTION_FIRST);
      d_smt->add_to_induction_solver(ind.F_fwd, solvers::INDUCTION_INTERMEDIATE);
      d_induction_frame.insert(ind);
      d_stats.frame_size->set_value(d_induction_frame.size());
      enqueue_induction_obligation(ind);
    }

    // Clear next frame info
    d_induction_obligations_next.clear();
    d_stats.frame_pushed->set_value(0);

    // Update stats
    d_stats.frame_index->set_value(d_induction_frame_index);
    d_stats.induction_depth->set_value(d_induction_frame_depth);

    // Do garbage collection
    d_smt->gc();
  }

  // Didn't prove or disprove, so unknown
  return engine::UNKNOWN;
}

engine::result pdkind_engine::query(const system::transition_system* ts, const system::state_formula* sf) {

  // Initialize
  result r = UNKNOWN;
  d_last_result = r;

  // Reset the engine
  reset();

  // Remember the input
  d_transition_system = ts;
  d_property = sf;

  // Make the trace
  d_trace = ts->get_trace_helper();
  d_trace->clear_model();

  // Initialize the solvers
  if (d_smt) { delete d_smt; }
  d_smt = new solvers(ctx(), ts, d_trace);

  // Initialize the reachability solver
  d_reachability.init(d_transition_system, d_smt);

  // Initialize the induction solver
  d_induction_frame_index = 0;
  d_induction_frame_depth = 1;
  d_smt->reset_induction_solver(1);

  // Add the property we're trying to prove (if not already invalid at frame 0)
  bool ok = add_property(d_property->get_formula());
  if (!ok) {
#ifndef NDEBUG
    // Check trace generation if not asked for explicityly
    if (!ctx().get_options().has_option("show-trace")) { get_trace(); }  
#endif
    return engine::INVALID;
  }

  while (r == UNKNOWN) {

    MSG(1) << "pdkind: starting search" << std::endl;

    // Search
    r = search();
    d_last_result = r;
  }

  MSG(1) << "pdkind: search done: " << r << std::endl;

  // Print cex graph if asked
  if (ctx().get_options().has_option("pdkind-output-cex-graph")) {
    std::string filename = ctx().get_options().get_string("pdkind-output-cex-graph");
    std::ofstream cex_out(filename.c_str());
    cex_out << expr::set_tm(tm()) << d_cex_manager;
  }

#ifndef NDEBUG
  // Check trace generation if not asked for explicityly
  if (r == engine::INVALID && !ctx().get_options().has_option("show-trace")) { get_trace(); }
#endif

  return r;
}

bool pdkind_engine::add_property(expr::term_ref P) {
  // Add to cex manager
  expr::term_ref P_cex = tm().mk_not(P);
  d_cex_manager.add_edge(P_cex, P_cex, 0, 0);

  // Check if sat in first frame
  smt::solver::result result = d_smt->query_at_init(tm().mk_not(P));
  if (result == smt::solver::UNSAT) {
    induction_obligation ind(tm(), P, P_cex, /* cex depth */ 0, /* score */ 1);
    if (d_induction_frame.find(ind) == d_induction_frame.end()) {
      // Add to induction frame, we know it holds at 0
      assert(d_induction_frame_depth == 1);
      d_induction_frame.insert(ind);
      d_stats.frame_size->set_value(d_induction_frame.size());
      d_smt->add_to_induction_solver(P, solvers::INDUCTION_FIRST);
      enqueue_induction_obligation(ind);
    }
    d_properties.insert(P);
    return true;
  } else {
    d_cex_manager.mark_root(P_cex, 0);
    return false;
  }
}

const system::trace_helper* pdkind_engine::get_trace() {

  MSG(1) << "pdkind: constructing counter-example" << std::endl;

  size_t property_id = 0;

  // Get the cex from the cex graph
  cex_manager::edge_vector cex_edges;
  expr::term_ref cex_start = d_cex_manager.get_full_cex(property_id, cex_edges);
  assert(!cex_start.is_null()); // Returns null if no path to property

  // Compute the size of the counter-example
  size_t cex_length = 0;
  for (size_t i = 0; i < cex_edges.size(); ++ i) { cex_length += cex_edges[i].edge_length; }

  // Construct the solver
  smt::solver* solver = smt::factory::mk_default_solver(tm(), ctx().get_options(), ctx().get_statistics());

  // Variables that are already assigned
  std::vector<expr::term_ref> assigned_variables;

  // Add all variables
  for (size_t k = 0; k <= cex_length; ++ k) {
    const std::vector<expr::term_ref>& x = d_trace->get_state_variables(k);
    solver->add_variables(x.begin(), x.end(), smt::solver::CLASS_A);
    if (k < cex_length) {
      const std::vector<expr::term_ref>& i = d_trace->get_input_variables(k);
      solver->add_variables(i.begin(), i.end(), smt::solver::CLASS_A);  
    }
  }

  // Scope for push/pop on the solver
  smt::solver_scope scope(solver);

  // Get the initial model at s0 
  scope.push();
  expr::term_ref I = d_transition_system->get_initial_states(); 
  I = d_trace->get_state_formula(I, 0);
  solver->add(I, smt::solver::CLASS_A);
  cex_start = d_trace->get_state_formula(cex_start, 0);
  TRACE("pdkind::cex") << "Starting from " << cex_start << std::endl;
  solver->add(cex_start, smt::solver::CLASS_A);
  smt::solver::result res = solver->check();
  (void)res;
  assert(res == smt::solver::SAT);
  expr::model::ref model = solver->get_model();
  TRACE("pdkind::cex") << "model = " << *model << std::endl;
  scope.pop();

  // Add model to trace
  d_trace->set_model(model, 0, 0);

  // Construct the counter-example
  size_t current_depth = 0;
  for (size_t i = 0; i < cex_edges.size(); ++ i) {
    
    // Next formula we are trying to reach
    expr::term_ref cex_next = cex_edges[i].B;
    // The steps we are taking
    size_t cex_step = cex_edges[i].edge_length;
    assert(cex_step > 0);

    TRACE("pdkind::cex") << "at " << current_depth << ", step = " << cex_step << std::endl;

    // Push the solver scope
    scope.push();

    // Assert the previous model as the starting point
    if (solver->supports(smt::solver::SMT_MODULO_MODELS)) {
      assigned_variables.clear();
      d_trace->add_vars(current_depth, current_depth, assigned_variables);
    } else {
      d_trace->add_model_to_solver(model, current_depth, current_depth, solver, smt::solver::CLASS_A);
    }

    // Add the transition relation
    expr::term_ref T = d_transition_system->get_transition_relation(); 
    for (size_t k = current_depth; k < current_depth + cex_step; ++ k) {
      expr::term_ref T_k = d_trace->get_transition_formula(T, k);
      solver->add(T_k, smt::solver::CLASS_A);
    }

    // Add the goal to reach 
    TRACE("pdkind::cex") << "cex_next = " << cex_next << std::endl;
    cex_next = d_trace->get_state_formula(cex_next, current_depth + cex_step);
    TRACE("pdkind::cex") << "cex_next at " << current_depth + cex_step << " = " << cex_next << std::endl;
    solver->add(cex_next, smt::solver::CLASS_A);

    // Check for satisfiability (it must be SAT)
    smt::solver::result res;
    if (solver->supports(smt::solver::SMT_MODULO_MODELS)) {
      res = solver->check(model, assigned_variables);
    } else {
      res = solver->check();
    }
    (void)res; // Unused
    assert(res == smt::solver::SAT);

    // Get the model and add it to the trace
    model = solver->get_model();
    d_trace->set_model(model, current_depth, current_depth + cex_step);

    // Done, pop the solver scope
    scope.pop();

    // Moving to next depth 
    current_depth += cex_step;
  }

  // Remove the solver
  delete solver;

  return d_trace;
}

void pdkind_engine::gc_collect(const expr::gc_relocator& gc_reloc) {
  assert(d_induction_obligations_next.size() == 0);
  d_smt->gc_collect(gc_reloc);
  d_reachability.gc_collect(gc_reloc);
}

engine::invariant pdkind_engine::get_invariant() {
  return d_invariant;
}

}
}
