#pragma once
#include <numeric>
#include <vector>
#include <iostream>
#include <iterator>
#include <limits>
#include <algorithm>
#include <set>
#include <tuple>
#include <queue>
using namespace std;

template <typename NodeType>
struct NodeAPI
{

    NodeAPI(){};

    float get_leaf_purity(NodeType *node, int tot_cnt)
    {
        float leaf_purity = 0;
        if (node->is_leaf())
        {
            int cnt_idxs = node->idxs.size();
            if (cnt_idxs == 0)
            {
                leaf_purity = 0.0;
            }
            else
            {
                int cnt_zero = 0;
                for (int i = 0; i < node->idxs.size(); i++)
                {
                    if (node->y[node->idxs[i]] == 0)
                    {
                        cnt_zero += 1;
                    }
                }
                leaf_purity = max(float(cnt_zero) / float(cnt_idxs),
                                  1 - float(cnt_zero) / float(cnt_idxs));
                leaf_purity = leaf_purity * (float(cnt_idxs) / float(tot_cnt));
            }
        }
        else
        {
            leaf_purity = get_leaf_purity(node->left, tot_cnt) + get_leaf_purity(node->right, tot_cnt);
        }
        return leaf_purity;
    }

    string print(NodeType *node, bool show_purity = false, bool binary_color = true, int target_party_id = -1)
    {
        pair<string, bool> result = recursive_print(node, "", false, show_purity, binary_color, target_party_id);
        if (result.second)
        {
            return "";
        }
        else
        {
            return result.first;
        }
    }

    string print_leaf(NodeType *node, bool show_purity, bool binary_color)
    {
        string node_info = to_string(node->get_val()[0]);
        if (show_purity)
        {
            int cnt_idxs = node->idxs.size();
            if (cnt_idxs == 0)
            {
                node_info += ", null";
            }
            else
            {
                int cnt_zero = 0;
                for (int i = 0; i < node->idxs.size(); i++)
                {
                    if (node->y[node->idxs[i]] == 0)
                    {
                        cnt_zero += 1;
                    }
                }
                float purity = max(float(cnt_zero) / float(cnt_idxs),
                                   1 - float(cnt_zero) / float(cnt_idxs));
                node_info += ", ";

                if (binary_color)
                {
                    if (purity < 0.7)
                    {
                        node_info += "\033[32m";
                    }
                    else if (purity < 0.9)
                    {
                        node_info += "\033[33m";
                    }
                    else
                    {
                        node_info += "\033[31m";
                    }
                    node_info += to_string(purity);
                    node_info += " (";
                    node_info += to_string(cnt_zero);
                    node_info += ", ";
                    node_info += to_string(cnt_idxs - cnt_zero);
                    node_info += ")";
                    node_info += "\033[0m";
                }
                else
                {
                    node_info += to_string(purity);
                }
            }
        }
        else
        {
            node_info += ", [";
            int temp_id;
            for (int i = 0; i < node->idxs.size(); i++)
            {
                temp_id = node->idxs[i];
                if (binary_color)
                {
                    if (node->y[temp_id] == 0)
                    {
                        node_info += "\033[32m";
                        node_info += to_string(temp_id);
                        node_info += "\033[0m";
                    }
                    else
                    {
                        node_info += to_string(temp_id);
                    }
                }
                else
                {
                    node_info += to_string(temp_id);
                }
                node_info += ", ";
            }
            node_info += "]";
        }

        return node_info;
    }

    pair<string, bool> recursive_print(NodeType *node, string prefix, bool isleft, bool show_purity,
                                       bool binary_color, int target_party_id = -1)
    {
        string node_info;
        bool skip_flag;
        if (node->is_leaf())
        {
            skip_flag = node->depth <= 0 && target_party_id != -1 && node->party_id != target_party_id;
            if (skip_flag)
            {
                node_info = "";
            }
            else
            {
                node_info = print_leaf(node, show_purity, binary_color);
            }
            node_info = prefix + "|-- " + node_info;
            node_info += "\n";
        }
        else
        {
            node_info += to_string(node->get_party_id());
            node_info += ", ";
            node_info += to_string(node->get_record_id());
            node_info = prefix + "|-- " + node_info;

            string next_prefix = "";
            if (isleft)
            {
                next_prefix += "|    ";
            }
            else
            {
                next_prefix += "     ";
            }

            pair<string, bool> left_node_info_and_skip_flag = recursive_print(node->left, prefix + next_prefix, true,
                                                                              show_purity, binary_color, target_party_id);
            pair<string, bool> right_node_info_and_skip_flag = recursive_print(node->right, prefix + next_prefix, false,
                                                                               show_purity, binary_color, target_party_id);
            if (left_node_info_and_skip_flag.second && right_node_info_and_skip_flag.second)
            {
                node_info += " -> " + print_leaf(node, show_purity, binary_color);
                node_info += "\n";
            }
            else
            {
                node_info += "\n";
                node_info += left_node_info_and_skip_flag.first;
                node_info += right_node_info_and_skip_flag.first;
            }

            skip_flag = false;
        }

        return make_pair(node_info, skip_flag);
    }

    vector<float> predict_row(NodeType *node, vector<float> &xi)
    {
        queue<NodeType *> que;
        que.push(node);

        NodeType *temp_node;
        while (!que.empty())
        {
            temp_node = que.front();
            que.pop();

            if (temp_node->is_leaf())
            {
                return temp_node->val;
            }
            else
            {
                if (node->parties[temp_node->party_id].is_left(temp_node->record_id, xi))
                {
                    que.push(temp_node->left);
                }
                else
                {
                    que.push(temp_node->right);
                }
            }
        }

        vector<float> nan_vec(node->num_classes, nan(""));
        return nan_vec;
    }

    vector<vector<float>> predict(NodeType *node, vector<vector<float>> &x_new)
    {
        int x_new_size = x_new.size();
        vector<vector<float>> y_pred(x_new_size);
        for (int i = 0; i < x_new_size; i++)
        {
            y_pred[i] = predict_row(node, x_new[i]);
        }
        return y_pred;
    }
};
