#pragma once
#include <vector>
#include <boost/numeric/odeint.hpp>
//#include "../myodeint/boost/numeric/odeint.hpp"

static constexpr double PI = 3.14159265358979323;

struct well_function
{
    double width = 1;
    double u0 = 1;

    double A = 1, sigma = 0.1, mean = 0;

    double operator() (double x)
    {
        if (x > -width / 2.0 && x < width / 2.0)
        return -u0 + A/sqrt(2*PI*sigma) * exp(-(x-mean)*(x - mean) / (2*sigma*sigma));
        else return 0;
    }

    well_function(double width, double u0, double A, double sigma, double mean) : width(width), u0(u0), A(A), sigma(sigma), mean(mean) {}

    double left_bound() { return -width / 2.0; }
    double right_bound() { return width / 2.0; }
};

struct phase_integrator
{
    double& E;
    well_function& well_f;

    void operator () (double& phase, double& d_phase_dx, double x)
    {
        d_phase_dx = (well_f(x) - E) * pow(cos(phase), 2) - pow(sin(phase), 2);
    }
};

struct wave_function_value : public std::array<double,2>
{
    double& r() { return (*this)[0]; }
    double& phase() { return (*this)[1]; }
};

struct radial_integrator
{
    double& E;
    well_function& well_f;

    void operator () (wave_function_value& phi, wave_function_value& d_phi_dx, double x)
    {
        d_phi_dx.r() = (1 + well_f(x) - E) * phi.r() * cos(phi.phase()) * sin(phi.phase());
        d_phi_dx.phase() = (well_f(x) - E) * pow(cos(phi.phase()), 2) - pow(sin(phi.phase()), 2);
    }
};

template<typename T = double>
std::pair<T, T> divide_by_two_method(std::function<T(T)> f, T left, T right, T epsilon, T compare_to)
{
    T c;
    do
    {
        c = (right + left) / 2.0;
        if ((f(c) - compare_to) * (f(left) - compare_to) < 0)
        {
            right = c;
        }
        else
        {
            left = c;
        }
    } while (abs(right - left) >= epsilon);
    return { f(c),c };
}

class Program
{
public:
    well_function well_f;

    std::vector<double> phase_E;
    std::vector<double> Ek;
    std::vector<double> appropriate_phases;

    size_t chosen_k = 0;
    std::vector<wave_function_value> wave_function;

    double E_loop_size = 2048;
    double x_step = 0.01;

    double R = 3.0;

    void evaluate_ek()
    {
        phase_E.resize(E_loop_size);

        boost::numeric::odeint::runge_kutta4<double> stepper;
        for (int i = 0; i < E_loop_size; i++)
        {
            double E = -well_f.u0 + i / (double)E_loop_size * well_f.u0;
            double phase = PI / 2.0;

            boost::numeric::odeint::integrate_const(
                stepper,
                phase_integrator{ E,well_f },
                phase,
                well_f.left_bound() - R,
                well_f.right_bound() + R,
                x_step
            );
            phase_E[i] = phase;
        }

        double phase_min, phase_max;
        int k_min, k_max;
        phase_min = *std::min_element(phase_E.begin(), phase_E.end());
        phase_max = *std::max_element(phase_E.begin(), phase_E.end());


        k_min = 0;
        k_max = floor((-2 * phase_min / PI - 1) / 2.0);
        k_max = k_max > 0 ? k_max : 0;

        Ek.resize(k_max - k_min);

        appropriate_phases.resize(k_max - k_min);
        for (size_t i = k_min; i < k_max; i++)
        {
            appropriate_phases[i] = (2 * i + 1) / 2.0 * (-PI);
        }

        for (size_t i = 0; i < Ek.size(); i++)
        {
            double appropriate_phase = (2 * i + 1) / 2.0 * (-PI);

            Ek[i] = divide_by_two_method<double>(
                [&](double E)
                {
                    double phase = PI / 2.0;
                    boost::numeric::odeint::integrate_const(
                        stepper,
                        phase_integrator{ E,well_f },
                        phase,
                        well_f.left_bound() - R,
                        well_f.right_bound() + R,
                        x_step
                    );
                    return phase;
                },
                -well_f.u0,
                0,
                1e-5,
                appropriate_phase
            ).second;
        }
    }

    void evaluate_wave_function()
    {
        wave_function.clear();

        wave_function_value start_value;
        start_value.r() = 1.0;
        start_value.phase() = PI / 2.0;

        wave_function.push_back(start_value);

        boost::numeric::odeint::runge_kutta4<wave_function_value> stepper;
        boost::numeric::odeint::integrate_const(
            stepper,
            radial_integrator{ Ek[chosen_k],well_f },
            start_value,
            well_f.left_bound() - R,
            well_f.right_bound() + R,
            x_step,
            [&](const wave_function_value& phi, double x) { wave_function.push_back(phi);}
        );

        auto norm = sqrt(std::accumulate(wave_function.begin(), wave_function.end(),0.0, [&](auto& a, auto& b) {return a + b.r() * b.r() * x_step; }));
        for (auto& i : wave_function) {i.r() *=well_f.u0 / norm;}
    }

    Program() : well_f(1,1,1,1,0){}
};