/*
 * 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 <boost/graph/strong_components.hpp>
#include <boost/graph/depth_first_search.hpp>
#include "WeakRefiner.h"
#include "Partition.h"

namespace parametric {
  
  using namespace std;
  using namespace boost;
  using namespace Rational;
  
  extern util::Timer nonSilentRefineTime;
  extern util::Timer silentRefineTime;
  
  template < typename Set > class dfs_non_div_visitor :
    public default_dfs_visitor {
  public:
    dfs_non_div_visitor(Set &_marking)
      : marking(_marking) {
    }
    template < typename Vertex, typename Graph >
    void discover_vertex(Vertex u, const Graph & g) const {
      marking.insert(u);
    }
    template < typename Vertex, typename Graph >
    void finish_vertex(Vertex u, const Graph & g) const {
      
    }
    
    Set &marking;
  };
  
  WeakRefiner::WeakRefiner(Quotient &quot)
    : Refiner(quot) {
  }
  
  /**
   * Checks whether @a u is a silent vertex.
   *
   * @param u vertex to be checked
   * @param A parition of @a u
   * @return true iff @a u is silent
   */
  bool WeakRefiner::weakIsSilent
  (vertex_descriptor u, EqClass &A) {
    pair<out_edge_iterator, out_edge_iterator> e_it;
    for (e_it = out_edges(u, graph); e_it.first != e_it.second; ++e_it.first) {
      edge_descriptor e = *e_it.first;
      vertex_descriptor v = target(e, graph);
      /* transition to outside partition */
      if (A.find(v) == A.end()) {
        return false;
      }
    }
    
    return true;
  }
  
  /**
   * Calculates P_Chi(u,C).
   *
   * @param u vertex
   * @param A partition of @a u
   * @param C target partition
   * @return P_Chi(u,C)
   */
  RationalFunction WeakRefiner::weakP_Chi
  (vertex_descriptor u, EqClass &A, EqClass &C) {
    pair<out_edge_iterator, out_edge_iterator> e_it;
    RationalFunction loop_el(0);
    RationalFunction out_el(0);
    for (e_it = out_edges(u, graph); e_it.first != e_it.second; ++e_it.first) {
      edge_descriptor e = *e_it.first;
      vertex_descriptor v = target(e, graph);
      if (A.find(v) != A.end()) {
        loop_el += graph[e].getValue();
      } else if (C.find(v) != C.end()) {
        out_el += graph[e].getValue();
      }
    }
    
    RationalFunction result(out_el / (RationalFunction(1) - loop_el));
    
    return result;
  }
  
  void WeakRefiner::silentSplit
  (EqClass &B, list<vertex_descriptor> &nonSilent,
   map<vertex_descriptor,vector<bool> > &label,
   map<vertex_descriptor,unsigned> &nonSilentLabel) {
    typedef graph_traits<reverse_graph<Graph> >::out_edge_iterator transp_out;
    
    reverse_graph<Graph> transp_graph(graph);
    for (list<vertex_descriptor>::iterator nonSilentIt = nonSilent.begin();
         nonSilentIt != nonSilent.end(); nonSilentIt++) {
      vertex_descriptor s = *nonSilentIt;
      unsigned classNr = nonSilentLabel[s];
      stack<vertex_descriptor> stack;
      stack.push(s);
      while (!stack.empty()) {
        vertex_descriptor u = stack.top();
        stack.pop();
        for (pair<transp_out, transp_out> it2 = out_edges(u, transp_graph);
             it2.first != it2.second; ++it2.first) {
          edge_descriptor e2 = *it2.first;
          vertex_descriptor v = target(e2, transp_graph);
          if ((B.find(v) != B.end())
              && weakIsSilent(v, B)
              && !label[v][classNr]) {
            label[v][classNr] = true;
            stack.push(v);
          }
        }
      }
    }
  }
  
  void WeakRefiner::refineClass(EqClass &eqClass) {
    nonSilentRefineTime.Start();
    
    /* enumerate neighbor partitions */
    map<EqClass*, unsigned> classMap;
    unsigned numClasses = 0;
    for (EqClass::iterator eqI = eqClass.begin(); eqI != eqClass.end(); eqI++) {
      vertex_descriptor u = *eqI;
      pair<out_edge_iterator, out_edge_iterator> oI;
      for (oI = out_edges(u, graph); oI.first != oI.second; ++oI.first) {
        edge_descriptor e = *oI.first;
        vertex_descriptor v = target(e, graph);
        EqClass *vClass = &*(*partition)[v];
        if ((vClass != &eqClass)
            && (classMap.find(vClass) == classMap.end())) {
          classMap[vClass] = numClasses;
          numClasses++;
        }
      }
    }
    
    /* for each non-silent vertex, each neighbored equivalence class, calculate
     * probabilities to go to this class. */
    HashMap<RationalFunction, unsigned> rMap;
    unsigned valueNumber = 1;
    for (EqClass::iterator eqI = eqClass.begin(); eqI != eqClass.end(); eqI++) {
      vertex_descriptor u = *eqI;
      if (!weakIsSilent(u, eqClass)) {
        pair<out_edge_iterator, out_edge_iterator> oI;
        map<EqClass*, RationalFunction> uMap;
        for (oI = out_edges(u, graph); oI.first != oI.second; ++oI.first) {
          edge_descriptor e = *oI.first;
          const RationalFunction &val(graph[e].getValue());
          vertex_descriptor v = target(e, graph);
          EqClass *vClass = &*(*partition)[v];
          map<EqClass*, RationalFunction>::iterator uMapIt(uMap.find(vClass));
          if (uMapIt == uMap.end()) {
            uMap.insert(make_pair(vClass, val));
          } else {
            RationalFunction oldVal(uMapIt->second);
            uMap.erase(uMapIt);
            uMap.insert(make_pair(vClass, oldVal + val));
          }
        }
        // TODO check following
        map<EqClass*, RationalFunction>::iterator uMapIt(uMap.find(&eqClass));
        if (uMap.end() != uMapIt) {
          RationalFunction dividor(RationalFunction(1) - uMapIt->second);
          uMap.erase(uMapIt);
          for (map<EqClass*, RationalFunction>::iterator it = uMap.begin();
               it != uMap.end(); it++) {
            it->second = it->second /  dividor;
          }
        }
        vector<unsigned> &vec = refined[u];
        vec.clear();
        vec.resize(numClasses, 0);
        
        for (uMapIt = uMap.begin(); uMapIt != uMap.end(); uMapIt++) {
          EqClass *eqClass = uMapIt->first;
          unsigned classNum = classMap[eqClass];
          const RationalFunction &val(uMapIt->second);
          HashMap<RationalFunction,unsigned>::iterator rMapIt =
            rMap.find(val);
          if (rMapIt == rMap.end()) {
            rMap.insert(make_pair(val, valueNumber));
            vec[classNum] = valueNumber;
            valueNumber++;
          } else {
            vec[classNum] = rMapIt->second;
          }
        }
      }
    }
    
    /* enumerate classes */
    int classNr = 0;
    map<vector<unsigned>,unsigned> classes;
    for (EqClass::iterator it = eqClass.begin(); it != eqClass.end(); it++) {
      vertex_descriptor u = *it;
      if (!weakIsSilent(u, eqClass)) {
        vector<unsigned> &eqp = refined[u];
        if (classes.find(eqp) == classes.end()) {
          classes.insert(make_pair(eqp, classNr));
          classNr++;
        }
      }
    }
    
    /* create label structure and remember non-silent states*/
    list<vertex_descriptor> nonSilent;
    map<vertex_descriptor,vector<bool> > label;
    map<vertex_descriptor,unsigned> nonSilentLabel;
    for (EqClass::iterator it = eqClass.begin(); it != eqClass.end(); it++) {
      vertex_descriptor u = *it;
      vector<bool> &vec = label[u];
      vec.resize(classes.size());
      if (!weakIsSilent(u, eqClass)) {
        vector<unsigned> &eqp = refined[u];
        unsigned classNr = classes[eqp];
        vec[classNr] = true;
        nonSilentLabel[u] = classNr;
        nonSilent.insert(nonSilent.begin(), u);
      }
    }
    nonSilentRefineTime.Stop();
    
    /* handle silent states */
    silentRefineTime.Start();
    silentSplit(eqClass, nonSilent, label, nonSilentLabel);
    silentRefineTime.Stop();
    
    /* create new classes */
    nonSilentRefineTime.Start();
    map<vector<bool>,EqClass> newClassMap;
    for (EqClass::iterator it = eqClass.begin(); it != eqClass.end(); it++) {
      vertex_descriptor u = *it;
      vector<bool> &vec = label[u];
      EqClass &eqClass = newClassMap[vec];
      eqClass.insert(u);
    }
    
    /* append these classes to list of new classes */
    map<vector<bool>,EqClass>::iterator it;
    if (newClassMap.size() == 1) {
      it = newClassMap.begin();
      partition->insert(it->second, false);
    } else {
      for (it = newClassMap.begin(); it != newClassMap.end(); it++) {
        partition->insert(it->second, true);
      }
    }
    nonSilentRefineTime.Stop();
  }
  
  void WeakRefiner::createQuotient(Partition &partition){
    PartitionList &P = partition.P;
    PartitionMap &P_map = partition.P_map;
    StateSet newInitStates;
    StateSet newTargetStates;
    
    PartitionList::iterator P_it;
    
    /* map states to their partitions */
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      set<vertex_descriptor> &entry = *P_it;
      set<vertex_descriptor>::iterator entry_it;
      for (entry_it = entry.begin(); entry_it != entry.end();
           entry_it++) {
        vertex_descriptor u = *entry_it;
        P_map.insert(make_pair(u, P_it));
      }
    }  
    
    /* find a representative non-silent state for each partition and create
     * vertices in the quotient graph. If in the partition an initial or
     * target state existed, the new state will be initial or target
     * respectively.
     */
    Graph *newGraph(new Graph);
    map<set<vertex_descriptor> *, vertex_descriptor> quotMap;
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      set<vertex_descriptor> &entry = *P_it;
      set<vertex_descriptor>::iterator entry_it;
      /* create new state in new graph representing partition */
      vertex_descriptor v;
      vertex_descriptor u;
      for (entry_it = entry.begin(); entry_it != entry.end(); entry_it++) {
        u = *entry_it;
        if (!weakIsSilent(u, entry)) {
          v = add_vertex(*newGraph);
          quotMap.insert(make_pair(&entry, v));
          break;
        }
      }
      if (entry_it == entry.end()) {
        v = add_vertex(*newGraph);
        quotMap.insert(make_pair(&entry, v));
      }
      
      /* find out whether new state is initial/target
       * and set reward if neccessary */
      bool is_init = false;
      bool is_target = false;
      bool rewardAlreadySet = false;
      for (entry_it = entry.begin(); entry_it != entry.end(); entry_it++) {
        vertex_descriptor u = *entry_it;
        if (mc.targetStates.find(u) != mc.targetStates.end()) {
          is_target = true;
        }
        if (mc.initStates.find(u) != mc.initStates.end()) {
          is_init = true;
        }
        if (mc.isRewardAnalysis()) {
          if (!rewardAlreadySet) {
            RationalFunction reward(mc.stateRewards[u]);
            mc.stateRewards.insert(make_pair(v, reward));
            rewardAlreadySet = true;
          }
          mc.stateRewards.erase(u);
        }
      }
      if (is_init) {
        newInitStates.insert(v);
      }
      if (is_target) {
        newTargetStates.insert(v);
      }
    }
    
    /* create edges between quotient graph nodes */
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      set<vertex_descriptor> &B_k = *P_it;
      vertex_descriptor B_k_node = quotMap[&B_k];
      vertex_descriptor xi;
      set<vertex_descriptor>::iterator entry_it;
      for (entry_it = B_k.begin(); entry_it != B_k.end(); entry_it++) {
        xi = *entry_it;
        if (!weakIsSilent(xi, B_k)) {
          break;
        }
      }
      if (entry_it != B_k.end()) {
        /* partition with non-silent states */
        pair<out_edge_iterator, out_edge_iterator> it2;
        for (it2 = out_edges(xi, graph); it2.first != it2.second; ++it2.first) {
          Graph::edge_descriptor e = *it2.first;
          vertex_descriptor xj = target(e, graph);
          set<vertex_descriptor> *B_l = &*P_map[xj];
          vertex_descriptor B_l_node = quotMap[B_l];
          RationalFunction newVal = weakP_Chi(xi, B_k, *B_l);
          pair<edge_descriptor,bool> e_p = edge(B_k_node, B_l_node, *newGraph);
          if (!e_p.second) {
            TransProp tp = graph[e];
            tp.setValue(newVal);
            add_edge(B_k_node, B_l_node, tp, *newGraph);
          }
        }
      } else {
        /* partition with only silent states */
        pair<edge_descriptor,bool> e_p = edge(B_k_node, B_k_node, *newGraph);
        if (!e_p.second
            && (newInitStates.find(B_k_node) == newInitStates.end())) {
          RationalFunction one(1);
          TransProp tp(one);
          RationalFunction reward(0);
          tp.setReward(reward);
          add_edge(B_k_node, B_k_node, tp, *newGraph);
        }
      }
    }
    mc.initStates = newInitStates;
    mc.targetStates = newTargetStates;
    delete mc.graph;
    mc.graph = newGraph;
  }
  
  void WeakRefiner::createInitialPartition(Partition &partition) {
    Refiner::createInitialPartition(partition);
    /* split off divergent states */
    reverse_graph<Graph> transp_graph(graph);
    PartitionList &P = partition.P;
    
    map<vertex_descriptor, int> color;
    associative_property_map<map<vertex_descriptor, int> > colorMap(color);
    PartitionList P_div;
    
    PartitionList::iterator P_it;
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      EqClass &A = *P_it;
      EqClass marked;
      dfs_non_div_visitor<set<vertex_descriptor> > non_div_vis(marked);
      for (EqClass::iterator A_it = A.begin(); A_it != A.end(); A_it++) {
        vertex_descriptor u = *A_it;
        bool non_silent  = false;
        /* check whether outgoing edges to other partitions exist */
        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;
          vertex_descriptor v = target(e, graph);
          /* transition to other partition? */
          if (A.find(v) == A.end()) {
            non_silent = true;
            break;
          }
        }
        
        /* mark vertex able to reach other partions */
        if (non_silent) {
          // TODO use subgraph of component here to improve performance
          // (may become important for reward models)
          depth_first_visit(transp_graph, u, non_div_vis, colorMap);
        }
      }
      /* place vertices not able to reach other partitions in new partition */
      set<vertex_descriptor> s;
      list<set<vertex_descriptor> >::iterator iA_div =
        P_div.insert(P_div.begin(), s);
      set<vertex_descriptor> &A_div = *iA_div;
      EqClass::iterator A_it;
      A_it = A.begin();
      while (A_it != A.end()) {
        vertex_descriptor u = *A_it;
        if (marked.find(u) == marked.end()) {
          set<vertex_descriptor>::iterator A_del_it = A_it;
          A_it++;
          A.erase(A_del_it);
          A_div.insert(u);
        } else {
          A_it++;
        }
      }
    }
    
    P.splice(P.begin(), P_div);
    /* remove empty partitions */
    P_it = P.begin();
    while (P_it != P.end()) {
      set<vertex_descriptor> &A = *P_it;
      if (A.empty()) {
        list<set<vertex_descriptor> >::iterator P_del_it = P_it;
        P_it++;
        P.erase(P_del_it);
      } else {
        P_it++;
      }
    }
  }
}
