#include <rofi/voxel_graph.hpp>

#include <algo/misc.hpp>

#include <algorithm>
#include <iostream>

namespace rofi
{

void VoxelGraph::joinSingular(side s, connector_ptr connector, voxel_ptr voxel)
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());

    // NOTE: Careful with ordering these two functions
    voxel->connectConnection(s, connector);
    connector->connectVoxel(s, voxel);
}

void VoxelGraph::join(glm::vec3 local_direction, connector_ptr connector, voxel_ptr voxel)
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());

    auto s = vec_to_enum(local_direction);

    if (voxel->getConnections().contains(s)
        && voxel->getConnectorAt(local_direction).lock()->getType(false) == connector_type::empty)
    {
        auto rot_con = std::dynamic_pointer_cast<RotationConnector>(voxel->getConnectorAt(local_direction).lock());
        ASSUMPTION(rot_con);
        ASSUMPTION(connector->getType(false) == connector_type::fixed || connector->getType(false) == connector_type::open);
        rot_con->setSemanticType(connector->getType(false));
        connector->connectVoxel(s, voxel);
        return;
    }

    joinSingular(s, connector, voxel);

    if (auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector)
        ; rot_con && rot_con->isHead() && rot_con->getType(voxel) == connector_type::rotation)
    {
        joinSingular(vec_to_enum(rot_con->getRotationAxis()), rot_con->getBoundConnectors().first.lock(), voxel);
        joinSingular(vec_to_enum(-rot_con->getRotationAxis()), rot_con->getBoundConnectors().second.lock(), voxel);
    }
}

void VoxelGraph::disjoinSingular(connector_ptr connector, voxel_ptr voxel)
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());

    connector->disconnectVoxel(voxel);
    voxel->disconnectConnection(connector);
}

void VoxelGraph::disjoinSingular(glm::vec3 direction, voxel_ptr voxel)
{
    auto s = vec_to_enum(direction);

    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());
    ASSUMPTION(voxel->getConnections().contains(s));

    auto connector = voxel->getConnections().find(s)->second.lock();

    disjoinSingular(connector, voxel);
}

void VoxelGraph::removeConnectorSingular(connector_ptr connector)
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());
    connections.erase(std::find(connections.begin(), connections.end(), connector));
}

bool VoxelGraph::contains(voxel_ptr voxel)
{
    return std::find(voxels.begin(), voxels.end(), voxel) != voxels.end();
}

bool VoxelGraph::contains(connector_ptr connector)
{
    return std::find(connections.begin(), connections.end(), connector) != connections.end();
}

auto find_new_from_orig(auto orig_new_pairs, auto orig_component)
{
    return std::find_if(orig_new_pairs.begin(), orig_new_pairs.end(), [&orig_component](const auto & pair) { return pair.first == orig_component; } )->second;
}

voxel_graph_ptr VoxelGraph::copy(const voxel_graph_ptr other_voxel_graph, const std::vector<node_ptr> &new_nodes)
{
    std::vector<std::pair<voxel_ptr, voxel_ptr>> voxel_pairs;
    std::vector<std::pair<connector_ptr, connector_ptr>> connection_pairs;

    for (const auto &orig_voxel : other_voxel_graph->voxels)
        voxel_pairs.emplace_back(orig_voxel, Voxel::copy(orig_voxel));
        // voxel_pairs[orig_voxel] = Voxel::copy(orig_voxel);
    for (const auto &orig_connector : other_voxel_graph->connections)
        connection_pairs.emplace_back(orig_connector, Connector::copy(orig_connector));
        // connection_pairs[orig_connector] = Connector::copy(orig_connector);

    for (const auto &orig_voxel : other_voxel_graph->voxels)
    {
        for (const auto &[side, weak_con] : orig_voxel->connections)
        {
            auto orig_connector = weak_con.lock();
            auto new_voxel = find_new_from_orig(voxel_pairs, orig_voxel);
            new_voxel->connections[side] = find_new_from_orig(connection_pairs, orig_connector);
            // voxel_pairs[orig_voxel]->connections[side] = connection_pairs[orig_connector];
        }
    }

    for (const auto &orig_connector : other_voxel_graph->connections)
    {
        if (!orig_connector->connection.first.expired())
        {
            const auto orig_voxel = orig_connector->connection.first.lock();
            find_new_from_orig(connection_pairs, orig_connector)->connection.first = find_new_from_orig(voxel_pairs, orig_voxel);
            // connection_pairs[orig_connector]->connection.first = find_new_from_orig(voxel_pairs, orig_voxel);
        }

        if (!orig_connector->connection.second.expired())
        {
            const auto orig_voxel = orig_connector->connection.second.lock();
            find_new_from_orig(connection_pairs, orig_connector)->connection.second = find_new_from_orig(voxel_pairs, orig_voxel);
            // connection_pairs[orig_connector]->connection.second = find_new_from_orig(voxel_pairs, orig_voxel);
        }

        if (const auto orig_rot_con = std::dynamic_pointer_cast<RotationConnector>(orig_connector))
        {
            auto rot_con = std::dynamic_pointer_cast<RotationConnector>(find_new_from_orig(connection_pairs, orig_connector));
            auto fst = find_new_from_orig(connection_pairs, orig_rot_con->getBoundConnectors().first.lock());
            auto snd = find_new_from_orig(connection_pairs, orig_rot_con->getBoundConnectors().second.lock());
            rot_con->setBoundConnectors(fst, snd);
        }
    }

    copy_nodes(voxel_pairs, connection_pairs, new_nodes);

    std::vector<voxel_ptr> copy_voxels;
    std::vector<connector_ptr> copy_connections;
    std::transform(voxel_pairs.begin(), voxel_pairs.end(), std::back_inserter(copy_voxels), 
                   [](std::pair<voxel_ptr, voxel_ptr> vp) { return vp.second; });
    std::transform(connection_pairs.begin(), connection_pairs.end(), std::back_inserter(copy_connections), 
                   [](std::pair<connector_ptr, connector_ptr> cp) { return cp.second; });

    return std::make_shared<VoxelGraph>(copy_voxels, copy_connections);
}

void VoxelGraph::copy_nodes(const std::vector<std::pair<voxel_ptr, voxel_ptr>> &voxel_pairs, 
                            const std::vector<std::pair<connector_ptr, connector_ptr>> &connection_pairs, 
                            const std::vector<node_ptr> &new_nodes)
{
    for (auto &[orig_voxel, new_voxel] : voxel_pairs)
    {
        auto new_node = find_node_with_rofi_object(new_nodes, orig_voxel);
        ASSUMPTION(new_node);
        new_voxel->node = new_node;

        std::erase_if(new_node->getObjects(),
                  [&orig_voxel](const objectbase_ptr &object) { return std::dynamic_pointer_cast<Voxel>(object) == orig_voxel; });
        new_node->addObject(new_voxel);
    }

    for (auto &[orig_connector, new_connector] : connection_pairs)
    {
        auto new_node = find_node_with_rofi_object(new_nodes, orig_connector);
        ASSUMPTION(new_node);
        new_connector->node = new_node;

        std::erase_if(new_node->getObjects(),
                  [&orig_connector](const objectbase_ptr &object) { return std::dynamic_pointer_cast<Connector>(object) == orig_connector; });
        new_node->addObject(new_connector);
    }
}

void VoxelGraph::linkImportVoxelGraphToNodes(const voxel_graph_ptr vg)
{
    for (const auto &voxel : vg->getVoxels())
        voxel->getNode()->addObject(voxel);

    for (const auto &connector : vg->getConnections())
        connector->getNode()->addObject(connector);
}

void VoxelGraph::extend(const voxel_graph_ptr other_voxel_graph)
{
    voxels.reserve(voxels.size() + other_voxel_graph->getVoxels().size());
    voxels.insert(voxels.end(), other_voxel_graph->getVoxels().begin(), other_voxel_graph->getVoxels().end());

    connections.reserve(connections.size() + other_voxel_graph->getConnections().size());
    connections.insert(connections.end(), other_voxel_graph->getConnections().begin(), other_voxel_graph->getConnections().end());
}

void VoxelGraph::erase(const voxel_graph_ptr other_voxel_graph)
{
    auto prev_voxels_size = voxels.size();
    auto prev_connections_size = connections.size();

    for (const auto &voxel : other_voxel_graph->voxels)
        std::erase(voxels, voxel);
    for (const auto &connector : other_voxel_graph->connections)
        std::erase(connections, connector);

    ASSUMPTION(prev_voxels_size - other_voxel_graph->voxels.size() == voxels.size());
    ASSUMPTION(prev_connections_size - other_voxel_graph->connections.size() == connections.size());
}

bool VoxelGraph::cellAvailable(glm::vec3 position) 
{
    auto epsilon = 0.0001f;
    for (auto &voxel : voxels)
    {
        if (glm::all(glm::epsilonEqual(position, voxel->getNode()->getPositionWorld(), epsilon)))
            return false;
    }
    return true;
}

voxel_ptr VoxelGraph::getVoxelAtPosition(glm::vec3 position)
{
    auto epsilon = 0.0001f;
    for (auto &voxel : voxels)
    {
        if (glm::all(glm::epsilonEqual(position, voxel->getNode()->getPositionWorld(), epsilon)))
            return voxel;
    }
    return nullptr;
}

voxel_ptr VoxelGraph::addVoxel(node_ptr node)
{
    voxels.emplace_back(std::make_shared<Voxel>(node));
    return voxels.back();
}

voxel_ptr VoxelGraph::addModuleLink(node_ptr node)
{
    voxels.emplace_back(std::make_shared<ModuleLink>(node));
    return voxels.back();
}

voxel_ptr VoxelGraph::addPadBoard(node_ptr node)
{
    voxels.emplace_back(std::make_shared<PadBoard>(node));
    return voxels.back();
}

connector_ptr VoxelGraph::addConnector(node_ptr node, connector_type t)
{
    using enum connector_type;
    ASSUMPTION(t == fixed || t == open);
    switch(t)
    {
        case fixed:
            connections.emplace_back(std::make_shared<FixedConnector>(node));
            break;
        case open:
            connections.emplace_back(std::make_shared<OpenConnector>(node));
            break;   
    }
    return connections.back();
}

connector_ptr VoxelGraph::addConnector(node_ptr node, connector_type t, glm::vec3 rotation_axis)
{
    using enum connector_type;
    ASSUMPTION(t == shared_rotation);

    connections.emplace_back(std::make_shared<SharedRotationConnector>(node, rotation_axis));

    return connections.back();
}

connector_ptr VoxelGraph::addConnector(node_ptr node, connector_type t, glm::vec3 rotation_axis, glm::vec3 highlight_direction)
{
    using enum connector_type;
    ASSUMPTION(t == rotation);

    glm::vec3 highlight_direction_world = node->getRotationMatWorld() * glm::vec4(highlight_direction, 0.0f);
    glm::vec3 rotation_axis_world = node->getRotationMatWorld() * glm::vec4(rotation_axis, 0.0f);

    ASSUMPTION(glm::length(rotation_axis) == 1.0f);

    auto fst_node = Node::create(node->getPositionWorld() - (highlight_direction_world / 2.0f) + (rotation_axis_world / 2.0f), glm::rotation(enum_to_vec(side::z_neg), rotation_axis_world));
    auto snd_node = Node::create(node->getPositionWorld() - (highlight_direction_world / 2.0f) - (rotation_axis_world / 2.0f), glm::rotation(enum_to_vec(side::z_neg), -rotation_axis_world));

    connections.emplace_back(std::make_shared<RotationConnector>(node, rotation_axis, true));
    auto head = connections.back();
    connections.emplace_back(std::make_shared<RotationConnector>(fst_node, rotation_axis, false));
    auto fst = connections.back();
    connections.emplace_back(std::make_shared<RotationConnector>(snd_node, rotation_axis, false));
    auto snd = connections.back();

    std::dynamic_pointer_cast<rofi::RotationConnector>(head)->setBoundConnectors(fst, snd);
    std::dynamic_pointer_cast<rofi::RotationConnector>(fst)->setBoundConnectors(head, snd);
    std::dynamic_pointer_cast<rofi::RotationConnector>(snd)->setBoundConnectors(head, fst);

    return head;
}

void VoxelGraph::removeVoxel(voxel_ptr voxel)
{
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());
    voxels.erase(std::find(voxels.begin(), voxels.end(), voxel));
}

void VoxelGraph::removeConnector(connector_ptr connector)
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());

    if (auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector)
        ; rot_con)
    {
        if (rot_con->isHead())
        {
            ASSUMPTION(!rot_con->getBoundConnectors().first.expired() && !rot_con->getBoundConnectors().second.expired());
            removeConnectorSingular(rot_con->getBoundConnectors().first.lock());
            removeConnectorSingular(rot_con->getBoundConnectors().second.lock());
            // NOTE: Cannot ask for .expired() without this due to ownership by 'to_remove_connectors' vector in editor
        }
        else 
        {
            removeConnector(rot_con->getHead().lock());
            return;
        }    
    }

    connections.erase(std::find(connections.begin(), connections.end(), connector));
}

void VoxelGraph::removeVoxels(const std::vector<voxel_ptr> &to_remove_voxels)
{
    for (auto &voxel : to_remove_voxels)
        removeVoxel(voxel);
}

void VoxelGraph::removeConnectors(const std::vector<connector_ptr> &to_remove_connectors)
{
    for (auto &connector : to_remove_connectors)
    {
        if (auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector)
            ; !rot_con || rot_con->isHead())
            removeConnector(connector);
    }
}

void VoxelGraph::join(connector_ptr connector, glm::vec3 connector_local_direction, voxel_ptr voxel)
{
    auto global_direction = connector->getNode()->getRotationMatWorld() * glm::vec4(connector_local_direction, 0.0f);
    auto voxel_local_direction = voxel->getNode()->getInverseRotationMatWorld() * global_direction;

    join(glm::vec3(voxel_local_direction), connector, voxel);
}

void VoxelGraph::join(connector_ptr connector, voxel_ptr voxel, glm::vec3 voxel_local_direction)
{
    join(voxel_local_direction, connector, voxel);
}

void VoxelGraph::joinConnectors(connector_ptr fst_connector, connector_ptr snd_connector)
{
    auto module_link = addModuleLink(Node::create(fst_connector->getNode()->getPositionWorld(), fst_connector->getNode()->getRotationWorld()));
    join(enum_to_vec(side::z_pos), fst_connector, module_link);
    join(enum_to_vec(side::z_neg), snd_connector, module_link);
}

void VoxelGraph::disjoin(connector_ptr connector, voxel_ptr voxel) 
{
    ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());

    if (auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector)
        ; rot_con && rot_con->isParentVoxel(voxel))
    {
        if (rot_con->isHead())
        {
            if (voxel->isConnectedAt(rot_con->getRotationAxis()))
                disjoinSingular(rot_con->getRotationAxis(), voxel);
            if (voxel->isConnectedAt(-rot_con->getRotationAxis()))
                disjoinSingular(-rot_con->getRotationAxis(), voxel);
        }
        else
        {
            ASSUMPTION(!rot_con->getHead().expired());
            disjoin(rot_con->getHead().lock(), voxel);
            return;
        }
    }

    disjoinSingular(connector, voxel);
}

void VoxelGraph::disjoinConnectors(connector_ptr fst_connector, connector_ptr snd_connector)
{
    auto fst_connection = fst_connector->getConnection();

    auto module_link = std::dynamic_pointer_cast<ModuleLink>(fst_connection.first.lock()) 
                       ? std::dynamic_pointer_cast<ModuleLink>(fst_connection.first.lock()) 
                       : std::dynamic_pointer_cast<ModuleLink>(fst_connection.second.lock());

    ASSUMPTION(module_link);

    disjoinAll(module_link);
    removeVoxel(module_link);
}

// void VoxelGraph::disjoin(glm::vec3 direction, voxel_ptr voxel)
// {
//     auto s = vec_to_enum(direction);

//     ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());
//     ASSUMPTION(voxel->getConnections().contains(s));

//     auto connector = voxel->getConnections().find(s)->second.lock();

//     disjoin(connector, voxel);
// }

void VoxelGraph::disjoinAll(voxel_ptr voxel)
{
    ASSUMPTION(std::find(voxels.begin(), voxels.end(), voxel) != voxels.end());

    for (auto it = voxel->getConnections().begin()
         ; it != voxel->getConnections().end()
         ; it = voxel->getConnections().begin())
    {
        disjoin(it->second.lock(), voxel);
    }     
}

void VoxelGraph::disjoinAll(connector_ptr connector)
{
    // TO DO: disjoinAll for a vector of connectors with assumption
    // ASSUMPTION(std::find(connections.begin(), connections.end(), connector) != connections.end());

    /* REMOVES THE CONNECTOR ENTRIES IN ASSOCIATED VOXELS*/
    if (!connector->getConnection().first.expired())
        disjoin(connector, connector->getConnection().first.lock());
    if (!connector->getConnection().second.expired())
        disjoin(connector, connector->getConnection().second.lock());
    
    // DESC: disjoins other two bound connectors from their secondary voxels (non-parent)
    if (auto rot_con = std::dynamic_pointer_cast<RotationConnector>(connector)
             ; rot_con)
    {
        if (!rot_con->getBoundConnectors().first.lock()->getConnection().second.expired())
            disjoinSingular(rot_con->getBoundConnectors().first.lock(), rot_con->getBoundConnectors().first.lock()->getConnection().second.lock());
        if (!rot_con->getBoundConnectors().second.lock()->getConnection().second.expired())
            disjoinSingular(rot_con->getBoundConnectors().second.lock(), rot_con->getBoundConnectors().second.lock()->getConnection().second.lock());   
    }
}

// void VoxelGraph::disjoin(connector_ptr connector, voxel_ptr fst_voxel, voxel_ptr snd_voxel)
// {
// }

void VoxelGraph::disjoinAll()
{
    for (auto &connector : connections)
    {
        if (connector->hasModuleLink())
        {
            auto opposite_connector = connector->getModuleLink().lock()->getOppositeConnector(connector).lock();
            disjoinConnectors(connector, opposite_connector);
        }
        
        disjoinAll(connector);
    }
}

void VoxelGraph::joinImplicit()
{
    for (auto con : connections)
    {
        if (con->isFull())
            continue;

        for (auto voxel : voxels)
        {
            if (voxel->connections.size() == 6)
                continue;

            side s = in_proximity(con, voxel);
            if (s == side::undefined)
                continue;

            auto s_voxel = vec_to_enum(voxel->getNode()->getInverseRotationMatWorld() * glm::vec4(enum_to_vec(s), 0.0f));

            if (!voxel->isConnectedAt(s_voxel)
                && voxel->slotCompatible(s_voxel, con->getType(voxel)))
                joinSingular(s_voxel, con, voxel);
            else if (voxel->isConnectedAt(s_voxel)
                     && con->getType(true) == connector_type::rotation
                     && con->getType(voxel) != connector_type::rotation
                     && con->getType(voxel) != connector_type::empty)
            {
                auto connector_at_voxel = voxel->getConnectorAt(s_voxel).lock();
                if (connector_at_voxel->getType(voxel) == connector_type::rotation
                    && connector_at_voxel->getType(false) != connector_type::empty)
                {
                    ASSUMPTION(!connector_at_voxel->isFull());
                    // TO DO: !!! FIX
                    // connector_at_voxel->connectVoxel(s_voxel, con->getOppositeConnection().lock());
                    // con->connectVoxel(s_voxel, voxel);
                }
            }
        }
    }
}

side in_proximity(connector_ptr connector, voxel_ptr voxel)
{   
    // TO DO: adjust for rotations
    auto con_position = connector->getNode()->getPositionWorld();
    auto voxel_position = voxel->getNode()->getPositionWorld();

    using enum side;

    if (glm::distance(con_position, voxel_position) > 0.6f)
        return undefined;

    auto diff = con_position - voxel_position;

    float epsilon = 0.0001f;

    if (glm::all(glm::epsilonEqual(diff, glm::vec3(0.5f, 0.0f, 0.0f), epsilon)))
        return x_pos;
    if (glm::all(glm::epsilonEqual(diff, glm::vec3(-0.5f, 0.0f, 0.0f), epsilon)))
        return x_neg;
    if (glm::all(glm::epsilonEqual(diff, glm::vec3(0.0f, 0.5f, 0.0f), epsilon)))
        return y_pos;
    if (glm::all(glm::epsilonEqual(diff, glm::vec3(0.0f, -0.5f, 0.0f), epsilon)))
        return y_neg;
    if (glm::all(glm::epsilonEqual(diff, glm::vec3(0.0f, 0.0f, 0.5f), epsilon)))
        return z_pos;
    if (glm::all(glm::epsilonEqual(diff, glm::vec3(0.0f, 0.0f, -0.5f), epsilon)))
        return z_neg;

    // ASSUMPTION(false);
    return undefined;
}

}
