/*
 * 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 "StrongRefiner.h"
#include "Partition.h"

namespace parametric {
  
  using namespace std;
  using namespace boost;
  using namespace Rational;
  
  extern util::Timer ringTime;
  
  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;
  };
  
  StrongRefiner::StrongRefiner(Quotient &quot) 
    : Refiner(quot) {
  }
  
  void StrongRefiner::refineClass(EqClass &eqClass) {
    /* 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 (classMap.find(vClass) == classMap.end()) {
          classMap[vClass] = numClasses;
          numClasses++;
        }
      }
    }
    
    /* for each vertex, each neighbored equivalence class, calculate
     * probabilities to go to this class. */
    typedef pair<RationalFunction,RationalFunction> RationalPair;
    HashMap<RationalPair, unsigned> rMap;
    unsigned valueNumber = 1;
    for (EqClass::iterator eqI = eqClass.begin(); eqI != eqClass.end(); eqI++) {
      vertex_descriptor u = *eqI;
      pair<out_edge_iterator, out_edge_iterator> oI;
      map<EqClass*, RationalPair> uMap;
      for (oI = out_edges(u, graph); oI.first != oI.second; ++oI.first) {
        edge_descriptor e = *oI.first;
        const RationalFunction &val(graph[e].getValue());
        const RationalFunction reward(graph[e].getReward());
        vertex_descriptor v = target(e, graph);
        EqClass *vClass = &*(*partition)[v];
        map<EqClass*, RationalPair>::iterator uMapIt(uMap.find(vClass));
        if (uMapIt == uMap.end()) {
          uMap.insert(make_pair(vClass, make_pair(val, reward)));
        } else {
          RationalPair oldVal(uMapIt->second);
          RationalPair newVal
            (make_pair(oldVal.first + val, oldVal.second + reward));
          uMap.erase(uMapIt);
          uMap.insert(make_pair(vClass, newVal));
        }
      }
      vector<unsigned> &vec = refined[u];
      vec.clear();
      vec.resize(numClasses, 0);
      
      map<EqClass*, RationalPair>::iterator uMapIt(uMap.find(&eqClass));
      for (uMapIt = uMap.begin(); uMapIt != uMap.end(); uMapIt++) {
        EqClass *eqClass = uMapIt->first;
        unsigned classNum = classMap[eqClass];
        const RationalPair &val(uMapIt->second);
        HashMap<RationalPair,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;
        }
      }
    }
    
    /* map vertices of same values to same new class*/
    map<vector<unsigned>,EqClass> newClassMap;
    for (EqClass::iterator it = eqClass.begin(); it != eqClass.end(); it++) {
      vertex_descriptor u = *it;
      vector<unsigned> &eqp = refined[u];
      EqClass &newEqClass = newClassMap[eqp];
      newEqClass.insert(u);
    }
    
    /* append these classes to list of new classes */
    map<vector<unsigned>,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);
      }
    }
  }
  
  void StrongRefiner::createQuotient(Partition &partition){
    PartitionList &P = partition.P;
    PartitionMap &P_map = partition.P_map;
    /* build new graph from blocks */
    PartitionList::iterator P_it;
    
    StateSet newInitStates;
    StateSet newTargetStates;
    /* create states of new graph */
    Graph *newGraph(new Graph);
    map<EqClass *, vertex_descriptor> quotMap;
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      EqClass &entry = *P_it;
      vertex_descriptor v = add_vertex(*newGraph);
      quotMap.insert(make_pair(&entry, v));
      EqClass::iterator entry_it;
      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.isRewardAnalysis()) {
          if (!rewardAlreadySet) {
            RationalFunction reward(mc.stateRewards[u]);
            mc.stateRewards.insert(make_pair(v, reward));
            rewardAlreadySet = true;
          }
          mc.stateRewards.erase(u);
        }
        if (mc.getTargetStates().find(u) != mc.getTargetStates().end()) {
          is_target = true;
        }
        if (mc.getInitStates().find(u) != mc.getInitStates().end()) {
          is_init = true;
        }
      }
      if (is_init) {
        newInitStates.insert(v);
      }
      if (is_target) {
        newTargetStates.insert(v);
      }
    }
    
    /* create transitions of new graph */
    for (P_it = P.begin(); P_it != P.end(); P_it++) {
      EqClass &B_k = *P_it;
      vertex_descriptor B_k_node = quotMap[&B_k];
      vertex_descriptor xi = *B_k.begin();
      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];
        const RationalFunction &Q_sum(graph[e].getValue());
        const RationalFunction &Q_reward(graph[e].getReward());
        pair<edge_descriptor,bool> e_p = edge(B_k_node, B_l_node, *newGraph);
        if (!e_p.second) {
          TransProp tp;
          tp.setValue(Q_sum);
          tp.setReward(Q_reward);
          add_edge(B_k_node, B_l_node, tp, *newGraph);
        } else {
          const RationalFunction &old_val((*newGraph)[e_p.first].getValue());
          const RationalFunction &old_reward((*newGraph)[e_p.first].getReward());
          RationalFunction newVal(old_val + Q_sum);
          RationalFunction newReward((old_reward * old_val + Q_reward * Q_sum) / newVal);
          (*newGraph)[e_p.first].setValue(newVal);
          (*newGraph)[e_p.first].setReward(newReward);
        }
      }
    }
    
    mc.initStates = newInitStates;
    mc.targetStates = newTargetStates;
    mc.graph = newGraph;
  }
}

