#include "interp_lb_profile.h"
#include <cstdio>
#include <cstring>
#include <algorithm>

namespace InterpLBProfile {

bool write_profile(const char *filepath, int nprocs,
                   const double *rank_times,
                   const int *heavy_ranks, int num_heavy,
                   double threshold_ratio)
{
    FILE *fp = fopen(filepath, "wb");
    if (!fp) return false;

    ProfileHeader hdr;
    hdr.magic = MAGIC;
    hdr.version = VERSION;
    hdr.nprocs = nprocs;
    hdr.num_heavy = num_heavy;
    hdr.threshold_ratio = threshold_ratio;

    fwrite(&hdr, sizeof(hdr), 1, fp);
    fwrite(rank_times, sizeof(double), nprocs, fp);
    fwrite(heavy_ranks, sizeof(int), num_heavy, fp);
    fclose(fp);
    return true;
}

bool read_profile(const char *filepath, int current_nprocs,
                  int *heavy_ranks, int &num_heavy,
                  double *rank_times, MPI_Comm comm)
{
    int myrank;
    MPI_Comm_rank(comm, &myrank);

    int valid = 0;
    ProfileHeader hdr;
    memset(&hdr, 0, sizeof(hdr));

    if (myrank == 0) {
        FILE *fp = fopen(filepath, "rb");
        if (fp) {
            if (fread(&hdr, sizeof(hdr), 1, fp) == 1 &&
                hdr.magic == MAGIC && hdr.version == VERSION &&
                hdr.nprocs == current_nprocs)
            {
                if (fread(rank_times, sizeof(double), current_nprocs, fp)
                    == (size_t)current_nprocs &&
                    fread(heavy_ranks, sizeof(int), hdr.num_heavy, fp)
                    == (size_t)hdr.num_heavy)
                {
                    num_heavy = hdr.num_heavy;
                    valid = 1;
                }
            } else if (fp) {
                printf("[InterpLB] Profile rejected: magic=0x%X version=%u "
                       "nprocs=%d (current=%d)\n",
                       hdr.magic, hdr.version, hdr.nprocs, current_nprocs);
            }
            fclose(fp);
        }
    }

    MPI_Bcast(&valid, 1, MPI_INT, 0, comm);
    if (!valid) return false;

    MPI_Bcast(&num_heavy, 1, MPI_INT, 0, comm);
    MPI_Bcast(heavy_ranks, num_heavy, MPI_INT, 0, comm);
    MPI_Bcast(rank_times, current_nprocs, MPI_DOUBLE, 0, comm);
    return true;
}

int identify_heavy_ranks(const double *rank_times, int nprocs,
                         double threshold_ratio,
                         int *heavy_ranks, int max_heavy)
{
    double sum = 0;
    for (int i = 0; i < nprocs; i++) sum += rank_times[i];
    double mean = sum / nprocs;
    double threshold = threshold_ratio * mean;

    // Collect candidates
    struct RankTime { int rank; double time; };
    RankTime *candidates = new RankTime[nprocs];
    int ncand = 0;

    for (int i = 0; i < nprocs; i++) {
        if (rank_times[i] > threshold)
            candidates[ncand++] = {i, rank_times[i]};
    }

    // Sort descending by time
    std::sort(candidates, candidates + ncand,
              [](const RankTime &a, const RankTime &b) {
                  return a.time > b.time;
              });

    int count = (ncand < max_heavy) ? ncand : max_heavy;
    for (int i = 0; i < count; i++)
        heavy_ranks[i] = candidates[i].rank;

    delete[] candidates;
    return count;
}

} // namespace InterpLBProfile