/*****************************************************************************
 * Project: RooFit                                                           *
 * Package: RooFitModels                                                     *
 *    File: $Id: RooNonCPEigenDecay.cc,v 1.23 2005/06/21 16:46:33 wverkerke Exp $
 * Authors:                                                                  *
 *   AH, Andreas Hoecker,  Orsay,            hoecker@slac.stanford.edu       *
 *   SL, Sandrine Laplace, Orsay,            laplace@slac.stanford.edu       *
 *   JS, Jan Stark,        Paris,            stark@slac.stanford.edu         *
 *   WV, Wouter Verkerke,  UC Santa Barbara, verkerke@slac.stanford.edu      *
 *                                                                           *
 * Copyright (c) 2000-2005, Regents of the University of California,         *
 *                          IN2P3. All rights reserved.                      *
 *                                                                           *
 * History                                                                   *
 *   Nov-2001   WV Created initial version                                   *
 *   Dec-2001   SL mischarge correction, direct CPV                          *
 *   Jan-2002   AH built dedicated generator + code cleaning                 *
 *   Mar-2002   JS committed debugged version to CVS                         *
 *   Apr-2002   AH allow prompt (ie, non-Pdf) mischarge treatment            *
 *   May-2002   JS Changed the set of CP parameters (mathematically equiv.)  *
 *                                                                           *
 * 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 [PDF] --
// Time-dependent RooAbsAnaConvPdf for CP violating decays 
// to Non-CP eigenstates (eg, B0 -> rho+- pi-+).
// For a description of the physics model see the 
// BaBar Physics Book, section 6.5.2.3 .
// The set of CP parameters used in this class is equivalent to
// the one used in the Physics Book, but it is not exactly the
// same. Starting from the set in the BaBar Book, in order to 
// get the parameters used here you have to change the sign of both
// a_c^+ and a_c^-, and then substitute:
//    a_s^Q = S + Q* deltaS
//    a_c^Q = C + Q*deltaC
// where Q denotes the charge of the rho.

#include "RooFit.h"

#include "Riostream.h"
#include "Riostream.h"
#include "RooRealVar.h"
#include "RooRandom.h"
#include "RooNonCPEigenDecay.h"
#include "TMath.h"

ClassImp(RooNonCPEigenDecay);

#define Debug_RooNonCPEigenDecay 1

RooNonCPEigenDecay::RooNonCPEigenDecay( const char *name, const char *title, 
					RooRealVar&     t,
					RooAbsCategory& tag,
					RooAbsReal&     tau,
					RooAbsReal&     dm,
					RooAbsReal&     avgW,
					RooAbsReal&     delW,
					RooAbsCategory& rhoQ,
					RooAbsReal&     correctQ,
					RooAbsReal&     wQ,
					RooAbsReal&     acp,
					RooAbsReal&     C,
					RooAbsReal&     delC,
					RooAbsReal&     S,
					RooAbsReal&     delS,
					const RooResolutionModel& model,
					DecayType       type )
  : RooAbsAnaConvPdf( name, title, model, t ),
  _acp      ( "acp",      "acp",                this, acp      ),
  _avgC        ( "C",        "C",                  this, C        ),
  _delC     ( "delC",     "delC",               this, delC     ),
  _avgS        ( "S",        "S",                  this, S        ),
  _delS     ( "delS",     "delS",               this, delS     ),
  _avgW     ( "avgW",     "Average mistag rate",this, avgW     ),
  _delW     ( "delW",     "Shift mistag rate",  this, delW     ),
  _t        ( "t",        "time",               this, t        ),
  _tau      ( "tau",      "decay time",         this, tau      ),
  _dm       ( "dm",       "mixing frequency",   this, dm       ),
  _tag      ( "tag",      "CP state",           this, tag      ),
  _rhoQ     ( "rhoQ",     "Charge of the rho",  this, rhoQ     ),
  _correctQ ( "correctQ", "correction of rhoQ", this, correctQ ),
  _wQ       ( "wQ",       "mischarge",          this, wQ       ),
  _genB0Frac     ( 0 ),
  _genRhoPlusFrac( 0 ),
  _type     ( type )
{

  // Constructor
  switch(type) {
  case SingleSided:
    _basisExp = declareBasis( "exp(-@0/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(-@0/@1)*sin(@0*@2)", RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(-@0/@1)*cos(@0*@2)", RooArgList( tau, dm ) );
    break;
  case Flipped:
    _basisExp = declareBasis( "exp(@0)/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(@0/@1)*sin(@0*@2)",  RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(@0/@1)*cos(@0*@2)",  RooArgList( tau, dm ) );
    break;
  case DoubleSided:
    _basisExp = declareBasis( "exp(-abs(@0)/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(-abs(@0)/@1)*sin(@0*@2)", RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(-abs(@0)/@1)*cos(@0*@2)", RooArgList( tau, dm ) );
    break;
  }
}

RooNonCPEigenDecay::RooNonCPEigenDecay( const char *name, const char *title, 
					RooRealVar&     t,
					RooAbsCategory& tag,
					RooAbsReal&     tau,
					RooAbsReal&     dm,
					RooAbsReal&     avgW,
					RooAbsReal&     delW,
					RooAbsCategory& rhoQ,
					RooAbsReal&     correctQ,
					RooAbsReal&     acp,
					RooAbsReal&     C,
					RooAbsReal&     delC,
					RooAbsReal&     S,
					RooAbsReal&     delS,
					const RooResolutionModel& model,
					DecayType       type )
  : RooAbsAnaConvPdf( name, title, model, t ),
  _acp      ( "acp",      "acp",                this, acp      ),
  _avgC        ( "C",        "C",                  this, C        ),
  _delC     ( "delC",     "delC",               this, delC     ),
  _avgS        ( "S",        "S",                  this, S        ),
  _delS     ( "delS",     "delS",               this, delS     ),
  _avgW     ( "avgW",     "Average mistag rate",this, avgW     ),
  _delW     ( "delW",     "Shift mistag rate",  this, delW     ),
  _t        ( "t",        "time",               this, t        ),
  _tau      ( "tau",      "decay time",         this, tau      ),
  _dm       ( "dm",       "mixing frequency",   this, dm       ),
  _tag      ( "tag",      "CP state",           this, tag      ),
  _rhoQ     ( "rhoQ",     "Charge of the rho",  this, rhoQ     ),
  _correctQ ( "correctQ", "correction of rhoQ", this, correctQ ),
  _genB0Frac     ( 0 ),
  _genRhoPlusFrac( 0 ),
  _type     ( type )
{

  // dummy mischarge (must be set to zero!)
  _wQ = RooRealProxy( "wQ", "mischarge", this, *(new RooRealVar( "wQ", "wQ", 0 )) );

  switch(type) {
  case SingleSided:
    _basisExp = declareBasis( "exp(-@0/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(-@0/@1)*sin(@0*@2)", RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(-@0/@1)*cos(@0*@2)", RooArgList( tau, dm ) );
    break;
  case Flipped:
    _basisExp = declareBasis( "exp(@0)/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(@0/@1)*sin(@0*@2)",  RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(@0/@1)*cos(@0*@2)",  RooArgList( tau, dm ) );
    break;
  case DoubleSided:
    _basisExp = declareBasis( "exp(-abs(@0)/@1)",            RooArgList( tau     ) );
    _basisSin = declareBasis( "exp(-abs(@0)/@1)*sin(@0*@2)", RooArgList( tau, dm ) );
    _basisCos = declareBasis( "exp(-abs(@0)/@1)*cos(@0*@2)", RooArgList( tau, dm ) );
    break;
  }
}

RooNonCPEigenDecay::RooNonCPEigenDecay( const RooNonCPEigenDecay& other, const char* name ) 
  : RooAbsAnaConvPdf( other, name ),
  _acp      ( "acp",      this, other._acp      ),
  _avgC        ( "C",        this, other._avgC        ),
  _delC     ( "delC",     this, other._delC     ),
  _avgS        ( "S",        this, other._avgS        ),
  _delS     ( "delS",     this, other._delS     ),
  _avgW     ( "avgW",     this, other._avgW     ),
  _delW     ( "delW",     this, other._delW     ),
  _t        ( "t",        this, other._t        ),
  _tau      ( "tau",      this, other._tau      ),
  _dm       ( "dm",       this, other._dm       ),
  _tag      ( "tag",      this, other._tag      ),
  _rhoQ     ( "rhoQ",     this, other._rhoQ     ),
  _correctQ ( "correctQ", this, other._correctQ ),
  _wQ       ( "wQ",       this, other._wQ       ),
  _genB0Frac     ( other._genB0Frac      ),
  _genRhoPlusFrac( other._genRhoPlusFrac ),
  _type          ( other._type           ),
  _basisExp      ( other._basisExp       ),
  _basisSin      ( other._basisSin       ),
  _basisCos      ( other._basisCos       )
{
  // Copy constructor
}

RooNonCPEigenDecay::~RooNonCPEigenDecay( void )
{
  // Destructor
}

Double_t RooNonCPEigenDecay::coefficient( Int_t basisIndex ) const 
{
  // B0    : _tag  == -1 
  // B0bar : _tag  == +1 
  // rho+  : _rhoQ == +1
  // rho-  : _rhoQ == -1
  // the charge corrrection factor "_correctQ" serves to implement mis-charges

  Int_t rhoQc = _rhoQ * int(_correctQ);
  assert( rhoQc == 1 || rhoQc == -1 );

  Double_t a_sin_p = _avgS + _delS;
  Double_t a_sin_m = _avgS - _delS;
  Double_t a_cos_p = _avgC + _delC;
  Double_t a_cos_m = _avgC - _delC;

  if (basisIndex == _basisExp) {
    if (rhoQc == -1 || rhoQc == +1)
      return (1 + rhoQc*_acp*(1 - 2*_wQ))*(1 + 0.5*_tag*(2*_delW));
    else
      return 1;
  }

  if (basisIndex == _basisSin) {

    if (rhoQc == -1)

      return - ((1 - _acp)*a_sin_m*(1 - _wQ) + (1 + _acp)*a_sin_p*_wQ)*(1 - 2*_avgW)*_tag;

    else if (rhoQc == +1)

      return - ((1 + _acp)*a_sin_p*(1 - _wQ) + (1 - _acp)*a_sin_m*_wQ)*(1 - 2*_avgW)*_tag;

    else
       return - _tag*((a_sin_p + a_sin_m)/2)*(1 - 2*_avgW);
  }

  if (basisIndex == _basisCos) {

    if ( rhoQc == -1)

      return + ((1 - _acp)*a_cos_m*(1 - _wQ) + (1 + _acp)*a_cos_p*_wQ)*(1 - 2*_avgW)*_tag;

    else if (rhoQc == +1)

      return + ((1 + _acp)*a_cos_p*(1 - _wQ) + (1 - _acp)*a_cos_m*_wQ)*(1 - 2*_avgW)*_tag;

    else
      return _tag*((a_cos_p + a_cos_m)/2)*(1 - 2*_avgW);
  }

  return 0;
}

// advertise analytical integration
Int_t RooNonCPEigenDecay::getCoefAnalyticalIntegral( RooArgSet& allVars, 
						     RooArgSet& analVars, const char* rangeName ) const
{
  if (rangeName) return 0 ;

  if (matchArgs( allVars, analVars, _tag, _rhoQ )) return 3;
  if (matchArgs( allVars, analVars, _rhoQ       )) return 2;
  if (matchArgs( allVars, analVars, _tag        )) return 1;

  return 0;
}

Double_t RooNonCPEigenDecay::coefAnalyticalIntegral( Int_t basisIndex, 
						     Int_t code, const char* /*rangeName*/ ) const
{
  // correct for the right/wrong charge...
  Int_t rhoQc = _rhoQ*int(_correctQ);

  Double_t a_sin_p = _avgS + _delS;
  Double_t a_sin_m = _avgS - _delS;
  Double_t a_cos_p = _avgC + _delC;
  Double_t a_cos_m = _avgC - _delC;

  switch(code) {

    // No integration
  case 0: return coefficient(basisIndex);

    // Integration over 'tag'
  case 1:
    if (basisIndex == _basisExp) return 2*(1 + rhoQc*_acp*(1 - 2*_wQ));
    if (basisIndex == _basisSin || basisIndex==_basisCos) return 0;
    assert( kFALSE );

    // Integration over 'rhoQ'
  case 2:
    if (basisIndex == _basisExp) return 2*(1 + 0.5*_tag*(2.*_delW));

    if (basisIndex == _basisSin)

      return - ( (1 - _acp)*a_sin_m + (1 + _acp)*a_sin_p )*(1 - 2*_avgW)*_tag;

    if (basisIndex == _basisCos)

      return + ( (1 - _acp)*a_cos_m + (1 + _acp)*a_cos_p )*(1 - 2*_avgW)*_tag;

    assert( kFALSE );

    // Integration over 'tag' and 'rhoQ'
  case 3:
    if (basisIndex == _basisExp) return 2*2; // for both: tag and charge
    if (basisIndex == _basisSin || basisIndex==_basisCos) return 0;
    assert( kFALSE );

  default:
    assert( kFALSE );
  }

  return 0;
}

Int_t RooNonCPEigenDecay::getGenerator( const RooArgSet& directVars, 
					RooArgSet&       generateVars, Bool_t staticInitOK ) const
{
  if (staticInitOK) {
    if (matchArgs( directVars, generateVars, _t, _tag, _rhoQ )) return 4;
    if (matchArgs( directVars, generateVars, _t, _rhoQ       )) return 3;
    if (matchArgs( directVars, generateVars, _t, _tag        )) return 2;
  }
  if (matchArgs( directVars, generateVars, _t              )) return 1;
  return 0;
}

void RooNonCPEigenDecay::initGenerator( Int_t code )
{

  if (code == 2 || code == 4) {
    // Calculate the fraction of mixed events to generate
    Double_t sumInt1 = RooRealIntegral( "sumInt1", "sum integral1", *this,
					RooArgSet( _t.arg(), _tag.arg(), _rhoQ.arg() )
				      ).getVal();
    _tag = -1;
    Double_t b0Int1 = RooRealIntegral( "mixInt1", "mix integral1", *this,
				       RooArgSet( _t.arg(), _rhoQ.arg() )
				     ).getVal();
    _genB0Frac = b0Int1/sumInt1;

    if (Debug_RooNonCPEigenDecay == 1)
      cout << "     o RooNonCPEigenDecay::initgenerator: genB0Frac     : "
	   << _genB0Frac
	   << ", tag dilution: " << (1 - 2*_avgW)
	   << endl;
  }

  if (code == 3 || code == 4) {
    // Calculate the fraction of positive rho's to generate
    Double_t sumInt2 = RooRealIntegral( "sumInt2", "sum integral2", *this,
					RooArgSet( _t.arg(), _tag.arg(), _rhoQ.arg() )
				      ).getVal();
    _rhoQ = 1;
    Double_t b0Int2 = RooRealIntegral( "mixInt2", "mix integral2", *this,
				       RooArgSet( _t.arg(), _tag.arg() )
				     ).getVal();
    _genRhoPlusFrac = b0Int2/sumInt2;

    if (Debug_RooNonCPEigenDecay == 1)
      cout << "     o RooNonCPEigenDecay::initgenerator: genRhoPlusFrac: "
	   << _genRhoPlusFrac << endl;
  }
}


void RooNonCPEigenDecay::generateEvent( Int_t code )
{

  // Generate delta-t dependent
  while (kTRUE) {

    // B flavor and rho charge (we do not use the integrated weights)
    if (code != 1) {
      if (code != 3) _tag  = (RooRandom::uniform()<=0.5) ? -1 : +1;
      if (code != 2) _rhoQ = (RooRandom::uniform()<=0.5) ?  1 : -1;
    }

    // opposite charge?
    // Int_t rhoQc = _rhoQ*int(_correctQ);

    Double_t a_sin_p = _avgS + _delS;
    Double_t a_sin_m = _avgS - _delS;
    Double_t a_cos_p = _avgC + _delC;
    Double_t a_cos_m = _avgC - _delC;

    // maximum probability density 
    double a1 = 1 + sqrt(TMath::Power(a_cos_m, 2) + TMath::Power(a_sin_m, 2));
    double a2 = 1 + sqrt(TMath::Power(a_cos_p, 2) + TMath::Power(a_sin_p, 2));

    Double_t maxAcceptProb = (1.10 + TMath::Abs(_acp)) * (a1 > a2 ? a1 : a2);
    // The 1.10 in the above line is a security feature to prevent crashes close to the limit at 1.00

    Double_t rand = RooRandom::uniform();
    Double_t tval(0);

    switch(_type) {

    case SingleSided:
      tval = -_tau*log(rand);
      break;

    case Flipped:
      tval = +_tau*log(rand);
      break;

    case DoubleSided:
      tval = (rand<=0.5) ? -_tau*log(2*rand) : +_tau*log(2*(rand-0.5));
      break;
    }

    // get coefficients
    Double_t expC = coefficient( _basisExp );
    Double_t sinC = coefficient( _basisSin );
    Double_t cosC = coefficient( _basisCos );

    // probability density
    Double_t acceptProb  = expC + sinC*sin(_dm*tval) + cosC*cos(_dm*tval);

    // sanity check...
    assert( acceptProb <= maxAcceptProb );

    // hit or miss...
    Bool_t accept = maxAcceptProb*RooRandom::uniform() < acceptProb ? kTRUE : kFALSE;

    if (accept && tval<_t.max() && tval>_t.min()) {
      _t = tval;
      break;
    }
  }
}