//
// ********************************************************************
// * License and Disclaimer                                           *
// *                                                                  *
// * The  Geant4 software  is  copyright of the Copyright Holders  of *
// * the Geant4 Collaboration.  It is provided  under  the terms  and *
// * conditions of the Geant4 Software License,  included in the file *
// * LICENSE and available at  http://cern.ch/geant4/license .  These *
// * include a list of copyright holders.                             *
// *                                                                  *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work  make  any representation or  warranty, express or implied, *
// * regarding  this  software system or assume any liability for its *
// * use.  Please see the license in the file  LICENSE  and URL above *
// * for the full disclaimer and the limitation of liability.         *
// *                                                                  *
// * This  code  implementation is the result of  the  scientific and *
// * technical work of the GEANT4 collaboration.                      *
// * By using,  copying,  modifying or  distributing the software (or *
// * any work based  on the software)  you  agree  to acknowledge its *
// * use  in  resulting  scientific  publications,  and indicate your *
// * acceptance of all terms of the Geant4 Software license.          *
// ********************************************************************
//
//
//
//
// class G4InterpolationDriver
//
// Class description:
//
// Driver class which uses Runge-Kutta stepper with interpolation property
// to integrate track with error control 

// History:
// - Created. D.Sorokin
// --------------------------------------------------------------------


#include "G4FieldUtils.hh"
#include "G4LineSection.hh"
#include "G4Exception.hh"

#include <algorithm>


template <class T>
G4InterpolationDriver<T>::
G4InterpolationDriver ( G4double hminimum, T* pStepper,
                      G4int numComponents, G4int statisticsVerbose )
    : G4RKIntegrationDriver<T>(pStepper),
      fMinimumStep(hminimum),
      fVerboseLevel(statisticsVerbose)
{
    if (numComponents != Base::GetStepper()->GetNumberOfVariables())
    {
        std::ostringstream message;
        message << "Driver's number of integrated components "
                << numComponents
                << " != Stepper's number of components "
                << pStepper->GetNumberOfVariables();
        G4Exception("G4InterpolationDriver","GeomField0002",
                    FatalException, message);
    }

    for (G4int i = 0; i < Base::GetMaxNoSteps(); ++i) 
    {
        fSteppers.push_back({
            std::unique_ptr<T>(new T(
                pStepper->GetEquationOfMotion(), 
                pStepper->GetNumberOfVariables())
            ),
            DBL_MAX,
            -DBL_MAX,
            0.0
        });
    }

    fLastStepper = fSteppers.end();
}

template <class T>
void G4InterpolationDriver<T>::
Interpolate(G4double curveLength, field_utils::State& y) const
{
    if (fLastStepper == fSteppers.end()) 
    {
        std::ostringstream message;
        message << "LOGICK ERROR: fLastStepper == end";
        G4Exception("G4InterpolationDriver::Interpolate()", 
                    "GeomField1001", FatalException, message);
        return;
    }

    ConstStepperIterator end = fLastStepper + 1;

    auto it = std::lower_bound(
            fSteppers.cbegin(), end, curveLength, 
            [](const InterpStepper& stepper, G4double value) {
                return stepper.end < value;
            }
        );

    if (it == end)
    {
        if (curveLength - fLastStepper->end > CLHEP::perMillion * mm)
        {
            std::ostringstream message;
            message << "curveLength = " << curveLength << " > " << fLastStepper->end;
            G4Exception("G4InterpolationDriver::Interpolate()", 
                        "GeomField1001", JustWarning, message);
        }

        return fLastStepper->stepper->Interpolate(1, y);
    }

    if (curveLength < it->begin) 
    {
        if (it->begin - curveLength > CLHEP::perMillion * mm)
        {
            std::ostringstream message;
            message << "curveLength = " << curveLength << " < " << it->begin;
            G4Exception("G4InterpolationDriver::Interpolate()", 
                        "GeomField1001", JustWarning, message);
        }

        return it->stepper->Interpolate(0, y);
    }

    return InterpolateImpl(curveLength, it, y);
}

template <class T>
void G4InterpolationDriver<T>::
InterpolateImpl(G4double curveLength, 
                ConstStepperIterator it, 
                field_utils::State& y) const
{
    const G4double tau = (curveLength - it->begin) * it->inverseLength;
    return it->stepper->Interpolate(field_utils::clamp(tau, 0., 1.), y);  
}

template <class T>
G4double G4InterpolationDriver<T>::
DistChord(const field_utils::State& yBegin, 
          G4double curveLengthBegin, 
          const field_utils::State& yEnd, 
          G4double curveLengthEnd) const
{
    // optimization check if it worth
    if (curveLengthBegin == fLastStepper->begin && 
        curveLengthEnd == fLastStepper->end) 
    {
        return fLastStepper->stepper->DistChord();
    }

    const G4double curveLengthMid = 0.5 * (curveLengthBegin + curveLengthEnd);
    field_utils::State yMid;
    
    Interpolate(curveLengthMid, yMid);

    return G4LineSection::Distline(
        field_utils::makeVector(yMid, field_utils::Value3D::Position),
        field_utils::makeVector(yBegin, field_utils::Value3D::Position),
        field_utils::makeVector(yEnd, field_utils::Value3D::Position)
    );
}

template <class T>
G4double G4InterpolationDriver<T>::
AdvanceChordLimited(G4FieldTrack& track,
                    G4double hstep,
                    G4double epsStep,
                    G4double chordDistance)
{
    ++fTotalStepsForTrack;

    const G4double curveLengthBegin = track.GetCurveLength();
    G4double hdid = 0;
    auto it = fSteppers.begin();
    G4double dChordStep = 0;

    field_utils::State yBegin, y;
    track.DumpToArray(yBegin);
    track.DumpToArray(y);

    if (fFirstStep) 
    {
        Base::GetEquationOfMotion()->RightHandSide(y, fdydx);
        fFirstStep = false;
    }

    if (fKeepLastStepper) 
    {
        std::swap(*fSteppers.begin(), *fLastStepper);
        it = fSteppers.begin(); //new begin, update iterator
        fLastStepper = it;
        hdid = it->end - curveLengthBegin;
        field_utils::copy(y, it->stepper->GetYOut());

        dChordStep = DistChord(
            yBegin, curveLengthBegin, y, curveLengthBegin + hdid
        );

        ++it;
    }

    // accurate advance & check chord distance
    G4double h = fhnext;
    const G4double hend = std::min(hstep * (1 - epsStep), fChordStepEstimate);
    for (; hdid < hend &&
           dChordStep < chordDistance &&
           it != fSteppers.end(); ++it)
    {
        h = std::min(h, hstep - hdid);

        // make one step
        hdid += OneGoodStep(it, y, fdydx, h, epsStep, curveLengthBegin + hdid);

        // update last stepper
        fLastStepper = it;

        // estimate chord distance
        dChordStep = DistChord(
            yBegin, curveLengthBegin, y, curveLengthBegin + hdid
        );
    }

    // update step estimation
    if (h * fTotalStepsForTrack > fSmallestCurveFraction * (curveLengthBegin + hdid) && 
        h > fMinimumStep)
    {
        fhnext = h;
    }

    //CheckState();

    // update chord step estimate
    hdid = FindNextChord(
        yBegin, curveLengthBegin, 
        y, curveLengthBegin + hdid, 
        dChordStep,
        chordDistance
    );

    const G4double curveLengthEnd = curveLengthBegin + hdid;
    fKeepLastStepper = fLastStepper->end - curveLengthEnd > fMinimumStep;
    track.LoadFromArray(y, fLastStepper->stepper->GetNumberOfVariables());
    track.SetCurveLength(curveLengthBegin + hdid);
    return hdid;
}

template <class T>
G4double G4InterpolationDriver<T>::FindNextChord(const field_utils::State& yBegin,
                                                 G4double curveLengthBegin,
                                                 field_utils::State& yEnd,
                                                 G4double curveLengthEnd,
                                                 G4double dChord,
                                                 G4double chordDistance)
{
    G4double hstep = curveLengthEnd - curveLengthBegin;
    G4double curveLength = curveLengthEnd;

    G4int i = 0;
    for (; i < fMaxTrials && dChord > chordDistance; ++i) 
    {
        // crop step size
        hstep = CalcChordStep(hstep, dChord, chordDistance);

        // hstep should be in the last stepper
        hstep = std::max(hstep, fLastStepper->begin - curveLengthBegin);
        curveLength = curveLengthBegin + hstep;

        // use fLastStepper!
        InterpolateImpl(curveLength, fLastStepper, yEnd);

        // update chord distance
        dChord = DistChord(yBegin, curveLengthBegin, yEnd, curveLength);
    }

    // dChord may be zero
    if (dChord > 0.0)
    {
        fChordStepEstimate = hstep * std::sqrt(chordDistance / dChord);
    }

    if (i == fMaxTrials) 
    {
        G4Exception("G4InterpolationDriver::FindNextChord()", 
                    "GeomField1001", JustWarning, "cannot converge");
    }

    return hstep;
}

// Is called to estimate the next step size, even for successful steps,
// in order to predict an accurate 'chord-sensitive' first step
// which is likely to assist in more performant 'stepping'.
template <class T>
G4double G4InterpolationDriver<T>::CalcChordStep(G4double stepTrialOld, 
                                                 G4double dChordStep,
                                                 G4double chordDistance)  
{
    G4double stepTrial;

    if (dChordStep > 0.0)
    {
        fChordStepEstimate =
            stepTrialOld * std::sqrt(chordDistance / dChordStep);
        stepTrial = fFractionNextEstimate * fChordStepEstimate;
    }
    else
    {
        // Should not update the Unconstrained Step estimate: incorrect!
        stepTrial = stepTrialOld * 2.; 
    }

    if (stepTrial <= 0.001 * stepTrialOld)
    {
        if (dChordStep > 1000.0 * chordDistance)
        {
            stepTrial = stepTrialOld * 0.03;   
        }
        else
        {
            if (dChordStep > 100. * chordDistance)
            {
                stepTrial = stepTrialOld * 0.1;   
            }
            else   // Try halving the length until dChordStep OK
            {
                stepTrial = stepTrialOld * 0.5;   
            }
        }
    }
    else if (stepTrial > 1000.0 * stepTrialOld)
    {
        stepTrial = 1000.0 * stepTrialOld;
    }

    if (stepTrial == 0.0)
    {
        stepTrial = 0.000001;
    }

    // A more sophisticated chord-finder could figure out a better
    // stepTrial, from dChordStep and the required d_geometry
    //   e.g.
    //      Calculate R, r_helix (eg at orig point)
    //      if( stepTrial < 2 pi  R )
    //          stepTrial = R arc_cos( 1 - chordDistance / r_helix )
    //      else    
    //          ??

    return stepTrial;
}

template <class T>
G4bool G4InterpolationDriver<T>::
AccurateAdvance(G4FieldTrack& track, G4double hstep,
                G4double /*eps*/, G4double /*hinitial*/)
{
    if (hstep == 0.0)
    {
        std::ostringstream message;
        message << "Proposed step is zero; hstep = " << hstep << " !";
        G4Exception("G4InterpolationDriver::AccurateAdvance()", 
                    "GeomField1001", JustWarning, message);
        return true; 
    }

    if (hstep < 0)
    {
        std::ostringstream message;
        message << "Invalid run condition." << G4endl
                << "Proposed step is negative; hstep = " << hstep << "." << G4endl
                << "Requested step cannot be negative! Aborting event.";
        G4Exception("G4InterpolationDriver::AccurateAdvance()", 
                    "GeomField0003", EventMustBeAborted, message);
        return false;
    }

    const G4double curveLength = track.GetCurveLength();
    const G4double curveLengthEnd = curveLength + hstep;

    field_utils::State y;
    Interpolate(curveLengthEnd, y);

    track.LoadFromArray(y, Base::GetStepper()->GetNumberOfVariables());
    track.SetCurveLength(curveLengthEnd);

    return true;
}

// Driver for one Runge-Kutta Step with monitoring of local truncation error
// to ensure accuracy and adjust stepsize. Input are dependent variable
// array y[0,...,5] and its derivative dydx[0,...,5] at the
// starting value of the independent variable x . Also input are stepsize
// to be attempted htry, and the required accuracy eps. On output y and x
// are replaced by their new values, hdid is the stepsize that was actually
// accomplished, and hnext is the estimated next stepsize. 
// This is similar to the function rkqs from the book:
// Numerical Recipes in C: The Art of Scientific Computing (NRC), Second
// Edition, by William H. Press, Saul A. Teukolsky, William T.
// Vetterling, and Brian P. Flannery (Cambridge University Press 1992),
// 16.2 Adaptive StepSize Control for Runge-Kutta, p. 719
template <class T>
G4double G4InterpolationDriver<T>::OneGoodStep(StepperIterator it,
                                               field_utils::State& y,
                                               field_utils::State& dydx,
                                               G4double& hstep,
                                               G4double epsStep,
                                               G4double curveLength)

{
    G4double error2 = DBL_MAX;
    field_utils::State yerr, ytemp, dydxtemp;
    G4double h = hstep;

    G4int i = 0;
    for (; i < fMaxTrials; ++i)
    {
        it->stepper->Stepper(y, dydx, h, ytemp, yerr, dydxtemp); 
        error2 = field_utils::relativeError2(y, yerr, h, epsStep);

        if (error2 <= 1.0)
        {
            hstep = Base::GrowStepSize2(h, error2);
            break; 
        }

        // don't control error for small steps
        if (h <= fMinimumStep)
        {
            hstep = fMinimumStep;
            break;
        }

        h = std::max(Base::ShrinkStepSize2(h, error2), fMinimumStep);
    }

    if (i == fMaxTrials) 
    {
        G4Exception("G4InterpolationDriver::OneGoodStep()", 
                    "GeomField1001", JustWarning, "cannot converge");
        h = std::max(Base::ShrinkStepSize2(h, error2), fMinimumStep);
    }

    // set interpolation inverval
    it->begin = curveLength;
    it->end = curveLength + h;
    it->inverseLength = 1. / h;

    // setup interpolation
    it->stepper->SetupInterpolation();

    field_utils::copy(dydx, dydxtemp);
    field_utils::copy(y, ytemp);

    return h;
}

template <class T>
void G4InterpolationDriver<T>::PrintState() const
{
    using namespace field_utils;
    State prevEnd, currBegin;
    auto prev = fSteppers.begin();

    G4cout << "====== curr state ========" << G4endl;
    for (auto i = fSteppers.begin(); i <= fLastStepper; ++i) {
        i->stepper->Interpolate(0, currBegin);

        G4cout << "cl_begin: " <<i->begin << " "
               << "cl_end: " << i->end << " ";

        if (prev != i) {
            prev->stepper->Interpolate(1, prevEnd);
            auto prevPos = makeVector(prevEnd, Value3D::Position);
            auto currPos = makeVector(currBegin, Value3D::Position);
            G4cout << "diff_begin: " << (prevPos - currPos).mag();
        }
                   
        G4cout << G4endl;
        prev = i;
    }

    const G4double clBegin = fSteppers.begin()->begin;
    const G4double clEnd = fLastStepper->end;
    const G4double hstep = (clEnd - clBegin) / 10.;
    State yBegin, yCurr;
    Interpolate(0, yBegin);
    for (G4double cl = clBegin; cl <= clEnd + 1e-12; cl += hstep) {
        Interpolate(cl, yCurr);
        auto d = DistChord(yBegin, clBegin, yCurr, cl);
        G4cout << "cl: " << cl << " chord_distance: " << d << G4endl;
    }

    G4cout << "==========================" << G4endl;
}


template <class T>
void G4InterpolationDriver<T>::CheckState() const
{
    G4int smallSteps = 0;
    for (auto i = fSteppers.begin(); i <= fLastStepper; ++i) 
    {
        G4double stepLength = i->end - i->begin;
        if (stepLength < fMinimumStep) {
            ++smallSteps;
        }
    }

    if (smallSteps > 1) 
    {
        std::ostringstream message;
        message << "====== curr state ========\n";
        for (auto i = fSteppers.begin(); i <= fLastStepper; ++i) {
            message << "cl_begin: " <<i->begin << " "
                    << "cl_end: " << i->end << "\n";
        }

        G4Exception("G4InterpolationDriver::CheckState()", 
                    "GeomField0003", FatalException, message);
    }
}