/*
 * This file is part of PARAM.
 *
 * PARAM 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.
 *
 * PARAM 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 PARAM.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Copyright 2009 Ernst Moritz Hahn (emh@cs.uni-sb.de)
 */

#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <sstream>
#include <iomanip>
#include <boost/property_map.hpp>
#include <boost/graph/copy.hpp>
#include <boost/graph/named_function_params.hpp>
#include "SparseMC.h"
#include "PASSStateExplorer.h"
#include "Quotient.h"
#include "Eliminator.h"
#include "BoundedIterator.h"

namespace parametric {
  
  extern lang::Model model;
  namespace po = boost::program_options;
  using namespace std;
  using namespace Rational;
  
  SparseMC::SparseMC(po::variables_map &__vm)
    : vm(__vm) {
    statistics.totalTime.Start();
    graph = new Graph();
  }
  
  SparseMC::~SparseMC() {
    delete graph;
  }
  
  void SparseMC::print()
  {
    cout << endl << "----------[ Printing MC ]----------" << endl;
    
    pair<vertex_iterator, vertex_iterator> vp;
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {
      vertex_descriptor v =  *vp.first;
      cout << v << " : " << endl;    
      pair<out_edge_iterator, out_edge_iterator> it;
      for (it = out_edges(v, *graph); it.first != it.second; ++it.first) {
        edge_descriptor e = *it.first;
        vertex_descriptor u = target(e, *graph);
        cout << " ---[" << (*graph)[e].getValue() << "]---> " << 
          u << endl;
      }
    }
    
    cout << "-------------[ Done ]--------------" << endl << endl;
  }
  
  void SparseMC::printResult
  (const vector<RationalFunction> &result, const string &filename) const {
    ofstream file(filename.c_str(), ios::out);
    printResult(result, file);
    file.close();
  }
  
  void SparseMC::printResult(const vector<RationalFunction> &result, ostream &stream) const {
    for (vector<RationalFunction>::const_iterator it = result.begin();
         it != result.end(); it++) {
      stream << *it << endl;
    }
  }
  
  void SparseMC::printResult(const vector<RationalFunction> &result) const {
    printResult(result, cout);
  }
  
  void SparseMC::printDOT(const string &filename) const {
    ofstream file(filename.c_str(), ios::out);
    printDOT(file);
    file.close();
  }
  
  /**
   * Print model to output stream @a stream in GraphViz format.
   *
   * @param stream stream to print model to
   */
  void SparseMC::printDOT(ostream &stream) const {
    map<vertex_descriptor,unsigned> stateEnum;
    stream << "digraph G {" << endl;
    pair<vertex_iterator, vertex_iterator> vp;
    
    /* print states include target/initial markings */
    unsigned stateNum = 0;
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {
      Graph::vertex_descriptor v =  *vp.first;
      stateEnum[v] = stateNum;
      stream << "  " << stateNum;
      if (targetStates.find(v) != targetStates.end()) {
        stream << "[shape=\"circle\", style=\"filled\", fillcolor=\"black\"]";
      }
      stream << ";" << endl;
      if (initStates.find(v) != initStates.end()) {
        stream << "init_" << stateNum <<
          " [label=\"\", fillcolor=\"black\", width=\"0.0\""
          " shape=\"circle\", style=\"filled\", root];" << endl;
        stream << "  init_" << stateNum << " -> "
               << stateNum << endl;
      }
      stateNum++;
    }
    
    /* print transitions between states */
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {    
      Graph::vertex_descriptor v =  *vp.first;
      pair<out_edge_iterator, out_edge_iterator> it;
      for (it = out_edges(v, *graph); it.first != it.second; ++it.first) {
        Graph::edge_descriptor e = *it.first;
        Graph::vertex_descriptor u = target(e, *graph);
        stream << "  " << stateEnum[v] << " -> "      
               << stateEnum[u] << " [label=\"" << (*graph)[e].getValue();
        if (isRewardAnalysis()) {
          stream << " // " << (*graph)[e].getReward();
        }
        stream << "\",labelfontsize=10];" << endl;
      }    
    }
    stream << "}" << endl;
  }
  
  /**
   * Print model to standard output in GraphViz format.
   */
  void SparseMC::printDOT() const {
    printDOT(cout);
  }
  
  void SparseMC::embed() {
    std::pair<vertex_iterator, vertex_iterator> vp;
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {
      Graph::vertex_descriptor v =  *vp.first;
      RationalFunction sum(0);
      pair<out_edge_iterator, out_edge_iterator> it;
      for (it = out_edges(v, *graph); it.first != it.second; ++it.first) {
        Graph::edge_descriptor e = *it.first;
        sum += (*graph)[e].getValue();
      } 
      for (it = out_edges(v, *graph); it.first != it.second; ++it.first) {
        Graph::edge_descriptor e = *it.first;
        RationalFunction val((*graph)[e].getValue() / sum);
        RationalFunction reward(0);
        if (isRewardAnalysis()) {
          reward = (*graph)[e].getReward();
        }
        (*graph)[e] = TransProp(val);
        if (isRewardAnalysis()) {
          (*graph)[e].setReward(reward);
        }
      }    
    }
  }
  
  void SparseMC::getParamsFromModel() {
    for (HashMap<string,CVC3::Expr>::const_iterator i = model.variables.begin();
         i != model.variables.end(); ++i) {
      if (model.isParameterVariable(i->second)) {
        RationalFunction::addSymbol(i->first);
      }
    }
  }
  
  void SparseMC::prepareSymbols() {
    getParamsFromModel();
    RationalFunction::start();
  }
  
  void SparseMC::unboundedUntil(vector<RationalFunction> &result) {
    Eliminator eliminator;
    eliminator.setRewardAnalysis(isRewardAnalysis());
    eliminator.setGraph(graph);
    eliminator.setInitStates(initStates);
    eliminator.setTargetStates(targetStates);
    eliminator.setStateRewards(stateRewards);
    Eliminator::EliminationOrder eliminationOrder;
    const string eliminationOrderString(vm["elimination-order"].as<string>());
    if ("forward" == eliminationOrderString) {
      eliminationOrder = Eliminator::forward;
    } else if ("backward" == eliminationOrderString) {
      eliminationOrder = Eliminator::backward;
    } else if ("random" == eliminationOrderString) {
      eliminationOrder = Eliminator::random;
    } else {
      throw new runtime_error("Invalid elimination order");
    }
    eliminator.setEliminationOrder(eliminationOrder);
    eliminator.eliminate(result);
  }

  void SparseMC::boundedUntil(vector<RationalFunction> &result) {
    BoundedIterator iterator;
    iterator.setGraph(graph);
    iterator.setInitStates(initStates);
    iterator.setTargetStates(targetStates);
    iterator.setTimeBound(time);
    iterator.iterate(result);
  }
  
  /**
   * Executes the analysis of SparseMC.
   * The model and formula must have been loaded. Statespace must have been
   * explored, model type and analysis type etc. must have been set.
   */
  void SparseMC::execute(vector<RationalFunction>& result) {
    if ((model_type == lang::CTMC)
        && (unboundedUntilAnalysis == analysisType)) {
      embed();
    }
    if (vm["lump"].as<bool>()) {
      Quotient quot(*this);
      
      if (reachabilityRewardAnalysis == analysisType) {
        //        transToStateRewards();
      }
      
      if ("strong" == vm["lump-method"].as<string>()) {
        quot.setBisim(Quotient::strong);
      } else if ("weak" == vm["lump-method"].as<string>()) {
        quot.setBisim(Quotient::weak);
      }
      quot.quot();
    }
    
    RationalFunction::setCleanupMethod(RationalFunction::nocache);
    RationalFunction::removeUnusedSymbols();
    statistics.analysisTime.Start();
    
    if (unboundedUntilAnalysis == analysisType) {
      unboundedUntil(result);
    } else if (boundedUntilAnalysis == analysisType) {
      boundedUntil(result);
    } else if (reachabilityRewardAnalysis == analysisType) {
      if (reachabilityRewardAnalysis == analysisType) {
        //        stateToTransRewards();
      }
      unboundedUntil(result);
    } else {
      throw runtime_error("Illegal analysis type");
    }
    
    statistics.analysisTime.Stop();
  }
  
  /**
   * Convert state to transition rewards.
   * This conversion is needed for state elimination based reachability reward
   * analysis. For each state, the state reward is added to the rewards of all
   * outgoing edges, if existent. State rewards are set to zero then.
   */
  void SparseMC::stateToTransRewards() {
    pair<vertex_iterator, vertex_iterator> vp;
    RationalFunction zero(0);
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {
      vertex_descriptor u =  *vp.first;
      if (0 != stateRewards.count(u)) {
        const RationalFunction stateReward(stateRewards[u]);
        pair<out_edge_iterator, out_edge_iterator> it;
        for (it = out_edges(u, *graph); it.first != it.second; ++it.first) {
          edge_descriptor e = *it.first;
          RationalFunction newVal((*graph)[e].getReward() + stateReward);
          (*graph)[e].setReward(newVal);
        }
        stateRewards.erase(u);
      }
    }
  }
  
  /**
   * Convert transition to state rewards using intermediate states.
   * This conversion is neccessary for bisimulation minimizing.
   */
  void SparseMC::transToStateRewards() {
    pair<vertex_iterator, vertex_iterator> vp;
    
    /* collect edges to convert to states */
    vector<edge_descriptor> extendEdges;
    for (vp = vertices(*graph); vp.first != vp.second; ++vp.first) {
      vertex_descriptor u =  *vp.first;
      pair<out_edge_iterator, out_edge_iterator> it;
      for (it = out_edges(u, *graph); it.first != it.second; ++it.first) {
        edge_descriptor e = *it.first;
        if (RationalFunction(0) != (*graph)[e].getReward()) {
          extendEdges.push_back(e);
        }
      }
    }
    
    /* convert edges to intermediate states */
    for (unsigned i = 0; i < extendEdges.size(); i++) {
      edge_descriptor e = extendEdges[i];
      const TransProp &transProp = (*graph)[e];
      vertex_descriptor rewardState = add_vertex(*graph);
      stateRewards.insert(make_pair(rewardState, transProp.getReward()));
      
      pair<edge_descriptor,bool> origToReward =
        add_edge(source(e, *graph), rewardState, transProp, *graph);
      (*graph)[origToReward.first].setReward(RationalFunction(0));
      pair<edge_descriptor,bool> rewardToSucc =
        add_edge(rewardState, target(e, *graph), transProp, *graph);
      (*graph)[rewardToSucc.first].setReward(RationalFunction(0));
      (*graph)[rewardToSucc.first].setValue(RationalFunction(1));
      remove_edge(e, *graph);
    }
  }
  
  void SparseMC::parse() {
    RationalFunction::setCleanupMethod(RationalFunction::never);
    PASSStateExplorer explorer(*this);
    explorer.explore();
  }
}
