/* Copyright (c) 2018 -- 2022
   Charles Jordan (Hokkaido University, Sapporo, Japan),
   Michael Joswig, Lars Kastner (Technische Universitaet Berlin, Germany)

   This program 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, or (at your option) any
   later version: http://www.gnu.org/licenses/gpl.txt.

   This program 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.
-----------------------------------------------------------------------------
*/

///////////////////////////////////////////////////////////////////////////////
// TriangulationOptimizer.hh
//
// This object computes the canonical (i.e. lex largest) class of a
// triangulation.
//
// The basic principle is the following: The GroupDecomposer decomposes the
// group in a way that allows one to find the canonical representative with a
// fast DFS. The TriangulationOptimizer object encapsulates this DFS.
///////////////////////////////////////////////////////////////////////////////

#ifndef TRIANGULATIONOPTIMIZER_HH
#define TRIANGULATIONOPTIMIZER_HH

#include "GroupDecomposer.hh"
#include "mptopcom_typedefs.hh"
#include <polymake/Bitset.h>
#include <utility>

class TriangulationOptimizerWrapper {
   public:
      virtual const TriangNode& get_optimal() = 0;
      virtual const gkz_vector_type& get_optimal_gkz() = 0;
      virtual const Symmetry& get_permutation() = 0;
};

template<bool trackPermutation, typename lexMinimal, typename checkFine>
class FullOrbitTriangulationOptimizer : public TriangulationOptimizerWrapper {
   private:
      TriangNode optimal;
      gkz_vector_type optimal_gkz;
      const GroupDecomposer& gd;
      const gkz_type& gkz;
      Symmetry mult;

      gkz_vector_type compute_gkz(const TriangNode& t){
         if(checkFine::value){
            gkz_vector_type result = gkz.evaluate(t);
            Scalar fineCoord = gkz.fine_coordinate(result);
            return fineCoord | result;
         } else {
            return gkz.evaluate(t);
         }
      }

      void optimize() {
         // If the group is small, it is cheaper to just iterate over all
         // of it.
         const SymmetryGroup& group(gd.get_group());
         TriangNode start(optimal), tmp;
         gkz_vector_type tmp_gkz;
         triang_vector_type optimal_tvec(gkz.triangulationAsVector(optimal)), tmp_tvec;
         for(SymmetryGroup::const_iterator sym_iter = group.begin();
               sym_iter != group.end();
               ++sym_iter){
            Symmetry g(*sym_iter);
            tmp = g(start);
            tmp_gkz = compute_gkz(tmp);
            Int comp;
            if(lexMinimal::value) {
               comp = -lex_compare(tmp_gkz, optimal_gkz);
            } else {
               comp = lex_compare(tmp_gkz, optimal_gkz);
            }
            if(comp == 1){
               optimal = tmp;
               optimal_gkz = tmp_gkz;
               optimal_tvec = gkz.triangulationAsVector(tmp);
               if(trackPermutation){
                  mult = g;
               }
            } else if (comp == 0){
               tmp_tvec = gkz.triangulationAsVector(tmp);
               if(optimal_tvec > tmp_tvec){
                  optimal = tmp;
                  optimal_gkz = tmp_gkz;
                  optimal_tvec = tmp_tvec;
                  if(trackPermutation){
                     mult = g;
                  }
               }
            }
         }
      }

   public:
      FullOrbitTriangulationOptimizer(const TriangNode& t, 
         const gkz_vector_type& gkzt, 
         const GroupDecomposer& gd_in, 
         const gkz_type& gkz_in) :
         optimal(t), optimal_gkz(gkzt),
         gd(gd_in), gkz(gkz_in)
   { 
      if(checkFine::value){
         optimal_gkz = compute_gkz(t);
      }
      if(trackPermutation){
         mult = array2symmetry(gd.get_identity());
      }
      optimize();
   }
      
      const TriangNode& get_optimal(){
         return optimal;
      }

      const gkz_vector_type& get_optimal_gkz(){
         if(checkFine::value){
            optimal_gkz = gkz.evaluate(optimal);
         }
         return optimal_gkz;
      }
      
      const Symmetry& get_permutation() {
         return mult;
      }

};


class PackagedTriangulation {
   private:
      triang_vector_type tvec;
      const gkz_type& gkz;
   public:
      PackagedTriangulation(const TriangNode& t_in, const gkz_type& gkz_in):
         gkz(gkz_in), tvec(gkz_in.triangulationAsVector(t_in))
      {}
      
      PackagedTriangulation(const triang_vector_type& t, const gkz_type& gkz_in):
         gkz(gkz_in), tvec(t)
      {}

      const pm::Set<Int> get_support(const Int& fixed, const pm::Set<Int>& in, bool& needsIdentity) const {
         pm::Set<Int> intersection(tvec*in);
         if(intersection.empty()){
            needsIdentity = true;
            return in;
         } else {
            if(tvec.contains(fixed)){
               needsIdentity = true;
            }
            return intersection;
         }
      }

      bool operator >(const PackagedTriangulation& other) const {
         return other.tvec > tvec;
      }

      PackagedTriangulation& operator= (const PackagedTriangulation& other) {
         tvec = other.tvec;
         return *this;
      }

      PackagedTriangulation mutate(const pm::Array<Int>* g) const {
         Symmetry tmp(array2symmetry(*g));
         triang_vector_type next;
         for(const auto& e: tvec){
            next += gkz.get_basis2int()[tmp(gkz.get_int2basis()[e])];
         }
         return PackagedTriangulation(next, gkz);
      }
};

class ReferenceCore {
   private:
      const pm::Array<Int>* identity;
      std::map<Int, std::map<Int, const pm::Array<Int>*>> switch_table;
      pm::Map<Int, pm::Set<Int>> supports;
      Int bound;

   public:
      ReferenceCore(std::map<Int, std::map<Int, const pm::Array<Int>*>> st, const pm::Array<Int>* id):
         identity(id), switch_table(st), bound(0)
      {
         for(const auto& level : switch_table){
            for(const auto& sublevel : level.second){
               supports[level.first] += sublevel.first;
            }
            bound = level.first > bound ? level.first : bound;
         }
         bound++;
      }

      void extract_switches(Int fixed, const pm::Set<Int>& desired, std::list<const pm::Array<Int>*>& switches) const {
         for(const auto& j : desired){
            switches.push_back(switch_table.at(fixed).at(j));
         }
      }

      Int get_bound() const {
         return bound;
      }

      const pm::Array<Int>& get_identity() const {
         return *identity;
      }

      bool support_exists(Int fixed) const {
         return supports.exists(fixed);
      }

      const pm::Set<Int>& get_support(Int fixed) const {
         return supports[fixed];
      }

};

template<bool trackPermutation, typename lexMinimal>
class TriangulationOptimizer : public TriangulationOptimizerWrapper {
   private:
      TriangNode optimal;
      gkz_vector_type optimal_gkz;
      const GroupDecomposer& gd;
      const gkz_type& gkz;
      Symmetry mult;


      // Find lex largest representative of class that also lex maximizes the
      // gkz vector.
      void optimize(){
         const polymake::group::SwitchTable& ST(gd.get_switch_table());
std::pair<gkz_vector_type, pm::Array<Int>> P;
         if(lexMinimal::value) {
         P = ST.lex_minimize_vector<Scalar>(optimal_gkz);
         } else {
         P = ST.lex_maximize_vector<Scalar>(optimal_gkz);
         }
         mult = Symmetry(array2symmetry(P.second));
         optimal = mult(optimal);
         optimal_gkz = P.first;
         optimize_representative();
      }

      // Every set in result corresponds to a subset of indices of the gkz
      // vector with the same value. These need to be stabilized if we want to
      // change the representative while preserving the gkz vector.
      const signature_vector_type signature() const{
         signature_vector_type result(optimal_gkz.dim());
         pm::Map<Scalar, Int> numbers, count;
         for(Int i=0; i<result.dim(); i++){
            if(numbers.exists(optimal_gkz[i])){
               result[i] = numbers[optimal_gkz[i]];
               count[optimal_gkz[i]]++;
            } else {
               result[i] = i;
               numbers[optimal_gkz[i]] = i;
               count[optimal_gkz[i]] = 1;
            }
         }
         return result;
      }


      // Find the lex largest representative, while keeping the gkz-vector.
      void optimize_representative(){
         std::map<Int, std::map<Int, const pm::Array<Int>*>> refST;
         signature_vector_type sig = signature();
         bool needsOpt = gd.populate_partition_stabilizer(sig, refST);
         if(needsOpt){
            ReferenceCore RC(std::move(refST), &(gd.get_identity()));
            PackagedTriangulation packaged(optimal, gkz);
            polymake::group::switchtable::Optimizer<ReferenceCore, PackagedTriangulation> opt(RC, packaged);
            opt.optimize();
            Symmetry g(array2symmetry(opt.get_optimal().second));
            optimal = g(optimal);
            if(trackPermutation){
               mult = g * mult;
            }
         }
      }



   public:
      TriangulationOptimizer(const TriangNode& t, 
         const gkz_vector_type& gkzt, 
         const GroupDecomposer& gd_in, 
         const gkz_type& gkz_in) :
         optimal(t), optimal_gkz(gkzt),
         gd(gd_in), gkz(gkz_in)
   { 
      if(trackPermutation){
         mult = array2symmetry(gd.get_identity());
      }
   }
      
      const TriangNode& get_optimal(){
         optimize();
         return optimal;
      }

      const gkz_vector_type& get_optimal_gkz(){
         return optimal_gkz;
      }

      const Symmetry& get_permutation() {
         return mult;
      }

};

#endif

