/*
 * Copyright 2020 Anton Khirnov <anton@khirnov.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef MG2D_STEP_CONTROL_H
#define MG2D_STEP_CONTROL_H

#include <float.h>

#include "common.h"

#define SC_HIST_SIZE 16

typedef struct StepControl {
    /* estimated value for a good step */
    double hint;
    /* minimum step size, fail if this value is reached */
    double step_min;

    int steps_since_init;

    double step[SC_HIST_SIZE];
    double fact[SC_HIST_SIZE];

    double conv_epsilon;

    int hist_len;
    int idx_bracket;
    int idx_diverge;

    int diverge_count;
} StepControl;

static void sci_hist_clear(StepControl *sc)
{
    for (int i = 0; i < SC_HIST_SIZE; i++) {
        sc->step[i] = DBL_MAX;
        sc->fact[i] = DBL_MAX;
    }

    sc->hist_len    = 0;
    sc->idx_bracket = -1;
    sc->idx_diverge = -1;
}

static void sc_reset(StepControl *sc)
{
    sci_hist_clear(sc);
    sc->hint      = DBL_MAX;
    sc->step_min  = DBL_MAX;
    sc->steps_since_init = 0;
}

static void sc_init(StepControl *sc, double hint, double step_min)
{
    sc->hint             = hint;
    sc->step_min         = step_min;
    sc->steps_since_init = 0;
    sc->conv_epsilon     = hint * 1e-3;
}

static double sc_step_get(StepControl *sc)
{
    // no history present, try the hint
    if (!sc->hist_len)
        return (sc->hint == DBL_MAX) ? 0.25 : sc->hint;

    // don't have the bracket yet
    if (sc->idx_bracket < 0) {
        const double high_step = sc->step[sc->hist_len - 1];
        const double high_fact = sc->fact[sc->hist_len - 1];

        // don't know where divergence happens
        // -> try a higher step
        if (high_fact > 1.0)
            return high_step * 1.2;

        // convergence factor should decrease towards smaller timesteps
        // so just try lower ones until we get a bracket
        return MAX(sc->step_min, sc->step[0] * 0.8);
    }

    // TODO periodic random steps to try escaping local minima

    // got a bracket, golden-section-search it until convergence
    mg2di_assert(sc->hist_len - sc->idx_bracket >= 3);
    mg2di_assert(sc->fact[sc->idx_bracket]     < sc->fact[sc->idx_bracket + 1]);
    mg2di_assert(sc->fact[sc->idx_bracket + 2] < sc->fact[sc->idx_bracket + 1]);
    {
        const double dist0 = sc->step[sc->idx_bracket + 1] - sc->step[sc->idx_bracket];
        const double dist1 = sc->step[sc->idx_bracket + 2] - sc->step[sc->idx_bracket + 1];

        // converged
        if (dist0 < sc->conv_epsilon && dist1 < sc->conv_epsilon)
            return sc->step[sc->idx_bracket + 1];

        if (dist0 > dist1)
            return sc->step[sc->idx_bracket] + 0.61803 * dist0;
        else
            return sc->step[sc->idx_bracket + 2] - 0.61803 * dist1;
    }
}

static void sci_find_bracket(StepControl *sc)
{
    for (int i = sc->hist_len - 2; i >= 0; i--) {
        if (sc->fact[i]    < sc->fact[i + 1] &&
            sc->fact[i + 2] < sc->fact[i + 1]) {
            sc->idx_bracket = i;
            return;
        }
    }
}

static int sc_step_confirm(StepControl *sc, const double step,
                           const double norm_old, const double norm_new)
{
    const double conv_fact = norm_old / norm_new;

    // signal failure if we reached minimum step
    if (step <= sc->step_min)
        return -EINVAL;

    // ignore the first relax step, it tends to be weird
    sc->steps_since_init++;
    //if (sc->steps_since_init < 2)
    //    return 0;

    // divergence
    if (conv_fact <= 1.0) {
        //if (step < sc->hint / 2.)
        //    return 0;

        // if the step is not larger then the largest known diverging step,
        // add it into history
        if (!sc->hist_len                    ||
            sc->idx_diverge < 0              ||
            sc->step[sc->idx_diverge] > step) {
            int idx = sc->idx_diverge >= 0 ? sc->idx_diverge : sc->hist_len;

            // history size overflow, shouldn't happen
            if (idx >= SC_HIST_SIZE)
                return -ENOMEM;

            // largest known coverging step doesn't converge anymore,
            // clear history
            if (idx > 0 && sc->step[idx - 1] <= step) {
                sci_hist_clear(sc);
                idx = 0;
            }

            sc->step[idx]   = step;
            sc->fact[idx]   = conv_fact;
            sc->idx_diverge = idx;
            sc->hist_len    = idx + 1;

            // check if we now have a bracket
            if (sc->idx_bracket < 0 && sc->hist_len >= 3 &&
                sc->fact[idx - 1] > conv_fact &&
                sc->fact[idx - 1] > sc->fact[idx - 2]) {
                sc->idx_bracket = idx - 2;
            }

        }

        return 1;
    }

    // convergence

    // history empty, add first element
    if (!sc->hist_len) {
        sc->step[0]  = step;
        sc->fact[0]  = conv_fact;
        sc->hist_len = 1;
        return 0;
    }

    // previously diverging step is now converging
    // drop all diverging steps from history
    if (sc->idx_diverge >= 0 &&
        step >= sc->step[sc->idx_diverge]) {
        const int idx = sc->idx_diverge;

        sc->idx_diverge = -1;
        sc->step[idx]   = step;
        sc->fact[idx]   = conv_fact;
        sc->hist_len    = idx + 1;

        // check if we just broke bracketing
        if (sc->idx_bracket >= 0 && sc->idx_bracket + 2 == idx &&
            sc->fact[sc->idx_bracket + 1] <= conv_fact)
            sc->idx_bracket = -1;

        // could not have created a bracket where there was none before, since
        // we increased the factor for topmost step

        return 0;
    }

    // no bracket
    if (sc->idx_bracket < 0) {
        int idx_floor = -1;

        for (int i = 0; i < sc->hist_len; i++) {
            if (sc->step[i] < step)
                idx_floor = i;
            else
                break;
        }
        if (idx_floor + 1 < SC_HIST_SIZE) {
            memmove(sc->step + idx_floor + 2, sc->step + idx_floor + 1,
                    sizeof(*sc->step) * (SC_HIST_SIZE - (idx_floor + 2)));
            memmove(sc->fact + idx_floor + 2, sc->fact + idx_floor + 1,
                    sizeof(*sc->fact) * (SC_HIST_SIZE - (idx_floor + 2)));

            sc->step[idx_floor + 1] = step;
            sc->fact[idx_floor + 1] = conv_fact;
            sc->hist_len = MIN(SC_HIST_SIZE, sc->hist_len + 1);
            if (sc->idx_diverge > idx_floor)
                sc->idx_diverge = (sc->idx_diverge + 1 < SC_HIST_SIZE) ? sc->idx_diverge + 1 : -1;
        }
        sci_find_bracket(sc);

        return 0;
    }

    // have bracket

    // got a step outside of the bracket for some reason
    if (step <= sc->step[sc->idx_bracket] ||
        step >= sc->step[sc->idx_bracket + 2])
        return 0;

    if (step < sc->step[sc->idx_bracket + 1]) {
        // step inside lower interval
        if (conv_fact > sc->fact[sc->idx_bracket + 1]) {
            // replace upper bound
            sc->step[sc->idx_bracket + 2] = sc->step[sc->idx_bracket + 1];
            sc->fact[sc->idx_bracket + 2] = sc->fact[sc->idx_bracket + 1];
            if (sc->idx_diverge == sc->idx_bracket + 2)
                sc->idx_diverge = -1;

            sc->step[sc->idx_bracket + 1] = step;
            sc->fact[sc->idx_bracket + 1] = conv_fact;
        } else {
            // replace lower bound
            sc->step[sc->idx_bracket] = step;
            sc->fact[sc->idx_bracket] = conv_fact;
        }
    } else {
        // step inside upper interval
        if (conv_fact > sc->fact[sc->idx_bracket + 1]) {
            // replace lower bound
            sc->step[sc->idx_bracket] = sc->step[sc->idx_bracket + 1];
            sc->fact[sc->idx_bracket] = sc->fact[sc->idx_bracket + 1];
            sc->step[sc->idx_bracket + 1] = step;
            sc->fact[sc->idx_bracket + 1] = conv_fact;
        } else {
            sc->step[sc->idx_bracket + 2] = step;
            sc->fact[sc->idx_bracket + 2] = conv_fact;
            if (sc->idx_diverge == sc->idx_bracket + 2)
                sc->idx_diverge = -1;
        }
    }

    return 0;
}

#endif // MG2D_STEP_CONTROL_H