// Copyright 2022 James Pace
// All Rights Reserved.
//
// For a license to this software contact
// James Pace at jpace121@gmail.com.
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#include "j7s-optimization/NelderMead.hpp"

#include <algorithm>
#include <cmath>
#include <stdexcept>

namespace j7s
{
NelderMead::NelderMead(
    const CostFunction & costFunction, const std::vector<double> initSimplex) :
    m_costFunction{costFunction}
{
    m_currentSimplex.reserve(initSimplex.size());

    for (const auto val : initSimplex)
    {
        m_currentSimplex.emplace_back(val, m_costFunction.eval(val));
    }
    std::sort(m_currentSimplex.begin(), m_currentSimplex.end());
}

Coordinate NelderMead::bestCoord() const
{
    return m_currentSimplex.front();
}

double NelderMead::newPoint() const
{
    if (m_currentSimplex.size() == 0)
    {
        throw std::runtime_error("Simplex can't be missing.");
    }

    // Calculate sum.
    double biggest = m_currentSimplex.back().input;
    // Calculate volume.
    double sum = 0.0;
    // All but the biggest, (would be adding 0...)
    for (unsigned int index = 0; index < m_currentSimplex.size() - 1; index++)
    {
        const double diff = m_currentSimplex[index].input - biggest;
        sum += diff;
    }

    const double newPoint = sum * (2.0 / m_currentSimplex.size());

    return newPoint;
}

double NelderMead::findCentroid() const
{
    double sum = 0.0;
    for (unsigned int index = 0; index < m_currentSimplex.size() - 1; index++)
    {
        sum += m_currentSimplex[index].input;
    }

    const double newPoint = sum * (1.0 / m_currentSimplex.size());

    return newPoint;
}

double NelderMead::secondNewPoint(const Coordinate& newPoint) const
{
    // FxN in paper.
    const auto biggest = m_currentSimplex.back();
    // Fx0 in paper.
    const auto smallest = m_currentSimplex.front();
    // FxN-1 in paper.
    // TODO: Assuming simplex size here.
    const auto secondBiggest = *(m_currentSimplex.end() - 2);

    const double centroid = findCentroid();

    if(newPoint < smallest)
    {
        return 2.0*newPoint.input - centroid;
    }
    else if(smallest <= newPoint and newPoint < secondBiggest)
    {
        return 0.5*(centroid + newPoint.input);
    }
    else if(secondBiggest <= newPoint)
    {
        return 0.5*(biggest.input + centroid);
    }
    else
    {
        // Shouldn't be able to get here?
        throw std::logic_error("Shouldn't be able to get here.");
    }
}

std::vector<Coordinate> NelderMead::contract() const
{
    const auto smallest = m_currentSimplex.front();

    std::vector<Coordinate> newVector;
    newVector.reserve(m_currentSimplex.size());
    newVector.emplace_back(smallest);

    // TODO: Really check size before I get here...
    for (auto it = m_currentSimplex.begin() + 1; it != m_currentSimplex.end(); it++)
    {
        const auto oldInput = it->input;
        const auto newInput = 0.5 * (oldInput + smallest.input);
        const auto newCost = m_costFunction.eval(newInput);
        newVector.emplace_back(newInput, newCost);
    }
    std::sort(newVector.begin(), newVector.end());

    return newVector;
}

double NelderMead::calcVolume() const
{
    // TODO: For reals do something like:
    // https://math.stackexchange.com/questions/337197/finding-the-volume-of-a-tetrahedron-by-given-vertices
    // For now:
    // Sort by input and find the difference squared between the first and last.

    const auto inputLess = [](const Coordinate & first, const Coordinate & second)
    { return first.input < second.input; };

    // Copy the vector so we don't sort the original.
    std::vector<Coordinate> simplexCopy = m_currentSimplex;
    std::sort(simplexCopy.begin(), simplexCopy.end(), inputLess);
    const auto smallest = simplexCopy.front();
    const auto biggest = simplexCopy.back();
    return std::pow(biggest.input - smallest.input, 2.0);
}

IterationState NelderMead::update()
{
    if (m_currentSimplex.size() < 3)
    {
        throw std::runtime_error("Simplex can't be a line.");
    }
    // Check for convergence and potentially early return.
    // TODO: Make configurable.
    const auto volume = calcVolume();
    if (volume < 1e-4)
    {
        return IterationState::CONVERGED;
    }

    // Do update.
    Coordinate potential;
    potential.input = newPoint();
    potential.cost = m_costFunction.eval(potential.input);
    Coordinate secondPotential;
    secondPotential.input = secondNewPoint(potential);
    secondPotential.cost = m_costFunction.eval(secondPotential.input);

    const auto minPotential = std::min(potential, secondPotential);

    const auto secondBiggest = *(m_currentSimplex.end() - 2);
    if (minPotential.cost < secondBiggest.cost)
    {
        // Replace the last simplex value with the better one.
        *(m_currentSimplex.end() - 1) = minPotential;
        std::sort(m_currentSimplex.begin(), m_currentSimplex.end());
    }
    else
    {
        // Do a contraction.
        m_currentSimplex = contract();
    }

    return IterationState::OK;
}

}  // namespace j7s
