// Copyright (C) 2010, Guy Barrand. All rights reserved.
// See the file tools.license for terms.

#ifndef tools_mpi_pntuple
#define tools_mpi_pntuple

// pntuple = for parallel ntupling.

#include <tools/wroot/base_pntuple>
#include <tools/S_STRING>
#include <tools/impi>

namespace tools {
namespace mpi {

class pntuple : public tools::wroot::base_pntuple {
  typedef tools::wroot::base_pntuple parent;
public:
  static tools::uint32 protocol_basket() {return 1;}
  static tools::uint32 protocol_end_fill() {return 2;}
public:
  pntuple(std::ostream& a_out,
          bool a_byte_swap,tools::uint32 a_compression,
          const std::string& a_name,const std::string& a_title,bool a_verbose)
  :parent(a_out,a_byte_swap,a_compression,a_name,a_title,a_verbose)
  {}
  pntuple(std::ostream& a_out,bool a_byte_swap,tools::uint32 a_compression,tools::wroot::seek a_seek_directory,
          const std::vector<tools::uint32>& a_basket_sizes,const tools::ntuple_booking& a_bkg,bool a_verbose)
  :parent(a_out,a_byte_swap,a_compression,a_seek_directory,a_basket_sizes,a_bkg,a_verbose)
  {}
  virtual ~pntuple() {}
protected:
  pntuple(const pntuple& a_from):parent(a_from){}
  pntuple& operator=(const pntuple& a_from){parent::operator=(a_from);return *this;}

protected:
  class basket_add : public virtual tools::wroot::branch::iadd_basket {
    typedef tools::wroot::branch::iadd_basket parent;
  public:
    virtual bool add_basket(tools::wroot::basket& a_basket) {
      m_mpi.pack_reset();
      if(!m_mpi.pack(protocol_basket())) return false;
      if(!m_mpi.pack(m_icol)) return false;

      if(!m_mpi.spack(a_basket.object_name())) return false;
      if(!m_mpi.spack(a_basket.object_title())) return false;

      if(!m_mpi.pack(a_basket.last())) return false;
      if(!m_mpi.pack(a_basket.nev_buf_size())) return false;
      if(!m_mpi.pack(a_basket.nev())) return false;
      
      if(a_basket.entry_offset()) {
        if(!m_mpi.bpack(true)) return false;
        if(!m_mpi.pack(a_basket.nev_buf_size(),a_basket.entry_offset())) return false;
      } else {
        if(!m_mpi.bpack(false)) return false;
      }

      if(a_basket.displacement()) {
        if(!m_mpi.bpack(true)) return false;
        if(!m_mpi.pack(a_basket.nev_buf_size(),a_basket.displacement())) return false;
      } else {
        if(!m_mpi.bpack(false)) return false;
      }

      if(!m_mpi.pack(a_basket.datbuf().length(),a_basket.datbuf().buf())) return false;

      if(!m_mpi.send_buffer(m_dest,m_tag)) return false;
      
      return true;
    }
  public:
    basket_add(file& a_file,tools::impi& a_mpi,int a_dest,int a_tag,tools::uint32 a_icol)
    :m_file(a_file),m_mpi(a_mpi),m_dest(a_dest),m_tag(a_tag),m_icol(a_icol)
    {}
  protected:
    basket_add(const basket_add& a_from):parent()
    ,m_file(a_from.m_file),m_mpi(a_from.m_mpi),m_dest(a_from.m_dest),m_tag(a_from.m_tag),m_icol(a_from.m_icol)
    {}
    basket_add& operator=(const basket_add& a_from){
      m_file = a_from.m_file;
      m_dest = a_from.m_dest;
      m_tag = a_from.m_tag;
      m_icol = a_from.m_icol;
      return *this;
    }
  protected:
    file m_file;
    tools::impi& m_mpi;
    int m_dest;
    int m_tag;
    tools::uint32 m_icol;
  };

public:
  bool add_row(tools::impi& a_mpi,int a_dest,int a_tag) {
    if(m_cols.empty()) return false;

    tools_vforit(tools::wroot::icol*,m_cols,it) (*it)->add();

    tools::uint32 icol = 0;
    tools_vforit(tools::wroot::icol*,m_cols,it) {
      basket_add _badd(m_file,a_mpi,a_dest,a_tag,icol);icol++;
      if(!(*it)->get_branch()->pfill(_badd)) return false;
    }

    tools_vforit(tools::wroot::icol*,m_cols,it) (*it)->set_def();
    return true;
  }
protected:
  void end_leaves(std::vector<tools::uint32>& a_lengths,std::vector<int>& a_mxs) const {
    a_lengths.clear();
    a_mxs.clear();
    tools_vforcit(tools::wroot::icol*,m_cols,it) {
      tools::wroot::base_leaf* _leaf = (*it)->get_leaf();
      tools::wroot::leaf_string* _leafs = _leaf?tools::id_cast<tools::wroot::base_leaf,tools::wroot::leaf_string>(*_leaf):0;
      if(_leafs) {
	a_lengths.push_back(_leafs->length());
        a_mxs.push_back(_leafs->get_max());        
      }      
    }
  }
public:
  bool end_fill(tools::impi& a_mpi,int a_dest,int a_tag) {
    tools::uint32 icol = 0;
    tools_vforit(tools::wroot::icol*,m_cols,it) {
      basket_add _badd(m_file,a_mpi,a_dest,a_tag,icol);icol++;
      if(!(*it)->get_branch()->end_pfill(_badd)) return false;
    }
    std::vector<tools::uint32> lengths;
    std::vector<int> mxs;
    end_leaves(lengths,mxs);

    a_mpi.pack_reset();
    if(!a_mpi.pack(protocol_end_fill())) return false;
    if(!a_mpi.vpack(lengths)) return false;
    if(!a_mpi.vpack(mxs)) return false;
    if(!a_mpi.send_buffer(a_dest,a_tag)) return false;

    return true;
  }

protected:
};

inline tools::wroot::basket* get_basket(std::ostream& a_out,tools::impi& a_mpi,
                                        bool a_byte_swap,tools::wroot::seek a_seek_directory,
                                        tools::uint32 a_basket_size) { 
  std::string oname;
  if(!a_mpi.sunpack(oname)) return 0;
  std::string otitle;
  if(!a_mpi.sunpack(otitle)) return 0;

  tools::uint32 last,nev_buf_size,nev;     
  if(!a_mpi.unpack(last)) return 0;
  if(!a_mpi.unpack(nev_buf_size)) return 0;
  if(!a_mpi.unpack(nev)) return 0;

  int* entry_offset = 0;
 {bool not_null;
  if(!a_mpi.bunpack(not_null)) return 0;
  if(not_null) {
    tools::uint32 n;
    if(!a_mpi.unpack(n,entry_offset)) return 0;
  }} 

  int* displacement = 0;
 {bool not_null;
  if(!a_mpi.bunpack(not_null)) return 0;
  if(not_null) {
    tools::uint32 n;
    if(!a_mpi.unpack(n,displacement)) return 0;
  }} 

  tools::uint32 _size;
  char* _buffer;
  if(!a_mpi.unpack(_size,_buffer)) {a_out << "unpack(buffer) failed."<< std::endl;return 0;}

  // wbasket fields set at file writing :
  //m_buf_size,m_buffer,m_nbytes,m_object_size,m_date,m_cycle,m_seek_key

  //::printf("debug :  col = %d, %s, %s, %lu\n",
  //    icol,rbasket.object_class().c_str(),rbasket.object_name().c_str(),rbasket.seek_key());

  tools::wroot::basket* basket = new tools::wroot::basket(a_out,a_byte_swap,a_seek_directory,
                                                          oname,otitle,"TBasket",a_basket_size,
                                                          false/*verbose*/);
  basket->datbuf().write_fast_array(_buffer,_size);
  basket->set_nev(last,nev_buf_size,nev,entry_offset,displacement);

  delete [] _buffer;

  return basket;
}

}}

#endif