/// \file
/// \ingroup tutorial_eve
/// Demonstrates usage of TEveTrackPRopagator with different magnetic
/// field configurations.
/// Needs to be run in compiled mode.
/// root
/// ~~~{.cpp}
///   .L track.C+
///   track(3, kTRUE)
/// ~~~
/// void track(Int_t mode = 5, Bool_t isRungeKutta = kTRUE)
/// Modes are
///  0. B = 0, no difference between signed and charge particles;
///  1. constant B field (along z, but could have arbitrary direction);
///  2. variable B field, sign change at  R = 200 cm;
///  3. magnetic field with a zero-field region;
///  4. CMS magnetic field - simple track;
///  5. CMS magnetic field - track with different path-marks.
///  6. Conceptual ILC detector, problematic track
///
/// \image html eve_track.png
/// \macro_code
///
/// \author Alja Mrak-Tadel


#include "TEveTrackPropagator.h"
#include "TEveTrack.h"
#include "TEveVSDStructs.h"
#include "TEveManager.h"
#include "TEveViewer.h"
#include "TSystem.h"
#include "TGLViewer.h"
#include "TMath.h"

#include "TEveViewer.h"
#include "TEvePointSet.h"

#include <iostream>

TEveTrackPropagator* g_prop = nullptr;

class GappedField : public TEveMagField
{
public:
   GappedField():TEveMagField(){}
   ~GappedField() override{};

   virtual Double_t    GetMaxFieldMagD() { return 4; }

   TEveVectorD GetFieldD(Double_t /*x*/, Double_t /*y*/, Double_t z) const override
   {
      if (TMath::Abs(z) < 300) return TEveVectorD(0, 0, -4);
      if (TMath::Abs(z) < 600) return TEveVectorD(0, 0, 0);
      return TEveVectorD(0, 0, 4);
   }
};

//==============================================================================

class CmsMagField: public TEveMagField
{
   bool m_magnetIsOn;
   bool m_reverse;
   bool m_simpleModel;

public:
   CmsMagField():
      m_magnetIsOn(true),
      m_reverse(false),
      m_simpleModel(true){}

   ~CmsMagField() override{}

   void setMagnetState( bool state )
   {
      if (state != m_magnetIsOn)
      {
         if ( state )
            std::cout << "Magnet state is changed to ON" << std::endl;
         else
            std::cout << "Magnet state is changed to OFF" << std::endl;
      }
      m_magnetIsOn = state;
   }

   bool isMagnetOn() const               { return m_magnetIsOn;}
   void setReverseState(bool state)      { m_reverse = state; }
   bool isReverse() const                { return m_reverse;}
   void setSimpleModel(bool simpleModel) { m_simpleModel = simpleModel; }
   bool isSimpleModel() const            { return m_simpleModel;}

   Double_t GetMaxFieldMagD() const override { return m_magnetIsOn ? 3.8 : 0.0; }

   TEveVectorD GetFieldD(Double_t x, Double_t y, Double_t z) const override
   {
      double R = sqrt(x*x+y*y);
      double field = m_reverse ? -GetMaxFieldMagD() : GetMaxFieldMagD();
      //barrel
      if ( TMath::Abs(z) < 724 )
      {
         //inside solenoid
         if ( R < 300) return TEveVectorD(0,0,field);
         // outside solenoid
         if ( m_simpleModel ||
              ( R>461.0 && R<490.5 ) ||
              ( R>534.5 && R<597.5 ) ||
              ( R>637.0 && R<700.0 ) )
            return TEveVectorD(0,0,-field/3.8*1.2);

      } else {
         // endcaps
         if (m_simpleModel)
         {
            if ( R < 50 ) return TEveVectorD(0,0,field);
            if ( z > 0 )
               return TEveVectorD(x/R*field/3.8*2.0, y/R*field/3.8*2.0, 0);
            else
               return TEveVectorD(-x/R*field/3.8*2.0, -y/R*field/3.8*2.0, 0);
         }
         // proper model
         if ( ( TMath::Abs(z)>724 && TMath::Abs(z)<786  ) ||
              ( TMath::Abs(z)>850 && TMath::Abs(z)<910  ) ||
              ( TMath::Abs(z)>975 && TMath::Abs(z)<1003 ) )
         {
            if ( z > 0 )
               return TEveVectorD(x/R*field/3.8*2.0, y/R*field/3.8*2.0, 0);
            else
               return TEveVectorD(-x/R*field/3.8*2.0, -y/R*field/3.8*2.0, 0);
         }
      }
      return TEveVectorD(0,0,0);
   }
};


//==============================================================================
//==============================================================================

//______________________________________________________________________________
TEveTrack* make_track(TEveTrackPropagator* prop, Int_t sign)
{
  // Make track with given propagator.
  // Add to math-marks to test fit.

  auto rc = new TEveRecTrackD();
  rc->fV.Set(0.028558, -0.000918, 3.691919);
  rc->fP.Set(0.767095, -2.400006, -0.313103);
  rc->fSign = sign;

  auto track = new TEveTrack(rc, prop);
  track->SetName(Form("Charge %d", sign));
  // daughter 0
  auto pm1 = new TEvePathMarkD(TEvePathMarkD::kDaughter);
  pm1->fV.Set(1.479084, -4.370661, 3.119761);
  track->AddPathMark(*pm1);
  // daughter 1
  auto pm2 = new TEvePathMarkD(TEvePathMarkD::kDaughter);
  pm2->fV.Set(57.72345, -89.77011, -9.783746);
  track->AddPathMark(*pm2);

  return track;
}


void track(Int_t mode = 1, Bool_t isRungeKutta = kTRUE)
{

   gSystem->IgnoreSignal(kSigSegmentationViolation, true);
   TEveManager::Create();

   auto list = new TEveTrackList();
   auto prop = g_prop = list->GetPropagator();
   prop->SetFitDaughters(kFALSE);
   prop->SetMaxZ(1000);

   if (isRungeKutta)
   {
      prop->SetStepper(TEveTrackPropagator::kRungeKutta);
      list->SetName("RK Propagator");
   }
   else
   {
      list->SetName("Heix Propagator");
   }

   TEveTrack *track = nullptr;
   switch (mode)
   {
      case 0:
      {
         // B = 0 no difference between signed and charge particles
         prop->SetMagField(0.);
         list->SetElementName(Form("%s, zeroB", list->GetElementName()));
         track = make_track(prop, 1);
         break;
      }

      case 1:
      {
         // constant B field, const angle
         prop->SetMagFieldObj(new TEveMagFieldConst(0., 0., -3.8));
         list->SetElementName(Form("%s, constB", list->GetElementName()));
         track = make_track(prop, 1);
         break;
      }

      case 2:
      {
         // variable B field, sign change at  R = 200 cm
         prop->SetMagFieldObj(new TEveMagFieldDuo(200, -4.4, 2));
         list->SetElementName(Form("%s, duoB", list->GetElementName()));
         track = make_track(prop, 1);
         break;
      }

      case 3:
      {
         // gapped field
         prop->SetMagFieldObj(new GappedField());
         list->SetElementName(Form("%s, gappedB", list->GetElementName()));


         auto rc = new TEveRecTrackD();
         rc->fV.Set(0.028558, -0.000918, 3.691919);
         rc->fP.Set(0.767095, -0.400006, 2.313103);
         rc->fSign = 1;
         track = new TEveTrack(rc, prop);

         auto marker = new TEvePointSet(2);
         marker->SetElementName("B field break points");
         marker->SetPoint(0, 0., 0., 300.f);
         marker->SetPoint(1, 0., 0., 600.f);
         marker->SetMarkerColor(3);
         gEve->AddElement(marker);
         break;
      }

      case 4:
      {
         // Magnetic field of CMS I.
         auto mf = new CmsMagField;
         mf->setReverseState(true);

         prop->SetMagFieldObj(mf);
         prop->SetMaxR(1000);
         prop->SetMaxZ(1000);
         prop->SetRnrDaughters(kTRUE);
         prop->SetRnrDecay(kTRUE);
         prop->RefPMAtt().SetMarkerStyle(4);
         list->SetElementName(Form("%s, CMS field", list->GetElementName()));


         auto rc = new TEveRecTrackD();
         rc->fV.Set(0.027667, 0.007919, 0.895964);
         rc->fP.Set(3.903134, 2.252232, -3.731366);
         rc->fSign = -1;
         track = new TEveTrack(rc, prop);

         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(3.576755e+00, 2.080579e+00, -2.507230e+00)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(8.440379e+01, 6.548286e+01, -8.788129e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.841321e+02, 3.915693e+02, -3.843072e+02)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.946167e+02, 4.793932e+02, -4.615060e+02)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDecay,
                  TEveVectorD(2.249656e+02, 5.835767e+02, -5.565275e+02)));

         track->SetRnrPoints(kTRUE);
         track->SetMarkerStyle(4);

         break;
      }

      case 5:
      {
         // Magnetic field of CMS I.
         auto mf = new CmsMagField;
         mf->setReverseState(true);
         mf->setSimpleModel(false);

         prop->SetMagFieldObj(mf);
         prop->SetMaxR(1000);
         prop->SetMaxZ(1000);
         prop->SetRnrReferences(kTRUE);
         prop->SetRnrDaughters(kTRUE);
         prop->SetRnrDecay(kTRUE);
         prop->RefPMAtt().SetMarkerStyle(4);
         list->SetElementName(Form("%s, CMS field", list->GetElementName()));

         auto rc = new TEveRecTrackD();
         rc->fV.Set(-16.426592, 16.403185, -19.782692);
         rc->fP.Set(3.631100, 3.643450, 0.682254);
         rc->fSign = -1;
         track = new TEveTrack(rc, prop);

         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kReference,
                  TEveVectorD(-1.642659e+01, 1.640318e+01, -1.978269e+01),
                  TEveVectorD(3.631100, 3.643450, 0.682254)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kReference,
                  TEveVectorD(-1.859987e+00, 3.172243e+01, -1.697866e+01),
                  TEveVectorD(3.456056, 3.809894, 0.682254)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kReference,
                  TEveVectorD(4.847579e+01, 9.871711e+01, -5.835719e+00),
                  TEveVectorD(2.711614, 4.409945, 0.687656)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.342045e+02, 4.203950e+02, 3.846268e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.483827e+02, 5.124750e+02, 5.064311e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.674676e+02, 6.167731e+02, 6.517403e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDecay,
                  TEveVectorD(1.884976e+02, 7.202000e+02, 7.919290e+01)));

         track->SetRnrPoints(kTRUE);
         track->SetMarkerStyle(4);

         break;
      }

      case 6:
      {
         // Problematic track from Druid
         prop->SetMagFieldObj(new TEveMagFieldDuo(350, -3.5, 2.0));
         prop->SetMaxR(1000);
         prop->SetMaxZ(1000);
         prop->SetRnrReferences(kTRUE);
         prop->SetRnrDaughters(kTRUE);
         prop->SetRnrDecay(kTRUE);
         prop->RefPMAtt().SetMarkerStyle(4);
         list->SetElementName(Form("%s, Some ILC Detector field",
                                   list->GetElementName()));

         auto rc = new TEveRecTrackD();
         rc->fV.Set(57.1068, 31.2401, -7.07629);
         rc->fP.Set(4.82895, 2.35083, -0.611757);
         rc->fSign = 1;
         track = new TEveTrack(rc, prop);

         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(1.692235e+02, 7.047929e+01, -2.064785e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDaughter,
                  TEveVectorD(5.806180e+02, 6.990633e+01, -6.450000e+01)));
         track->AddPathMark(TEvePathMarkD(TEvePathMarkD::kDecay,
                  TEveVectorD(6.527213e+02, 1.473249e+02, -8.348498e+01)));

         track->SetRnrPoints(kTRUE);
         track->SetMarkerStyle(4);

         break;
      }
   };

   if (isRungeKutta)
      list->SetLineColor(kMagenta);
   else
      list->SetLineColor(kCyan);

   track->SetLineColor(list->GetLineColor());

   gEve->AddElement(list);
   list->AddElement(track);

   track->MakeTrack();

   TEveViewer *ev = gEve->GetDefaultViewer();
   TGLViewer  *gv = ev->GetGLViewer();
   gv->SetGuideState(TGLUtil::kAxesOrigin, kTRUE, kFALSE, nullptr);

   gEve->Redraw3D(kTRUE);
   gSystem->ProcessEvents();

   gv->CurrentCamera().RotateRad(-0.5, 1.4);
   gv->RequestDraw();
}