/**
 * CopyRight F. J. Madrid-Cuevas <fjmadrid@uco.es>
 *
 * Sólo se permite el uso de este código en la docencia de las asignaturas sobre
 * Estructuras de Datos de la Universidad de Córdoba.
 *
 * Está prohibido su uso para cualquier otro objetivo.
 */

#include <cassert>
#include <algorithm>
#include <limits>
#include <cmath>
#include <tuple>

#include "kdtree.hpp"

bool operator==(const Pattern &a, const Pattern &b)
{
    return a.class_label() == b.class_label() &&
           a.dim() == b.dim() &&
           (a.values() == b.values()).min();
}

bool KDTree::is_empty() const
{
    // El KDTree está vacío si su árbol binario interno lo está
    return _tree.is_empty();
}

Pattern KDTree::item() const
{
    assert(!is_empty());
    return _tree.item();
}

KDTree KDTree::left() const
{
    assert(!is_empty());
    return KDTree(_tree.left(), _k);
}

KDTree KDTree::right() const
{
    assert(!is_empty());
    return KDTree(_tree.right(), _k);
}

size_t KDTree::get_k() const
{
    return _k;
}

KDTree::KDTree(BTree<Pattern> tree, size_t k)
{
    _k = k;
    _tree = tree;
}

BTree<Pattern> KDTree::btree() const
{
    return _tree;
}

void KDTree::set_btree(BTree<Pattern> new_tree)
{
    _tree = std::move(new_tree);
}

void KDTree::set_k(size_t k)
{
    _k = k;
    assert(get_k() == k);
}

static BTree<Pattern>
create_kdtree(std::vector<Pattern>::iterator begin,
              std::vector<Pattern>::iterator end,
              size_t level, size_t k)
{
    BTree<Pattern> btree;
    const size_t size = std::distance(begin, end);

    if (size > 0)
    {
        // 1. Calcular eje de corte
        size_t axis = level % k;

        // 2. Localizar mediana
        auto median_it = begin + (size / 2);

        // Lambda de comparación usando el valarray de Pattern
        auto compare_patterns = [axis](const Pattern &a, const Pattern &b) {
            return a.values()[axis] < b.values()[axis];        
        };

        std::nth_element(begin, median_it, end, compare_patterns);

        // 3. Crear raíz con la mediana
        btree = BTree<Pattern>(*median_it);

        // 4. Construcción recursiva de subárboles
        if (begin != median_it) {
            btree.set_left(create_kdtree(begin, median_it, level + 1, k));
        }

        if ((median_it + 1) != end) {
            btree.set_right(create_kdtree(median_it + 1, end, level + 1, k));
        }
    }
    return btree;
}

void KDTree::fit(std::vector<Pattern> &dataset)
{
    assert(dataset.size() > 0);

    _k = dataset[0].dim();
    _tree = create_kdtree(dataset.begin(), dataset.end(), 0, _k);

    assert(get_k() == dataset[0].dim());
}

KDTree::KDTree(std::vector<Pattern> &dataset)
{
    if (dataset.size() > 0) {
        this->fit(dataset);
    }
    assert(dataset.size() == 0 || !is_empty());
}

std::istream &
operator>>(std::istream &in, KDTree &kdtree) noexcept(false)
{
    std::string token;
    in >> token;

    if (token == "[]") {
        kdtree = KDTree();
    } else if (token == "[") {
        size_t k;
        BTree<Pattern> tree;
        
        if (!(in >> k)) throw std::runtime_error("Wrong input format.");
        in >> tree;
        
        in >> token;
        if (token != "]") throw std::runtime_error("Wrong input format.");
        
        kdtree.set_k(k);
        kdtree.set_btree(tree);
    } else {
        throw std::runtime_error("Wrong input format.");
    }
    return in;
}

std::ostream &
operator<<(std::ostream &out, const KDTree &kdtree) noexcept
{
    if (kdtree.is_empty()) {
        out << "[]";
    } else {
        out << "[ " << kdtree.get_k() << " " << kdtree.btree() << " ]";
    }
    return out;
}

std::tuple<float, Pattern>
KDTree::find_nn(Pattern const &p, const KDTree::distance_function_t &dist,
                size_t level) const
{
    assert(!is_empty());

    size_t axis = level % get_k();

    Pattern best = item();
    float best_dist = dist(p, best);

    KDTree next_tree, opposite_tree;

    // STEP 1: Determinar subárboles para continuar la búsqueda
    if (p.values()[axis] < item().values()[axis]) {
        next_tree = left();
        opposite_tree = right();
    } else {
        next_tree = right();
        opposite_tree = left();
    }

    // STEP 2: Buscar en el subárbol más cercano
    if (!next_tree.is_empty()) {
        float tmp_dist;
        Pattern tmp_p;
        std::tie(tmp_dist, tmp_p) = next_tree.find_nn(p, dist, level + 1);
        if (tmp_dist < best_dist) {
            best_dist = tmp_dist;
            best = tmp_p;
        }
    }

    // STEP 3: Evaluar si debemos explorar el subárbol opuesto
    float dist_to_plane = std::abs(p.values()[axis] - item().values()[axis]);

    if (dist_to_plane < best_dist && !opposite_tree.is_empty()) {
        float tmp_dist;
        Pattern tmp_p;
        std::tie(tmp_dist, tmp_p) = opposite_tree.find_nn(p, dist, level + 1);
        if (tmp_dist < best_dist) {
            best_dist = tmp_dist;
            best = tmp_p;
        }
    }

    return std::make_tuple(best_dist, best);
}

std::tuple<float, Pattern>
KDTree::find_nn(Pattern const &p, const KDTree::distance_function_t &dist) const
{
    assert(!is_empty());
    assert(p.dim() == get_k());
    
    return find_nn(p, dist, 0);
}