/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitCore                                                       *
 *    File: $Id: RooAbsCategory.cc,v 1.45 2005/06/20 15:44:44 wverkerke Exp $
 * Authors:                                                                  *
 *   WV, Wouter Verkerke, UC Santa Barbara, verkerke@slac.stanford.edu       *
 *   DK, David Kirkby,    UC Irvine,         dkirkby@uci.edu                 *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California          *
 *                          and Stanford University. All rights reserved.    *
 *                                                                           *
 * Redistribution and use in source and binary forms,                        *
 * with or without modification, are permitted according to the terms        *
 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)             *
 *****************************************************************************/

// -- CLASS DESCRIPTION [CAT] --
// RooAbsCategory is the common abstract base class for objects that represent a
// discrete value with a finite number of states. Each state consist of a 
// label/index pair, which is stored in a RooCatType object.
// 
// Implementation of RooAbsCategory may be derived, there no interface
// is provided to modify the contents, nor a public interface to define states.

#include "RooFit.h"

#include "Riostream.h"
#include "Riostream.h"
#include <stdlib.h>
#include "TString.h"
#include "TH1.h"
#include "TTree.h"
#include "TLeaf.h"
#include "RooAbsCategory.h"
#include "RooArgSet.h"
#include "Roo1DTable.h"
#include "RooCategory.h"

ClassImp(RooAbsCategory) 
;

RooAbsCategory::RooAbsCategory(const char *name, const char *title) : 
  RooAbsArg(name,title)
{
  // Constructor
  _typeIter = _types.MakeIterator() ;
  setValueDirty() ;  
  setShapeDirty() ;  
}

RooAbsCategory::RooAbsCategory(const RooAbsCategory& other,const char* name) :
  RooAbsArg(other,name), _value(other._value) 
{
  // Copy constructor, copies the registered category states from the original.
  _typeIter = _types.MakeIterator() ;

  other._typeIter->Reset() ;
  TObject* obj ;
  while ((obj=other._typeIter->Next())) {
    _types.Add(obj->Clone()) ;
  }

  setValueDirty() ;
  setShapeDirty() ;
}


 RooAbsCategory::~RooAbsCategory()
{
  // Destructor

  // We own the contents of _types 
  delete _typeIter ;
  _types.Delete() ;
}


 Int_t RooAbsCategory::getIndex() const
{
  // Return index number of current state 

  if (isValueDirty() || isShapeDirty()) {
    _value = traceEval() ;

    clearValueDirty() ;
    clearShapeDirty() ;
  }

  return _value.getVal() ;
}


 const char* RooAbsCategory::getLabel() const
{
  // Return label string of current state 

  if (isValueDirty() || isShapeDirty()) {
    _value = traceEval() ;

    clearValueDirty() ;
    clearShapeDirty() ;
  }

  return _value.GetName() ;
}




 RooCatType RooAbsCategory::traceEval() const
{
  // Recalculate current value and check validity of new result.

  RooCatType value = evaluate() ;
  
  // Standard tracing code goes here
  if (!isValid(value)) {
  }

  // Call optional subclass tracing code
  traceEvalHook(value) ;

  return value ;
}



 TIterator* RooAbsCategory::typeIterator() const
{
  // Return iterator over all defined states
  return _types.MakeIterator() ;
}


Bool_t RooAbsCategory::operator==(Int_t index) const
{
  // Equality operator with a integer (compares with state index number)
  return (index==getIndex()) ;
}

Bool_t RooAbsCategory::operator==(const char* label) const
{
  // Equality operator with a string (compares with state label string)
  return !TString(label).CompareTo(getLabel()) ;
}



Bool_t RooAbsCategory::operator==(const RooAbsArg& other) 
{
  const RooAbsCategory* otherCat = dynamic_cast<const RooAbsCategory*>(&other) ;
  return otherCat ? operator==(otherCat->getIndex()) : kFALSE ;
}


 Bool_t RooAbsCategory::isValidIndex(Int_t index) const
{
  // Check if state with given index is defined
  return lookupType(index)?kTRUE:kFALSE ;
}

 Bool_t RooAbsCategory::isValidLabel(const char* label) const
{
  // Check if state with given name is defined
  return lookupType(label)?kTRUE:kFALSE ;
}



 const RooCatType* RooAbsCategory::defineType(const char* label)
{
  // Define a new state with given name. The lowest available
  // integer number is assigned as index value

  // Find lowest unused index
  Int_t index(-1) ;
  while(lookupType(++index,kFALSE)) ;
  
  // Assign this index to given label 
  return defineType(label,index) ;
}


 const RooCatType* RooAbsCategory::defineTypeUnchecked(const char* label, Int_t index) 
{
  Bool_t first = _types.GetEntries()?kFALSE:kTRUE ;
  RooCatType *newType = new RooCatType(label,index) ;
  _types.Add(newType) ;

  if (first) _value = RooCatType(label,index) ;
  setShapeDirty() ;

  return newType ;  
}



 const RooCatType* RooAbsCategory::defineType(const char* label, Int_t index) 
{
  // Define new state with given name and index number.

  if (isValidIndex(index)) {
    cout << "RooAbsCategory::defineType(" << GetName() << "): index " 
	 << index << " already assigned" << endl ;
    return 0 ;
  }

  if (isValidLabel(label)) {
    cout << "RooAbsCategory::defineType(" << GetName() << "): label " 
	 << label << " already assigned or not allowed" << endl ;
    return 0 ;
  }

  return defineTypeUnchecked(label,index) ;
}



 void RooAbsCategory::clearTypes() 
{
  // Delete all currently defined states

  _types.Delete() ;
  _value = RooCatType("",0) ;
  setShapeDirty() ;
}



 const RooCatType* RooAbsCategory::lookupType(const RooCatType &other, Bool_t printError) const 
{
  // Find our type that matches the specified type, or return 0 for no match.

  RooCatType* type ;
  _typeIter->Reset() ;
  while((type=(RooCatType*)_typeIter->Next())){
    if((*type) == other) return type; // delegate comparison to RooCatType
  }

  if (printError) {
    cout << ClassName() << "::" << GetName() << ":lookupType: no match for ";
    other.printToStream(cout,OneLine);
  }
  return 0 ;
}

 const RooCatType* RooAbsCategory::lookupType(Int_t index, Bool_t printError) const
{
  // Find our type corresponding to the specified index, or return 0 for no match.

  RooCatType* type ;
  _typeIter->Reset() ;
  while((type=(RooCatType*)_typeIter->Next())){  
    if((*type) == index) return type; // delegate comparison to RooCatType
  }
  if (printError) {
    cout << ClassName() << "::" << GetName() << ":lookupType: no match for index "
	 << index << endl;
  }
  return 0 ;
}

 const RooCatType* RooAbsCategory::lookupType(const char* label, Bool_t printError) const 
{
  // Find our type corresponding to the specified label, or return 0 for no match.
 
  RooCatType* type ;
  _typeIter->Reset() ;
  while((type=(RooCatType*)_typeIter->Next())){  
    if((*type) == label) return type; // delegate comparison to RooCatType
  }

  // Try if label represents integer number
  char* endptr ;
  Int_t idx=strtol(label,&endptr,10)  ;
  if (endptr==label+strlen(label)) {
    _typeIter->Reset() ;
    while((type=(RooCatType*)_typeIter->Next())){  
       if((*type) == idx) return type; // delegate comparison to RooCatType
     }
  }

  if (printError) {
    cout << ClassName() << "::" << GetName() << ":lookupType: no match for label "
	 << label << endl;
  }
  return 0 ;
}

 Bool_t RooAbsCategory::isValid() const
{
  // Check if current value is a valid state
  return isValid(_value) ;
}

 Bool_t RooAbsCategory::isValid(RooCatType value)  const
{
  // Check if given state is defined for this object
  return isValidIndex(value.getVal()) ;
}

 Roo1DTable* RooAbsCategory::createTable(const char *label)  const
{
  // Create a table matching the shape of this category
  return new Roo1DTable(GetName(),label,*this) ;
}

 Bool_t RooAbsCategory::readFromStream(istream&, Bool_t, Bool_t) 
{
  // Read object contents from stream (dummy for now)

  return kFALSE ;
} 

 void RooAbsCategory::writeToStream(ostream& os, Bool_t compact) const
{
  // Write object contents to stream 
  if (compact) {
    os << getLabel() ;
  } else {
    os << getLabel() ;
  }
}

 void RooAbsCategory::printToStream(ostream& os, PrintOption opt, TString indent) const
{
  // Print info about this object to the specified stream. In addition to the info
  // from RooAbsArg::printToStream() we add:
  //
  //     Shape : label, index, defined types

  RooAbsArg::printToStream(os,opt,indent);
  if(opt >= Shape) {
    os << indent << "--- RooAbsCategory ---" << endl;
    if (_types.GetEntries()==0) {
      os << indent << "  ** No values defined **" << endl;
      return;
    }
    os << indent << "  Value is \"" << getLabel() << "\" (" << getIndex() << ")" << endl;
    os << indent << "  Has the following possible values:" << endl;
    indent.Append("    ");
    opt= lessVerbose(opt);
    RooCatType *type;
    _typeIter->Reset() ;
    while((type=(RooCatType*)_typeIter->Next())) {
      os << indent;
      type->printToStream(os,opt,indent);
    }
  }
}


 void RooAbsCategory::attachToTree(TTree& t, Int_t bufSize)
{
  // Attach the category index and label to as branches
  // to the given TTree. The index field will be attached
  // as integer with name <name>_idx, the label field will be attached
  // as char[] with label <name>_lbl.

  // First check if there is an integer branch matching the category name
  TBranch* branch = t.GetBranch(GetName()) ;
  if (branch) {

    TString typeName(((TLeaf*)branch->GetListOfLeaves()->At(0))->GetTypeName()) ;
    if (!typeName.CompareTo("Int_t")) {
      // Imported TTree: attach only index field as branch

      cout << "RooAbsCategory::attachToTree(" << GetName() << ") TTree branch " << GetName() 
	   << " will be interpreted as category index" << endl ;

      t.SetBranchAddress(GetName(),&((Int_t&)_value._value)) ;
      setAttribute("INTIDXONLY_TREE_BRANCH",kTRUE) ;      
      return ;
    } else if (!typeName.CompareTo("UChar_t")) {
      cout << "RooAbsReal::attachToTree(" << GetName() << ") TTree UChar_t branch " << GetName() 
	   << " will be interpreted as category index" << endl ;
      t.SetBranchAddress(GetName(),&((Bool_t&)_value._value)) ;
      setAttribute("UCHARIDXONLY_TREE_BRANCH",kTRUE) ;
      return ;
    } 

    if (branch->GetCompressionLevel()<0) {
      cout << "RooAbsCategory::attachToTree(" << GetName() << ") Fixing compression level of branch " << GetName() << endl ;
      branch->SetCompressionLevel(1) ;
    }
  }

  // Native TTree: attach both index and label of category as branches  
  TString idxName(GetName()) ;
  TString lblName(GetName()) ;  
  idxName.Append("_idx") ;
  lblName.Append("_lbl") ;
  
  // First determine if branch is taken
  if ((branch = t.GetBranch(idxName))) {    

    t.SetBranchAddress(idxName,&((Int_t&)_value._value)) ;
    if (branch->GetCompressionLevel()<0) {
      cout << "RooAbsCategory::attachToTree(" << GetName() << ") Fixing compression level of branch " << idxName << endl ;
      branch->SetCompressionLevel(1) ;
    }
    
  } else {    
    TString format(idxName);
    format.Append("/I");
    void* ptr = &(_value._value) ;
    branch = t.Branch(idxName, ptr, (const Text_t*)format, bufSize);
    branch->SetCompressionLevel(1) ;
  }
  
  // First determine if branch is taken
  if ((branch = t.GetBranch(lblName))) {

    t.SetBranchAddress(lblName,_value._label) ;
    if (branch->GetCompressionLevel()<0) {
      cout << "RooAbsCategory::attachToTree(" << GetName() << ") Fixing compression level of branch " << lblName << endl ;
      branch->SetCompressionLevel(1) ;
    }

  } else {    
    TString format(lblName);
    format.Append("/C");
    void* ptr = _value._label ;
    branch = t.Branch(lblName, ptr, (const Text_t*)format, bufSize);
    branch->SetCompressionLevel(1) ;
  }
}


 void RooAbsCategory::fillTreeBranch(TTree& t) 
{
  // Attach object to a branch of given TTree

  TString idxName(GetName()) ;
  TString lblName(GetName()) ;  
  idxName.Append("_idx") ;
  lblName.Append("_lbl") ;

  // First determine if branch is taken
  TBranch* idxBranch = t.GetBranch(idxName) ;
  TBranch* lblBranch = t.GetBranch(lblName) ;
  if (!idxBranch||!lblBranch) { 
    cout << "RooAbsCategory::fillTreeBranch(" << GetName() << ") ERROR: not attached to tree" << endl ;
    assert(0) ;
  }

  idxBranch->Fill() ;
  lblBranch->Fill() ;  
}


 void RooAbsCategory::setTreeBranchStatus(TTree& t, Bool_t active) 
{
  // (De)Activate associate tree branch
  TBranch* branch = t.GetBranch(Form("%s_idx",GetName())) ;
  if (branch) { 
    t.SetBranchStatus(Form("%s_idx",GetName()),active?1:0) ;
    t.SetBranchStatus(Form("%s_lbl",GetName()),active?1:0) ;
  }
}


 void RooAbsCategory::syncCache(const RooArgSet*) 
{ 
  getIndex() ; 
}


 void RooAbsCategory::copyCache(const RooAbsArg* source) 
{
  // Copy the cached value from given source and raise dirty flag.
  // It is the callers responsability to unsure that the sources
  // cache is clean before this function is called, e.g. by
  // calling syncCache() on the source.

  RooAbsCategory* other = dynamic_cast<RooAbsCategory*>(const_cast<RooAbsArg*>(source)) ;
  assert(other!=0) ;

  if (source->getAttribute("INTIDXONLY_TREE_BRANCH")) {
    // Lookup cat state from other-index because label is missing
    const RooCatType* type = lookupType(other->_value._value) ;
    if (type) {
      _value = *type ;
    } else {
      cout << "RooAbsCategory::copyCache(" << GetName() 
	   << ") ERROR: index of source arg " << source->GetName() 
	   << " is invalid (" << other->_value._value 
	   << "), value not updated" << endl ;
    }
  } if (source->getAttribute("UCHARIDXONLY_TREE_BRANCH")) {
    // Lookup cat state from other-index because label is missing
    Int_t tmp = *reinterpret_cast<UChar_t*>(&(other->_value._value)) ;
    const RooCatType* type = lookupType(tmp) ;
    if (type) {
      _value = *type ;
    } else {
      cout << "RooAbsCategory::copyCache(" << GetName() 
	   << ") ERROR: index of source arg " << source->GetName() 
	   << " is invalid (" << tmp
	   << "), value not updated" << endl ;
    }
  } else {
    _value = other->_value ;
  }

  setValueDirty() ;
}


 const RooCatType* RooAbsCategory::getOrdinal(UInt_t n) const 
{
  // Return state definition of ordinal nth defined state,
  // needed by the generator mechanism.
  
  return (const RooCatType*)_types.At(n);
}

 RooAbsArg *RooAbsCategory::createFundamental(const char* newname) const 
{
  // Create a RooCategory fundamental object with our properties.

  // Add and precalculate new category column 
  RooCategory *fund= new RooCategory(newname?newname:GetName(),GetTitle()) ; 

  // Copy states
  TIterator* tIter = typeIterator() ;
  RooCatType* type ;
  while ((type=(RooCatType*)tIter->Next())) {
    ((RooAbsCategory*)fund)->defineType(type->GetName(),type->getVal()) ;
  }
  delete tIter;

  return fund;
}


 Bool_t RooAbsCategory::isSignType(Bool_t mustHaveZero) const 
{
  // Determine if category has 2 or 3 states with index values -1,0,1

  if (numTypes()>3||numTypes()<2) return kFALSE ;
  if (mustHaveZero&&numTypes()!=3) return kFALSE ;

  Bool_t ret(kTRUE) ;
  TIterator* tIter = typeIterator() ;
  RooCatType* type ;
  while((type=(RooCatType*)tIter->Next())) {
    if (abs(type->getVal())>1) ret=kFALSE ;
  }
  
  delete tIter ;
  return ret ;
}