#include "find_subsets.h"
#include <stdexcept>
#include <vector>
#include <unordered_map>
#include <cmath>
#include <thread>
#include <mutex>
#include <atomic>

#include <iostream>

void FindSum(std::unordered_map<int64_t, size_t>& mp, size_t i, size_t j,
             const std::vector<int64_t>& data) {
    if (i == j) {
        return;
    }
    int n = j - i;
    uint64_t bound = pow(3, n);
    for (uint64_t k = 1; k < bound; ++k) {
        int64_t sum = 0;
        uint64_t kk = k;
        int ind = i;
        while (kk != 0) {
            if (kk % 3 == 2) {
                sum += data[ind];
            } else if (kk % 3 == 1) {
                sum -= data[ind];
            }
            ind++;
            kk = kk / 3;
        }
        mp[sum] = k;
    }
}

void FindIndexes(size_t num, std::vector<size_t>& first, std::vector<size_t>& second, int i) {
    int ind = i;
    while (num != 0) {
        if (num % 3 == 2) {
            first.push_back(ind);
        } else if (num % 3 == 1) {
            second.push_back(ind);
        }
        ind++;
        num /= 3;
    }
}

Subsets subset;
std::atomic<bool> fl;
std::mutex mtx;

template <class InputIt>
void my_advance(InputIt& it, InputIt end, int n ) {
    while (n-- > 0 && it != end) {
        std::advance(it, 1);
    }
}

bool Check(size_t first, size_t second) {
    int a = 0, b = 0, c = 0, d = 0;
    while (first != 0) {
        if (first % 3 == 2) {
            a++;
        } else if (first % 3 == 1) {
            b++;
        }
        first /= 3;
    }

    while (second!= 0) {
        if (second % 3 == 2) {
            c++;
        } else if (second % 3 == 1) {
            d++;
        }
        second /= 3;
    }
    return (a + d != 0 && b + c != 0);
}

void Result(const std::unordered_map<int64_t, size_t>& mp1,
            const std::unordered_map<int64_t, size_t>& mp2, size_t sz, int i) {
    auto st = mp1.begin();
    my_advance(st, mp1.end(), i);
    for (auto& it = st; it != mp1.end() && !fl.load(); my_advance(it, mp1.end(), 4)) {
        int64_t key = it->first;
        if (mp2.find(key) != mp2.end()) {
            if (Check(mp1.at(key), mp2.at(key))) {
                std::vector<size_t> first, second;
                FindIndexes(mp1.at(key), first, second, 0);
                FindIndexes(mp2.at(key), second, first, sz);
                fl.store(true);
                mtx.lock();
                subset = Subsets{first, second, true};
                mtx.unlock();
                return;
            }
        }
    }
}

Subsets FindEqual(const std::unordered_map<int64_t, size_t>& mp1,
                  const std::unordered_map<int64_t, size_t>& mp2, size_t sz) {
    fl.store(false);
    std::thread t1([&mp1, &mp2, sz] { Result(mp1, mp2, sz, 0); });
    std::thread t2([&mp1, &mp2, sz] { Result(mp1, mp2, sz, 1); });
    std::thread t3([&mp1, &mp2, sz] { Result(mp1, mp2, sz, 2); });
    std::thread t4([&mp1, &mp2, sz] { Result(mp1, mp2, sz, 3); });
    t1.join();
    t2.join();
    t3.join();
    t4.join();

    if (fl) {
        return subset;
    }
    return Subsets{{}, {}, false};
}

Subsets FindEqualSumSubsets(const std::vector<int64_t>& data) {
    size_t sz = data.size() / 2;
    std::unordered_map<int64_t, size_t> mp1;
    std::unordered_map<int64_t, size_t> mp2;
    std::thread th1([&mp1, sz, &data] { FindSum(mp1, 0, sz, data); });
    std::thread th2([&mp2, sz, &data] { FindSum(mp2, sz, data.size(), data); });
    th1.join();
    th2.join();
    std::vector<size_t> first, second;
    if (mp1.find(0) != mp1.end()) {
        FindIndexes(mp1[0], first, second, 0);
        return Subsets{first, second, true};
    }
    if (mp2.find(0) != mp2.end()) {
        FindIndexes(mp2[0], first, second, sz);
        return Subsets{first, second, true};
    }
    return FindEqual(mp1, mp2, sz);
}