#include "TMap.h"
#include "TVectorD.h"
#include "TTree.h"
#include "TFile.h"
#include "TRandom3.h"

#include <iostream>

const double keyvals[] = {12,5,100};

void writetree(void)
{
  TFile outfile("outfile.root","RECREATE");
  outfile.cd();
  TTree * outtree = new TTree("outtree","Test tree with TMap branch");

  TMap * map = new TMap(2800);
  map->SetOwnerKeyValue(kTRUE,kTRUE);
  outtree->Branch("map","TMap",&map,32000,0);

  TRandom3 tr;

  for(int j = 0;j<10;j++)
    {
      map->DeleteAll();

      TVectorD * key = new TVectorD(3,keyvals);
      TVectorD * value = new TVectorD(10);
      for(int i = 0;i<10;i++) (*value)[i] = tr.Gaus();

      map->Add(key,value);
      outtree->Fill();
    }
  outfile.cd();
  gDirectory->Write();
  outfile.Close();
}

void readtree(void)
{
  TFile infile("outfile.root","READ");
  infile.cd();
  TTree * intree = (TTree*) infile.Get("outtree");
  TMap * map = new TMap(2800);
  intree->SetBranchAddress("map",&map);
  
  TVectorD * key = new TVectorD(3,keyvals);
  for(int j = 0;j<intree->GetEntries();j++)
    {
      intree->GetEntry(j);
      TVectorD * value = (TVectorD*) map->GetValue(key);
      if(! value) std::cout << "Could not get value for key";
      else std::cout << "Found value for key";
      key->Print();
      std::cout << "TMap printout: ";
      map->Print();
      std::cout << std::endl;
    }
}

void maptest(void)
{
  writetree();
  readtree();
}