#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <limits>
#include <new>
#include <exception>
#include <random>
#include <utility>
#include <vector>
#include <omp.h>

namespace {
using Real = long double;
using Index = std::size_t;
using Input = std::uint32_t;
using RNG = std::mt19937_64;

struct Key {
    Real score;
    Index id;
};

bool less(Key a, Key b) {
    return a.score < b.score || (a.score == b.score && a.id < b.id);
}

struct Band {
    Index first = 0, last = 0;
    Real before = 0, mass = 0;

    Index size() const { return last - first; }
    Real gains() const { return mass * (2 * before + mass); }
};

// The two nonempty bands are adjacent in both array position and score order.
Band join(Band a, Band b) {
    if (!a.size()) return b;
    if (!b.size()) return a;
    return {a.first, b.last, a.before, a.mass + b.mass};
}

struct Parts {
    Band low, equal, high;
};

struct RowBand {
    Index channel;
    Band band;
};

struct HeadBand {
    std::vector<RowBand> rows;
    Index count = 0;
    Real before = 0;  // Channel-gain mass below this interval in the full head.
    Real mass = 0;

    void add(Index channel, Band band) {
        if (band.size()) {
            rows.push_back({channel, band});
            count += band.size();
            mass += band.gains();
        }
    }
};

// Merge metadata only. Each channel's partitioned indices stay in place.
HeadBand join(HeadBand a, HeadBand b) {
    if (!a.count) return b;
    if (!b.count) return a;
    HeadBand result;
    result.before = a.before;
    result.rows.reserve(a.rows.size() + b.rows.size());
    Index i = 0, j = 0;
    while (i < a.rows.size() && j < b.rows.size()) {
        const auto& left = a.rows[i];
        const auto& right = b.rows[j];
        if (left.channel < right.channel) {
            result.add(left.channel, left.band);
            ++i;
        } else if (right.channel < left.channel) {
            result.add(right.channel, right.band);
            ++j;
        } else {
            result.add(left.channel, join(left.band, right.band));
            ++i;
            ++j;
        }
    }
    for (; i < a.rows.size(); ++i) result.add(a.rows[i].channel, a.rows[i].band);
    for (; j < b.rows.size(); ++j) result.add(b.rows[j].channel, b.rows[j].band);
    return result;
}

struct HeadParts {
    HeadBand low, equal, high;
};

struct Pivot {
    Index row_slot;
    Input input;
};

Index random_index(RNG& rng, Index size) {
    return std::uniform_int_distribution<Index>(0, size - 1)(rng);
}

struct Projection {
    const void* weights;
    const double* moments;
    std::uint8_t* mask;
    Index heads, channels, inputs;
    int dtype;
};

Real weight_at(const unsigned char* weights, int dtype, Index i) {
    if (dtype == 4) {
        float x;
        std::memcpy(&x, weights + 4*i, sizeof(x));
        return x;
    }
    if (dtype == 8) {
        double x;
        std::memcpy(&x, weights + 8*i, sizeof(x));
        return x;
    }
    std::uint16_t bits;
    std::memcpy(&bits, weights + 2*i, sizeof(bits));
    if (dtype == 3) {
        const std::uint32_t wide = static_cast<std::uint32_t>(bits) << 16;
        float x;
        std::memcpy(&x, &wide, sizeof(x));
        return x;
    }
#ifdef __FLT16_MANT_DIG__
    _Float16 x;
    std::memcpy(&x, &bits, sizeof(x));
    return static_cast<float>(x);
#else
    const unsigned exponent = (bits >> 10) & 31;
    const unsigned fraction = bits & 1023;
    Real x;
    if (exponent == 31) {
        x = fraction ? std::numeric_limits<Real>::quiet_NaN()
                     : std::numeric_limits<Real>::infinity();
    } else if (exponent) {
        x = std::ldexp(static_cast<Real>(1024 + fraction), static_cast<int>(exponent) - 25);
    } else {
        x = std::ldexp(static_cast<Real>(fraction), -24);
    }
    return bits & 32768 ? -x : x;
#endif
}

struct Channel {
    const unsigned char* weights;
    const double* moments;
    Input* indices;
    std::uint8_t* mask;
    Index base;
    int dtype;

    Key key(Input input) const {
        const Real w = weight_at(weights, dtype, input);
        return {w * w * static_cast<Real>(moments[input]), base + input};
    }

    Parts split_raw(Band band, Input pivot) const {
        const Key cut = key(pivot);
        Index lower = band.first, pos = band.first, upper = band.last;
        Real low_mass = 0, high_mass = 0;
        while (pos < upper) {
            const Key value = key(indices[pos]);
            if (less(value, cut)) {
                low_mass += value.score;
                std::swap(indices[lower++], indices[pos++]);
            } else if (less(cut, value)) {
                high_mass += value.score;
                std::swap(indices[pos], indices[--upper]);
            } else {
                ++pos;
            }
        }
        return {
            {band.first, lower, band.before, low_mass},
            {lower, upper, band.before + low_mass, cut.score},
            {upper, band.last, band.before + low_mass + cut.score, high_mass}
        };
    }

    Parts split_gain(Band active, Key cutoff, RNG& rng) const {
        Band low, high;
        while (active.size()) {
            const Input pivot = indices[active.first + random_index(rng, active.size())];
            const Parts parts = split_raw(active, pivot);
            const Key gain{parts.equal.gains(), base + pivot};
            if (less(gain, cutoff)) {
                low = join(join(low, parts.low), parts.equal);
                active = parts.high;
            } else if (less(cutoff, gain)) {
                high = join(join(parts.equal, parts.high), high);
                active = parts.low;
            } else {
                return {join(low, parts.low), parts.equal, join(parts.high, high)};
            }
        }
        return {low, {}, high};
    }
};

struct HeadView {
    const unsigned char* weights;
    const double* moments;
    std::uint8_t* mask;
    Input* indices;
    Index base, channels, inputs, row_base;
    int dtype;
};

struct Tensor {
    std::vector<HeadView> heads;

    Channel row(Index head, Index channel) {
        const auto& view = heads[head];
        const Index offset = channel * view.inputs;
        const Index item_size = view.dtype == 3 ? 2 : view.dtype;
        return {view.weights + offset * item_size, view.moments,
                view.indices ? view.indices + offset : nullptr, view.mask + offset,
                view.base + offset, view.dtype};
    }

    Pivot pick(Index head, const HeadBand& active, Index rank) {
        Index slot = 0;
        for (; rank >= active.rows[slot].band.size(); ++slot) {
            rank -= active.rows[slot].band.size();
        }
        const auto& item = active.rows[slot];
        return {slot, row(head, item.channel).indices[item.band.first + rank]};
    }

    std::pair<HeadParts, Key> split_pivot(
        Index head, const HeadBand& active, Pivot pivot, RNG& rng
    ) {
        const auto& pivot_row = active.rows[pivot.row_slot];
        const auto channel = row(head, pivot_row.channel);
        const Parts pivot_parts = channel.split_raw(pivot_row.band, pivot.input);
        const Key channel_gain{pivot_parts.equal.gains(), channel.base + pivot.input};

        HeadParts result;
        result.low.before = active.before;
        result.low.rows.reserve(active.rows.size());
        result.high.rows.reserve(active.rows.size());
        for (Index j = 0; j < active.rows.size(); ++j) {
            const auto& item = active.rows[j];
            const Parts parts = j == pivot.row_slot ? pivot_parts
                : row(head, item.channel).split_gain(item.band, channel_gain, rng);
            result.low.add(item.channel, parts.low);
            result.high.add(item.channel, parts.high);
        }
        result.equal.before = active.before + result.low.mass;
        result.equal.add(pivot_row.channel, pivot_parts.equal);
        result.high.before = result.equal.before + channel_gain.score;
        const Key final_gain{
            channel_gain.score * (2 * result.equal.before + channel_gain.score),
            channel_gain.id
        };
        return {std::move(result), final_gain};
    }

    HeadParts split_head(Index head, const HeadBand& active, Key cutoff, RNG& rng) {
        if (!active.count) return {};
        const Pivot pivot = pick(head, active, random_index(rng, active.count));
        auto [parts, gain] = split_pivot(head, active, pivot, rng);
        if (less(gain, cutoff)) {
            HeadParts next = split_head(head, parts.high, cutoff, rng);
            next.low = join(join(std::move(parts.low), std::move(parts.equal)),
                            std::move(next.low));
            return next;
        }
        if (less(cutoff, gain)) {
            HeadParts next = split_head(head, parts.low, cutoff, rng);
            next.high = join(join(std::move(next.high), std::move(parts.equal)),
                             std::move(parts.high));
            return next;
        }
        return parts;
    }

    void keep(Index head, const HeadBand& band) {
        for (const auto& item : band.rows) {
            const auto channel = row(head, item.channel);
            for (Index j = item.band.first; j < item.band.last; ++j) {
                channel.mask[channel.indices[j]] = 1;
            }
        }
    }
};

int prune(const Projection* projections, Index count, Input* indices,
          Index budget, int threads, std::uint64_t seed) {
    Tensor tensor;
    Index total = 0, total_rows = 0;
    std::vector<std::pair<Index, Index>> rows;
    for (Index p = 0; p < count; ++p) {
        const auto& view = projections[p];
        const Index head_size = view.channels * view.inputs;
        const Index item_size = view.dtype == 3 ? 2 : view.dtype;
        const auto* weights = static_cast<const unsigned char*>(view.weights);
        for (Index h = 0; h < view.heads; ++h) {
            const Index offset = h * head_size;
            for (Index c = 0; c < view.channels; ++c) rows.emplace_back(tensor.heads.size(), c);
            tensor.heads.push_back({weights + offset * item_size, view.moments,
                view.mask + offset, indices ? indices + total + offset : nullptr,
                total + offset, view.channels, view.inputs, total_rows, view.dtype});
            total_rows += view.channels;
        }
        total += view.heads * head_size;
    }

    std::vector<Real> masses(indices ? rows.size() : 0);
    int invalid = 0;
    const int workers = threads ? threads : omp_get_max_threads();
    const int row_workers = static_cast<int>(std::min<Index>(workers, rows.size()));
    #pragma omp parallel for schedule(static) num_threads(row_workers) if(row_workers > 1 && total >= 8192) reduction(|:invalid)
    for (Index r = 0; r < rows.size(); ++r) {
        const auto [h, c] = rows[r];
        const auto channel = tensor.row(h, c);
        const auto inputs = tensor.heads[h].inputs;
        Real mass = 0;
        for (Index i = 0; i < inputs; ++i) {
            const Real w = weight_at(channel.weights, channel.dtype, i);
            if (!std::isfinite(w)) invalid = 1;
            if (indices) {
                channel.indices[i] = static_cast<Input>(i);
                mass += w * w * static_cast<Real>(channel.moments[i]);
            }
        }
        std::fill_n(channel.mask, inputs, budget == total ? 1 : 0);
        if (indices) {
            masses[r] = mass;
            if (!std::isfinite(mass) || !std::isfinite(mass * mass * mass * mass)) invalid = 1;
        }
    }
    if (invalid) return 1;
    if (!indices) return 0;

    struct ActiveHead {
        Index id;
        HeadBand band;
        RNG rng;
    };
    std::vector<ActiveHead> active;
    active.reserve(tensor.heads.size());
    for (Index h = 0; h < tensor.heads.size(); ++h) {
        const auto& view = tensor.heads[h];
        active.push_back({h, {}, RNG(seed + h + 1)});
        auto& band = active.back().band;
        band.rows.reserve(view.channels);
        for (Index c = 0; c < view.channels; ++c) {
            band.add(c, {0, view.inputs, 0, masses[view.row_base + c]});
        }
        if (!std::isfinite(band.mass * band.mass)) return 1;
    }

    RNG rng(seed);
    Index remaining = total;
    while (budget && remaining) {
        if (budget == remaining) {
            const int team = static_cast<int>(std::min<Index>(workers, active.size()));
            #pragma omp parallel for schedule(static) num_threads(team) if(team > 1 && remaining >= 8192)
            for (Index j = 0; j < active.size(); ++j) tensor.keep(active[j].id, active[j].band);
            break;
        }
        Index rank = random_index(rng, remaining), chosen = 0;
        for (; rank >= active[chosen].band.count; ++chosen) rank -= active[chosen].band.count;
        auto& anchor = active[chosen];
        const Pivot pivot = tensor.pick(anchor.id, anchor.band, rank);
        auto [anchor_parts, cutoff] = tensor.split_pivot(anchor.id, anchor.band, pivot, anchor.rng);

        std::vector<HeadParts> parts(active.size());
        std::vector<std::exception_ptr> errors(active.size());
        parts[chosen] = std::move(anchor_parts);
        const int team = static_cast<int>(std::min<Index>(workers, active.size() - 1));
        const int split_workers = std::max(team, 1);
        #pragma omp parallel for schedule(dynamic, 1) num_threads(split_workers) if(split_workers > 1 && remaining >= 8192)
        for (Index j = 0; j < active.size(); ++j) {
            if (j == chosen) continue;
            try {
                auto& head = active[j];
                parts[j] = tensor.split_head(head.id, head.band, cutoff, head.rng);
            } catch (...) {
                errors[j] = std::current_exception();
            }
        }
        Index above = 0;
        for (Index j = 0; j < parts.size(); ++j) {
            if (errors[j]) std::rethrow_exception(errors[j]);
            above += parts[j].high.count;
        }

        const bool choose_upper = budget <= above;
        if (!choose_upper) {
            const int team = static_cast<int>(std::min<Index>(workers, active.size()));
            #pragma omp parallel for schedule(static) num_threads(team) if(team > 1 && remaining >= 8192)
            for (Index j = 0; j < active.size(); ++j) {
                tensor.keep(active[j].id, parts[j].high);
                tensor.keep(active[j].id, parts[j].equal);
            }
            budget -= above + 1;  // Global index makes the pivot key unique.
        }
        remaining = 0;
        Index next = 0;
        for (Index j = 0; j < active.size(); ++j) {
            HeadBand& band = choose_upper ? parts[j].high : parts[j].low;
            if (band.count) {
                remaining += band.count;
                active[j].band = std::move(band);
                if (next != j) active[next] = std::move(active[j]);
                ++next;
            }
        }
        active.resize(next);
    }
    return 0;
}
}  // namespace

extern "C" int tensor_prune_many(const Projection* projections, Index count, Input* indices,
                                 Index budget, int threads, std::uint64_t seed) {
    try {
        return prune(projections, count, indices, budget, threads, seed);
    } catch (const std::bad_alloc&) {
        return 2;
    } catch (...) {
        return 3;
    }
}