/****************************************************************\
| This file is part of the PAX example files for CDF             |
| -------------------------------------------------------------- | 
|   Author of the last change: $Author: erdmann $                 
|   Date of last change:       $Date: 2003/11/18 15:13:17 $   
|   Revision number:           $Revision: 1.2 $        
\****************************************************************/
//
//! example algorithm
//
//! MakeW reconstructs the decay chain W -> lepton + neutrino
//! from the W-mass constraint.
//
//! The algorithm uses on input an event interpretation, a lepton 
//! fourvector of this input event interpretation, the missing
//! transverse momentum components px, py in the event, and a value 
//! for the W-boson mass.
//
//! On output it delivers two event interpretations corresponding
//! to two solutions of the reconstructed W boson.
//! If the solutions are identical, only one event interpretation
//! is filled, the second remains empty.
//
//! On return, the number of W solutions is given
//
#include "TPAXAnaModule.hh"

#include "PAX/PAXexperimentclass.hh"
#include "PAX/PAXfourvector.hh"
#include "PAX/PAXinterpret.hh"
#include "PAX/PAXiterator.hh"
#include "PAX/PAXprintlevel.hh"
#include "PAX/PAXvertex.hh"

#include <iostream>

int TPAXAnaModule::MakeW( PaxEventInterpret* event_interpret_in,
			  PaxEventInterpret* event_interpret_W1, 
			  PaxEventInterpret* event_interpret_W2, 
			  PaxFourVector*     lepton_in,
			  double             missing_px_in = 0,
			  double             missing_py_in = 0,
			  double             W_mass = 80.42 )
{
  // choose pax names

  string name_neutrino             = "neutrino";
  string name_W                    = "W";
  string name_W_decay_vertex       = "W decay vertex";
  string name_event_interpretation = "W event interpretation";

  // number of reconstructed W bosons

  int number_of_W = 0;

  // don't do anything if proper input missing

  if (        event_interpret_in->IsEmpty() 
       ||   ! lepton_in 
       || ( ! missing_px_in && ! missing_py_in) )
    {
      cout << "makeW: proper input missing. "
	   << "Input event interpret not empty (1) = "
	   << !event_interpret_in->IsEmpty() 
	   << " lepton_in = "  << lepton_in
	   << " px = " << missing_px_in
	   << " py = " << missing_py_in
	   << endl;
      return number_of_W;
    }

  // if output event interpretations are not empty: clear them

  if ( ! event_interpret_W1->IsEmpty() )  event_interpret_W1->clear();
  if ( ! event_interpret_W2->IsEmpty() )  event_interpret_W2->clear();

  // find lepton in input event interpretation
 
  PaxFourVector* lepton = 0;

  lepton = event_interpret_in->
    fourvector->findItemByKey( lepton_in->getPaxId() );

  // continue only if fourvector was found

  if ( !lepton ) 
    {
      cout << "makeW: lepton not found in input event interpretation" << endl;
      return number_of_W;
    }

  // particle data group codes

  int pdg_lepton   = lepton->getParticleId();
  int pdg_neutrino = 0;
  int pdg_W        = 0;

  switch ( abs(pdg_lepton) )
    {
    case 11:
      pdg_neutrino = -12;
      pdg_W        = -24;
      break;
    case 13:
      pdg_neutrino = -14;
      pdg_W        = -24;
      break;
    case 15:
      pdg_neutrino = -16;
      pdg_W        = -24;
      break;
    }
  if ( lepton->getCharge() > 0 )
    {
      pdg_neutrino = -pdg_neutrino;
      pdg_W        = -pdg_W;
    }

  // -------------------------------------------------------
  // calculate solutions for the W

  double px = missing_px_in;
  double py = missing_py_in;

  // missing transverse momentum

  double MET = sqrt( pow( px, 2 ) + pow( py, 2 ) );

  // solutions of neutrino pz using W-mass constraint

  double pz1 = 0;
  double pz2 = 0;

  double phi     = atan2(py,px);
  double epsilon = cos( lepton->phi() ) * cos(phi) 
                 + sin( lepton->phi() ) * sin(phi);
  double mu      = pow(W_mass,2)/2 + epsilon * lepton->perp() * MET;
	  
  if  ( (  mu*mu*pow(lepton->pz(),2)-pow(MET,2) 
	 * pow(lepton->e(),2)*pow(lepton->perp(),2)
	 + mu*mu*pow(lepton->perp(),2))<0 )
    {
      pz1 = 1/pow(lepton->perp(),2)*mu*lepton->pz();
      pz2 = 1/pow(lepton->perp(),2)*mu*lepton->pz();
    }
  else
    {
      pz1 = 1/pow(lepton->perp(),2)
  	   *(mu*lepton->pz()+sqrt(mu*mu*pow(lepton->pz(),2)-pow(MET,2)
	   *pow(lepton->e(),2)*pow(lepton->perp(),2)
	   +mu*mu*pow(lepton->perp(),2)));
      
      pz2 = 1/pow(lepton->perp(),2)
	   *(mu*lepton->pz()-sqrt(mu*mu*pow(lepton->pz(),2)-pow(MET,2)
	   *pow(lepton->e(),2)*pow(lepton->perp(),2)
	   +mu*mu*pow(lepton->perp(),2)));
    }

  // return if no meaningful solution

  if ( pz1 == 0 && pz2 == 0 )
    return number_of_W;

  // ---------------------------------------------------------------

  // register the unique PAX identifier of the lepton in the user record 
  // of the input event interpretation: "id_" at the beginning of the name 
  // tells PAX that the value is a unique PAX identifier. It is automatically
  // updated when copying the event interpretation below.

  string lepton_name = "id_lepton temporarily used by MakeW";
  event_interpret_in->add( lepton_name, lepton->getPaxId() );

  // find primary vertex by name, register the unique PAX identifier
  // of the vertex in the user record of the input event interpretation.

  string vertex_name = "id_primary vertex temporarily used by MakeW";

  PaxVertex* primary_vertex = 0;

  PaxIterator<PaxVertex*>* itervx = event_interpret_in->vertex->GetIterator();

  itervx->First(); 
  while ( !primary_vertex && !itervx->IsDone() )
    {
      string s[3];  s[0]="primary"; s[1]="Primary"; s[2]="PRIMARY";
      if ( itervx->CurrentItem()->getPaxName().find(s[0]) < s[0].length() 
        || itervx->CurrentItem()->getPaxName().find(s[1]) < s[1].length() 
        || itervx->CurrentItem()->getPaxName().find(s[2]) < s[2].length() )
	primary_vertex = itervx->CurrentItem();
      itervx->Next();
    }

  if ( primary_vertex )
    event_interpret_in->add( vertex_name, primary_vertex->getPaxId() );

  // ---------------------------------------------------------------

  // here we start filling the first output event interpretation W1
  // fill it with a copy of the input event interpretation

  event_interpret_W1->copy( event_interpret_in );

  // set name of event interpretation

  event_interpret_W1->setPaxName( name_event_interpretation );

  // find the new PaxFourVector instance of the lepton in W1 (the unique PAX 
  // identifiers in the user record are automatically updated when copying 
  // an event interpretation)

  lepton           = 0;
  int lepton_paxid = 0;

  lepton_paxid     = event_interpret_W1->user_record->findItemByKey( lepton_name  );
  if ( lepton_paxid )
    lepton         = event_interpret_W1->fourvector ->findItemByKey( lepton_paxid );

  // remove temporary lepton entry in the user record

  event_interpret_W1->remove( lepton_name );

  // find the new PaxVertex instance of the primary vertex in W1 (the unique 
  // PAX identifiers in the user record are automatically updated when copying 
  // an event interpretation)

  primary_vertex    = 0;
  int  vertex_paxid = 0;

  vertex_paxid      = event_interpret_W1->user_record->findItemByKey( vertex_name  );
  if ( vertex_paxid )
    primary_vertex  = event_interpret_W1->vertex     ->findItemByKey( vertex_paxid );

  // remove temporary primary vertex entry in the user record

  event_interpret_W1->remove( vertex_name );

  // construction of the W decay chain 

  if ( lepton )
    {
      // make neutrino fourvector

      PaxFourVector* neutrino = new PaxFourVector;
      event_interpret_W1->add(neutrino);

      double e = sqrt( pow(px,2) + pow(py,2) + pow(pz1,2) );

      neutrino->setPx(px);
      neutrino->setPy(py);
      neutrino->setPz(pz1);
      neutrino->setE (e );
      neutrino->setPaxName(name_neutrino);
      neutrino->setParticleId( pdg_neutrino );

      // make W

      PaxFourVector* Wboson = new PaxFourVector;
      event_interpret_W1->add( Wboson );

      Wboson->setPx        ( lepton->px() + neutrino->px() );
      Wboson->setPy        ( lepton->py() + neutrino->py() );
      Wboson->setPz        ( lepton->pz() + neutrino->pz() );
      Wboson->setMass      ( W_mass               );
      Wboson->setPaxName   ( name_W               );
      Wboson->setCharge    ( lepton->getCharge()  );
      Wboson->setParticleId( pdg_W                );
      Wboson->addExperimentClassRelations( lepton );

      number_of_W++;

      // make W decay vertex

      PaxVertex* decay_vertex = new PaxVertex;
      event_interpret_W1->add ( decay_vertex );

      decay_vertex->setPaxName( name_W_decay_vertex );

      // change the begin vertex of the lepton 

      itervx = lepton->begin_vertex_relations->GetIterator();

      if ( ! itervx->IsDone() )
	{
	  // fill values of the begin vertex of the lepton 
	  // as the W decay vertex

	  decay_vertex->setX( itervx->CurrentItem()->x() );
	  decay_vertex->setY( itervx->CurrentItem()->y() );
	  decay_vertex->setZ( itervx->CurrentItem()->z() );
	  
	  // check if primary vertex was found above, 
	  // if not, use begin vertex of lepton
	  
	  if ( ! primary_vertex )
	    primary_vertex = itervx->CurrentItem();

	  // remove begin vertex from lepton, 

	  lepton->begin_vertex_relations
	    ->remove( itervx->CurrentItem(), lepton );
	}

      // add primary vertex for Wboson

      if ( primary_vertex )
	Wboson  ->begin_vertex_relations->add   ( primary_vertex, Wboson );

      // add W decay vertex to Wboson, lepton, and neutrino

      Wboson  ->  end_vertex_relations->add( decay_vertex, Wboson   );
      lepton  ->begin_vertex_relations->add( decay_vertex, lepton   );
      neutrino->begin_vertex_relations->add( decay_vertex, neutrino );
    }
  else
    {
      // remove temporary lepton entry and vertex entry in user record
      event_interpret_in->remove( lepton_name );
      event_interpret_in->remove( vertex_name );

      cout << "makeW: lepton not found in W1 event interpretation" << endl;
      return number_of_W;
    }

  // ---------------------------------------------------------------

  // identical solution for neutrino

  if ( pz1 == pz2 )
    {
      // remove temporary lepton and vertex entries in user record
      event_interpret_in->remove( lepton_name );
      event_interpret_in->remove( vertex_name );

      return number_of_W;
    };

  // copy of the above code to fill event interpretation W2, except for the
  // neutrino kinematics pz

  // here we start filling the first output event interpretation W2
  // fill it with a copy of the input event interpretation

  event_interpret_W2->copy( event_interpret_in );

  // set name of event interpretation

  event_interpret_W2->setPaxName( name_event_interpretation );

  // find the new PaxFourVector instance of the lepton in W2 (the unique PAX 
  // identifiers in the user record are automatically updated when copying 
  // an event interpretation)

  lepton           = 0;
  lepton_paxid     = 0;

  lepton_paxid     = event_interpret_W2->user_record->findItemByKey( lepton_name  );
  if ( lepton_paxid )
    lepton         = event_interpret_W2->fourvector ->findItemByKey( lepton_paxid );

  // remove temporary lepton entry in the user record

  event_interpret_W2->remove( lepton_name );

  // find the new PaxVertex instance of the primary vertex in W2 (the unique 
  // PAX identifiers in the user record are automatically updated when copying 
  // an event interpretation)

  primary_vertex    = 0;
  vertex_paxid      = 0;

  vertex_paxid      = event_interpret_W2->user_record->findItemByKey( vertex_name  );
  if ( vertex_paxid )
    primary_vertex  = event_interpret_W2->vertex     ->findItemByKey( vertex_paxid );

  // remove temporary primary vertex entry in the user record

  event_interpret_W2->remove( vertex_name );

  // construction of the W decay chain 

  if ( lepton )
    {
      // make neutrino fourvector

      PaxFourVector* neutrino = new PaxFourVector;
      event_interpret_W2->add(neutrino);

      double e = sqrt( pow(px,2) + pow(py,2) + pow(pz2,2) );

      neutrino->setPx(px);
      neutrino->setPy(py);
      neutrino->setPz(pz2);
      neutrino->setE (e );
      neutrino->setPaxName(name_neutrino);
      neutrino->setParticleId( pdg_neutrino );

      // make W

      PaxFourVector* Wboson = new PaxFourVector;
      event_interpret_W2->add( Wboson );

      Wboson->setPx        ( lepton->px() + neutrino->px() );
      Wboson->setPy        ( lepton->py() + neutrino->py() );
      Wboson->setPz        ( lepton->pz() + neutrino->pz() );
      Wboson->setMass      ( W_mass               );
      Wboson->setPaxName   ( name_W               );
      Wboson->setCharge    ( lepton->getCharge()  );
      Wboson->setParticleId( pdg_W                );
      Wboson->addExperimentClassRelations( lepton );

      number_of_W++;

      // make W decay vertex

      PaxVertex* decay_vertex = new PaxVertex;
      event_interpret_W2->add ( decay_vertex );

      decay_vertex->setPaxName( name_W_decay_vertex );

      // change the begin vertex of the lepton 

      itervx = lepton->begin_vertex_relations->GetIterator();

      if ( ! itervx->IsDone() )
	{
	  // fill values of the begin vertex of the lepton 
	  // as the W decay vertex

	  decay_vertex->setX( itervx->CurrentItem()->x() );
	  decay_vertex->setY( itervx->CurrentItem()->y() );
	  decay_vertex->setZ( itervx->CurrentItem()->z() );
	  
	  // check if primary vertex was found above, 
	  // if not, use begin vertex of lepton
	  
	  if ( ! primary_vertex )
	    primary_vertex = itervx->CurrentItem();

	  // remove begin vertex from lepton, 

	  lepton->begin_vertex_relations
	    ->remove( itervx->CurrentItem(), lepton );
	}

      // add primary vertex for Wboson

      if ( primary_vertex )
	Wboson  ->begin_vertex_relations->add   ( primary_vertex, Wboson );

      // add W decay vertex to Wboson, lepton, and neutrino

      Wboson  ->  end_vertex_relations->add( decay_vertex, Wboson   );
      lepton  ->begin_vertex_relations->add( decay_vertex, lepton   );
      neutrino->begin_vertex_relations->add( decay_vertex, neutrino );
    }
  else
    {
      // remove temporary lepton entry and vertex entry in user record
      event_interpret_in->remove( lepton_name );
      event_interpret_in->remove( vertex_name );

      cout << "makeW: lepton not found in W2 event interpretation" << endl;
      return number_of_W;
    }

  // remove temporary lepton and vertex entries in user record
      
  event_interpret_in->remove( lepton_name );
  event_interpret_in->remove( vertex_name );

  return number_of_W;
}

//_____________________________________________________________________________