#include "aner/ml.hpp"
#include "aner/parallel_jobs.hpp"
#include "aner/random.hpp"
#include <algorithm>
#include <bit>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>
#if defined(__unix__) || defined(__APPLE__)
#include <sys/resource.h>
#endif
namespace ml=aner::ml; namespace nn=aner::nn; namespace par=aner::parallel;
namespace {
using Clock=std::chrono::steady_clock;
struct Arguments{std::size_t n=4096,q=32768,d=8,k=4,cap=20,workers=1,repeats=7,warmups=2;std::int64_t seed=42;std::string operation="fit",mode="serial",distribution="balanced";bool workspace=false,rss=false;};
Arguments parse(int argc,char**argv){Arguments a;for(int i=1;i<argc;i+=2){if(i+1==argc)throw std::runtime_error("missing flag value");const std::string f=argv[i],v=argv[i+1];if(f=="--operation")a.operation=v;else if(f=="--mode")a.mode=v;else if(f=="--distribution")a.distribution=v;else{std::size_t used=0;auto n=std::stoull(v,&used);if(used!=v.size())throw std::runtime_error("invalid integer");if(f=="--n")a.n=n;else if(f=="--q")a.q=n;else if(f=="--d")a.d=n;else if(f=="--k")a.k=n;else if(f=="--cap")a.cap=n;else if(f=="--workers")a.workers=n;else if(f=="--repeats")a.repeats=n;else if(f=="--warmups")a.warmups=n;else if(f=="--workspace")a.workspace=n!=0;else if(f=="--rss-only")a.rss=n!=0;else if(f=="--seed")a.seed=static_cast<std::int64_t>(n);else throw std::runtime_error("unknown flag");}}
 if(!a.n||!a.q||!a.d||!a.k||a.k>a.n||!a.cap||a.cap>1000||!a.workers||!a.repeats||a.repeats>100||a.warmups>100)throw std::runtime_error("invalid dimensions");if(a.n>1'000'000/a.d||a.q>1'000'000/a.d||(a.cap+1)>50'000'000/a.n/a.k/a.d||a.q>50'000'000/a.k/a.d)throw std::runtime_error("work limit");if(a.operation!="fit"&&a.operation!="predict")throw std::runtime_error("unknown operation");if(a.mode!="serial"&&a.mode!="context"&&a.mode!="cold")throw std::runtime_error("unknown mode");if(a.distribution!="balanced"&&a.distribution!="overlap"&&a.distribution!="skew")throw std::runtime_error("unknown distribution");if(a.workers>par::discover_cpu_budget().hard_limit)throw std::runtime_error("CPU allocation");
#ifdef ANER_KMEANS_BASELINE
 if(a.mode!="serial"||a.workspace)throw std::runtime_error("legacy supports serial without byte admission");
#endif
 return a;}
void exact(const std::vector<double>&a,const std::vector<double>&b){if(a.size()!=b.size())throw std::runtime_error("shape mismatch");for(std::size_t i=0;i<a.size();++i)if(std::bit_cast<std::uint64_t>(a[i])!=std::bit_cast<std::uint64_t>(b[i]))throw std::runtime_error("bitwise result mismatch");}
void same(const ml::KMeansModel&a,const ml::KMeansModel&b){exact(ml::centers(a).values(),ml::centers(b).values());exact(ml::labels(a).values(),ml::labels(b).values());if(std::bit_cast<std::uint64_t>(ml::inertia(a))!=std::bit_cast<std::uint64_t>(ml::inertia(b))||ml::iterations(a)!=ml::iterations(b)||ml::converged(a)!=ml::converged(b))throw std::runtime_error("model metadata mismatch");}
double median(std::vector<double>a){std::sort(a.begin(),a.end());const auto m=a.size()/2;return a.size()%2?a[m]:(a[m-1]+a[m])/2;}
std::uint64_t hash(const std::vector<double>&values,std::uint64_t h=14695981039346656037ULL){for(const auto v:values){auto b=std::bit_cast<std::uint64_t>(v);for(int i=0;i<8;++i){h^=(b>>(8*i))&255;h*=1099511628211ULL;}}return h;}
std::size_t peak_rss(){
#if defined(__unix__) || defined(__APPLE__)
 rusage r{};if(getrusage(RUSAGE_SELF,&r))throw std::runtime_error("getrusage failed");
#if defined(__APPLE__)
 return static_cast<std::size_t>(r.ru_maxrss);
#else
 return static_cast<std::size_t>(r.ru_maxrss)*1024;
#endif
#else
 return 0;
#endif
}
}
int main(int argc,char**argv){try{const auto a=parse(argc,argv);std::uint64_t random=20261030;std::vector<double>data(a.n*a.d),query(a.q*a.d);auto generate=[&](std::vector<double>&v,std::size_t rows){for(std::size_t r=0;r<rows;++r){auto cluster=r%a.k;if(a.distribution=="skew"&&r%10!=0)cluster=0;for(std::size_t c=0;c<a.d;++c){const auto bits=aner::rng::mix(cluster*131+c*17);const auto center=static_cast<double>(bits%1025)/32-16;const auto noise=static_cast<double>(aner::rng::next(random)%513)/256-1;v[r*a.d+c]=center+(a.distribution=="overlap"?8:1)*noise;}}};generate(data,a.n);generate(query,a.q);const auto x=nn::tensor(a.n,a.d,std::move(data)),input=nn::tensor(a.q,a.d,std::move(query));par::Context context;
#ifndef ANER_KMEANS_BASELINE
 if(a.mode=="context")context=par::local_context(a.workers);
 ml::KMeansOptions options;if(a.mode=="context")options.execution=&context;
#endif
 const auto fit=[&]()->ml::KMeansModel{
#ifdef ANER_KMEANS_BASELINE
  return ml::fit_kmeans(x,a.k,a.cap,0,a.seed);
#else
  if(a.mode=="cold"){auto pool=par::local_context(a.workers);auto o=options;o.execution=&pool;return ml::fit_kmeans(x,a.k,a.cap,0,a.seed,o);}return ml::fit_kmeans(x,a.k,a.cap,0,a.seed,options);
#endif
 };
 // RSS-only children avoid correctness copies and repeated retained results;
 // prediction lifetime peak still includes its prerequisite serial fit.
 if(a.rss){auto m=a.operation=="fit"?fit():ml::fit_kmeans(x,a.k,a.cap,0,a.seed);if(a.operation=="predict"){
#ifdef ANER_KMEANS_BASELINE
 (void)ml::predict(m,input);
#else
 if(a.mode=="cold"){auto pool=par::local_context(a.workers);auto o=options;o.execution=&pool;(void)ml::predict(m,input,o);}else (void)ml::predict(m,input,options);
#endif
 }std::cout<<"{\"peak_rss_bytes\":"<<peak_rss()<<",\"iterations\":"<<ml::iterations(m)<<"}\n";return 0;}
 const auto expected=ml::fit_kmeans(x,a.k,a.cap,0,a.seed);const auto prediction=ml::predict(expected,input).values();same(expected,fit());exact(ml::labels(expected).values(),ml::predict(expected,x).values());const auto predict=[&]()->nn::Tensor{
#ifdef ANER_KMEANS_BASELINE
 return ml::predict(expected,input);
#else
 if(a.mode=="cold"){auto pool=par::local_context(a.workers);auto o=options;o.execution=&pool;return ml::predict(expected,input,o);}return ml::predict(expected,input,options);
#endif
 };exact(predict().values(),prediction);std::size_t charged=0;
#ifndef ANER_KMEANS_BASELINE
 if(a.workspace){if(a.mode=="cold")throw std::runtime_error("workspace probe uses reusable or serial context");std::size_t low=1,high=64*1024*1024;while(low<high){const auto middle=low+(high-low)/2;auto o=options;o.max_workspace_bytes=middle;try{if(a.operation=="fit")(void)ml::fit_kmeans(x,a.k,a.cap,0,a.seed,o);else (void)ml::predict(expected,input,o);high=middle;}catch(const ml::Error&e){if(std::string(e.what()).find("workspace")==std::string::npos)throw;low=middle+1;}}charged=low;auto o=options;o.max_workspace_bytes=low;if(a.operation=="fit")same(expected,ml::fit_kmeans(x,a.k,a.cap,0,a.seed,o));else exact(ml::predict(expected,input,o).values(),prediction);o.max_workspace_bytes=low-1;bool failed=false;try{if(a.operation=="fit")(void)ml::fit_kmeans(x,a.k,a.cap,0,a.seed,o);else (void)ml::predict(expected,input,o);}catch(const ml::Error&e){failed=std::string(e.what()).find("workspace")!=std::string::npos;}if(!failed)throw std::runtime_error("budget minus one was accepted");}
#endif
 std::vector<double>samples;for(std::size_t i=0;i<a.warmups+a.repeats;++i){const auto start=Clock::now();if(a.operation=="fit"){const auto model=fit();const auto end=Clock::now();if(i>=a.warmups)samples.push_back(std::chrono::duration<double>(end-start).count());same(expected,model);}else{const auto result=predict();const auto end=Clock::now();if(i>=a.warmups)samples.push_back(std::chrono::duration<double>(end-start).count());exact(result.values(),prediction);}}
 const auto med=median(samples);std::vector<double>deviations;for(auto v:samples)deviations.push_back(std::abs(v-med));const auto digest=hash(ml::labels(expected).values(),hash(ml::centers(expected).values()));
 std::cout<<std::setprecision(17)<<"{\"flavor\":\""
#ifdef ANER_KMEANS_BASELINE
 <<"legacy"
#else
 <<"corrected"
#endif
 <<"\",\"mode\":\""<<a.mode<<"\",\"workers\":"<<(a.mode=="serial"?0:a.workers)<<",\"raw_seconds\":[";for(std::size_t i=0;i<samples.size();++i)std::cout<<(i?",":"")<<samples[i];std::cout<<"],\"median_seconds\":"<<med<<",\"mad_seconds\":"<<median(deviations)<<",\"iterations\":"<<ml::iterations(expected)<<",\"converged\":"<<(ml::converged(expected)?"true":"false")<<",\"inertia\":"<<ml::inertia(expected)<<",\"result_digest_fnv64\":\""<<std::hex<<digest<<std::dec<<"\",\"bitwise_serial_equal\":"
#ifdef ANER_KMEANS_BASELINE
 <<"null"
#else
 <<"true"
#endif
 <<",\"charged_workspace_bytes\":"<<charged<<",\"exact_budget_accepted\":"<<(a.workspace?"true":"null")<<",\"minus_one_rejected\":"<<(a.workspace?"true":"null")<<",\"active_lanes\":"<<(a.mode=="serial"?1:std::min(a.workers,a.operation=="fit"?std::max(a.n,a.k*a.d):a.q))<<",\"cpu_hard_limit\":"<<par::discover_cpu_budget().hard_limit<<"}\n";
 }catch(const std::exception&e){std::cerr<<e.what()<<'\n';return 1;}}