// etabench
// Copyright (C) 2022 xaizek <xaizek@posteo.net>
//
// This file is part of etabench.
//
// etabench is free software: you can redistribute it and/or modify
// it under the terms of version 3 of the GNU General Public License
// as published by the Free Software Foundation.
//
// etabench is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with etabench.  If not, see <https://www.gnu.org/licenses/>.

#include "algs.hpp"

#include <cmath>

#include <map>
#include <memory>
#include <utility>
#include <vector>

#include <fmt/core.h>

#include "utils/float.hpp"
#include "core.hpp"

struct FactorWeight
{
    float factor;

    FactorWeight(float factor) : factor(factor)
    { }

    float operator()(int idx, float lastWeight) const
    {
        return (idx == 0 ? 1.0f : lastWeight*factor);
    }
};

template <float factor>
float
factorWeight(int idx, float lastWeight)
{ return FactorWeight(factor)(idx, lastWeight); }

static float
indexWeight(int idx, float lastWeight)
{ return 1 + idx; }

int
AverageAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    float speed = float(current)/time;
    return (total - current)/speed;
}

void
ImmediateAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
}

int
ImmediateAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    float speed = float(current - lastProgress)/(time - lastTime);
    lastProgress = current;
    lastTime = time;
    return (speed == 0 ? 0 : (total - current)/speed);
}

CombinedAlg::CombinedAlg(std::vector<std::unique_ptr<EtaAlg>> algs) :
    algs(std::move(algs))
{ }

std::string
CombinedAlg::getName() const
{
    std::string combined = "Combined[ ";
    for (auto &alg : algs) {
        combined += alg->getName();
        combined += ' ';
    }
    combined += ']';
    return combined;
}

void
CombinedAlg::reset() {
    for (auto &alg : algs) {
        alg->reset();
    }
}

int
CombinedAlg::estimate(int current, int total, int time)
{
    int eta = 0;
    for (auto &alg : algs) {
        eta += alg->estimate(current, total, time);
    }

    eta /= algs.size();
    return eta;
}

void
ImmChangeLimitAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    lastSpeed = 0.0f;
}

int
ImmChangeLimitAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    float speed = float(current - lastProgress)/(time - lastTime);
    lastProgress = current;
    lastTime = time;

    if (lastSpeed < 0.05f) {
        lastSpeed = speed;
    } else {
        float p = 0.02f;
        lastSpeed = std::clamp(speed*1.0f,
                               lastSpeed*(1 - p),
                               lastSpeed*(1 + p));
    }
    return (total - current)/lastSpeed;
}

void
AveChangeLimitAlg::reset()
{
    lastSpeed = 0.0f;
}

int
AveChangeLimitAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    float speed = float(current)/time;
    if (lastSpeed < 0.05f) {
        lastSpeed = speed;
    } else {
        float p = 0.02f;
        lastSpeed = std::clamp(speed*1.0f,
                               lastSpeed*(1 - p),
                               lastSpeed*(1 + p));
    }
    return (total - current)/lastSpeed;
}

LookBackAlg::LookBackAlg(int delay) : delay(delay)
{ }

void
LookBackAlg::reset()
{
    points.clear();
}

int
LookBackAlg::estimate(int current, int total, int time)
{
    const int cutOff = time - delay;
    while (!points.empty() && points.front().time < cutOff) {
        points.erase(points.begin());
    }

    points.emplace_back(time, current);

    if (points.size() == 1) {
        if (current == 0 || time == 0) {
            return 0;
        }
        return (total - current)/(1.0f*current/time);
    }

    const Point &p = points.front();
    float speed = float(current - p.current)/(time - p.time);
    return (speed == 0 ? 0 : (total - current)/speed);
}

AccelerationAlg::AccelerationAlg() :
    speeds(10, &factorWeight<2.0f>), etas(10, &indexWeight)
{ }

void
AccelerationAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    lastEta = -1;
    lastV = 0.0f;
    speeds.clear();
    etas.clear();
}

// s = t*v0 + a*t*t/2
//
// (a/2)*t*t + v0*t - s = 0
// d = v0*v0 + 2*a*s
// positive root = (sqrt(d) - v0)/a
int
AccelerationAlg::estimate(int current, int total, int time)
{
    if (time == 0) {
        return 0;
    }

    // Immediate speed, used as initial speed after applying damping.
    float v = float(current - lastProgress)/(time - lastTime);

    // We store actual speed, but use averaged one.
    v = speeds.update(v);

    // Second round of adjusting speed to not deviate from previous value too
    // much.
    if (v != 0 && lastV != 0) {
        auto [minV, maxV] = std::minmax(v, lastV);
        v = lastV + (v - lastV)*minV/maxV;
    }

    // Acceleration.
    float a = (v - lastV)/(lastTime == 0 ? 1 : time - lastTime);
    // Remains this much.
    int r = total - current;

    int eta;
    if (!floatEq(a, 0.0f) && v*v + 2*a*r > 0) {
        eta = (std::sqrt(v*v + 2*a*r) - v)/a;
    } else if (!floatEq(v, 0.0f)) {
        eta = r/v;
    } else {
        eta = lastEta;
    }

    // The result is also smoothed to avoid jumps.
    eta = etas.update(eta);

    lastProgress = current;
    lastTime = time;
    lastV = v;
    lastEta = eta;

    return eta;
}

void
SwitchAlg::reset()
{
    history.clear();
    lastProgress = 0;
}

int
SwitchAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    int progress = current*100/total;
    if (progress != lastProgress) {
        history.emplace(progress, time);
    }

    if (progress <= 50) {
        float speed = float(current)/time;
        return (total - current)/speed;
    }

    int from = progress - (100 - progress);
    auto it = history.lower_bound(from);
    if (it != history.end() && it->first != from && it != history.begin()) {
        std::advance(it, -1);
    }
    return time - it->second;
}

GravityAlg::GravityAlg() : speeds(-1, &indexWeight)
{ }

void
GravityAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    avgProgressRate = 0;
    avgSeekRate = 0;
    speeds.clear();
}

int
GravityAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    int lastTimeElapsed = time - lastTime;

    float v = speeds.update(float(current - lastProgress)/lastTimeElapsed);

    const float mass = (lastTime == 0 ? 1.0f : 4.0f);
    const float maxSpeedCoeff = 0.25f;
    const float frictionCoeff = 0.75f;
    const float dampingCoeff = 0.25f;

    // `lastTime == 0` checks are to work around jumps at the start.

    // Lose 25% of our speed per sec, simulating friction.
    avgSeekRate *= std::pow(frictionCoeff, lastTimeElapsed);

    float delta = v - avgProgressRate;

    // Update the velocity.
    float oldSpeed = avgSeekRate;
    float accel = delta/mass;
    avgSeekRate += accel*lastTimeElapsed;  // v += at

    // Clamp the top speed to 25% of our current value.
    if (avgProgressRate > 0) {
        float maxVal = avgProgressRate*maxSpeedCoeff;
        if (std::abs(avgSeekRate) > maxVal) {
            float sign = (avgSeekRate > 0.0f ? 1.0f : -1.0f);
            avgSeekRate = sign*maxVal;
        }
    }

    // Make sure they have the same sign before doing up update.
    if ((avgSeekRate > 0.0f) == (delta > 0.0f)) {
        float f = (lastTime == 0 ? 1 : 2);
        float adjust = (oldSpeed + avgSeekRate)/f*lastTimeElapsed;

        // Don't overshoot.
        if (std::abs(adjust) > std::abs(delta)) {
            adjust = delta;
            // Apply damping.
            avgSeekRate *= dampingCoeff;
        }

        avgProgressRate += adjust;
    }

    lastTime = time;
    lastProgress = current;
    return (total - current)/avgProgressRate;
}

void
FirefoxAlg::reset()
{
    lastEta = 0;
    lastProgress = 0;
    lastTime = 0;
    speed = 0;
}

int
FirefoxAlg::estimate(int current, int total, int time)
{
    float v = float(current - lastProgress)/(time - lastTime);
    if (lastTime == 0)
        speed = v;
    else
        speed = v*0.1f + speed*0.9f;

    float eta = (total - current)/speed;

    // Apply hysteresis to favor downward over upward swings 30% of down and 10%
    // of up (exponential smoothing).
    float etaDiff = eta - lastEta;
    eta = lastEta + (lastTime == 0 ? 1.0f : etaDiff < 0 ? 0.3f : 0.1f)*etaDiff;

    // If the new ETA is similar, reuse something close to the last ETA, but
    // subtract a little to provide forward progress.
    float diffPercent = (etaDiff/lastEta)*100;
    if (std::abs(etaDiff) < 5 || std::abs(diffPercent) < 5) {
        eta = lastEta - (etaDiff < 0 ? 0.4 : 0.2);
    }

    lastEta = eta;
    lastProgress = current;
    lastTime = time;

    return eta;
}

void
SmoothingAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    lastEta = 0;
    speed = 0;
}

int
SmoothingAlg::estimate(int current, int total, int time)
{
    float v = float(current - lastProgress)/(time - lastTime);
    float a = alpha/100.0f;

    if (lastTime == 0)
        speed = v;
    else
        speed = v*a + speed*(1.0f - a);

    int eta = (floatEq<0.1f>(speed, 0) ? lastEta : (total - current)/speed);

    lastProgress = current;
    lastTime = time;
    lastEta = eta;

    return eta;
}

void
SlownessAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    totalEta = 0;
    lastEta = 0;
}

int
SlownessAlg::estimate(int current, int total, int time)
{
    if (current == lastProgress) {
        return lastEta + time - lastTime;
    }

    float dt = float(time - lastTime)/(current - lastProgress);

    // Slowness is the amount of time per unit of progress.
    float slowness = total*dt;

    float eta;
    if (lastProgress == 0) {
        // If first invocation, initialize the estimate.
        totalEta = slowness;
        eta = slowness;
    } else {
        float weight = std::exp(-1.0f/(total*decay/100.f));
        totalEta = totalEta*weight + slowness*(1.0f - weight);
        eta = (1.0f - 1.0f*current/total)*totalEta;
    }

    lastProgress = current;
    lastTime = time;
    lastEta = eta;

    return eta;
}

WindowAlg::WindowAlg(int windowSize, float alpha) :
    windowSize(windowSize),
    alpha(alpha),
    speeds(windowSize, FactorWeight(alpha))
{ }

std::string
WindowAlg::getName() const
{ return fmt::format("Window {} {:.2}", windowSize, alpha); }

void
WindowAlg::reset()
{
    lastTime = 0;
    lastProgress = 0;
    lastEta = 0;
    speeds.clear();
}

int
WindowAlg::estimate(int current, int total, int time)
{
    if (time == 0 || current == 0) {
        return 0;
    }

    float v = speeds.update(float(current - lastProgress)/(time - lastTime));
    if (floatEq(v, 0)) {
        return lastEta + time - lastTime;
    }

    lastTime = time;
    lastProgress = current;
    lastEta = (total - current)/v;

    return lastEta;
}

void
ExponentialAlg::reset()
{
    lastProgress = 0;
    lastTime = 0;
    rateEst = 0;
    lastEta = 0;
}

int
ExponentialAlg::estimate(int current, int total, int time)
{
    if (current == lastProgress) {
        return lastEta + time - lastTime;
    }

    float dt = float(time - lastTime)/(current - lastProgress);

    // Rate equals normalized progress per unit of time.
    float rate = 1/(total*dt);

    float eta = 0;
    if (lastProgress == 0) {
        // If first invocation, initialize the estimate.
        rateEst = rate;
    } else {
        float weight = std::exp(-dt/decay);
        rateEst = rateEst*weight + rate*(1.0f - weight);
    }

    eta = (1.0f - 1.0f*current/total)/rateEst;

    lastProgress = current;
    lastTime = time;
    lastEta = eta;

    return eta;
}