#include "aner/gen.hpp"
#include "aner/nn.hpp"
#include <chrono>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>
namespace gen=aner::gen;namespace nn=aner::nn;namespace token=aner::token;
int main(int argc,char**argv){try{
 std::string device="cuda",dtype="Float32",name="tiny";int batch=8,reps=8,warmups=2;bool graph=false;
 for(int i=1;i<argc;i+=2){if(i+1>=argc)throw std::runtime_error("flag needs value");std::string f=argv[i],v=argv[i+1];if(f=="--device")device=v;else if(f=="--dtype")dtype=v;else if(f=="--batch-size")batch=std::stoi(v);else if(f=="--case")name=v;else if(f=="--repetitions")reps=std::stoi(v);else if(f=="--warmups")warmups=std::stoi(v);else if(f=="--cuda-graph")graph=v=="true";else throw std::runtime_error("unknown flag");}
 if((batch!=1&&batch!=8)||(name!="tiny"&&name!="longdecode")||reps<2||warmups<1||!gen::device_supported(device,dtype))throw std::runtime_error("unsupported explicit native benchmark configuration");
 const auto model=gen::to(gen::decoder(token::byte_tokenizer(),{96,32,4,2,20261015,0}),device,dtype);
 std::vector<std::vector<std::int64_t>> prompts(batch);std::vector<std::int64_t> seeds(batch);
 for(int row=0;row<batch;++row){prompts[row]={token::bos_id};for(int j=1;j<8;++j)prompts[row].push_back(97+(37*j+11*row+5)%26);seeds[row]=7301+101*row;}
 std::vector<std::vector<std::int64_t>> expected;const std::int64_t budget=name=="tiny"?8:64;
 using Clock=std::chrono::steady_clock;
 for(int rep=-warmups;rep<reps;++rep){nn::synchronize(device);const auto before=nn::memory_stats();nn::reset_memory_peaks();const auto start=Clock::now();
  std::vector<std::vector<std::int64_t>> out;
  if(batch==1&&!graph)out={gen::generate(model,prompts[0],budget,0.,0,1.,seeds[0],true)};
  else out=gen::generate_batch(model,prompts,budget,0.,0,1.,seeds,true,graph);
  nn::synchronize(device);const auto seconds=std::chrono::duration<double>(Clock::now()-start).count();const auto after=nn::memory_stats();
  if(out.size()!=std::size_t(batch)||before.all_buffers.current_bytes!=after.all_buffers.current_bytes)throw std::runtime_error("row inventory or lifetime mismatch");
  for(const auto&r:out)if(r.empty()||r.size()>std::size_t(budget))throw std::runtime_error("generated token budget mismatch");
  if(expected.empty())expected=out;else if(out!=expected)throw std::runtime_error("repeat generated IDs differ");
  if(rep<0)continue;
  std::cout<<std::setprecision(17)<<"{\"schema\":\"aner.torchtpu-natural/1\",\"device\":\""<<device<<"\",\"dtype\":\""<<dtype<<"\",\"case\":\""<<name<<"\",\"batch_size\":"<<batch<<",\"cuda_graph\":"<<(graph?"true":"false")<<",\"iteration\":"<<rep<<",\"seconds\":"<<seconds<<",\"tracked_peak_increment_bytes\":"<<after.all_buffers.peak_bytes-before.all_buffers.current_bytes<<",\"activation_readback_peak_increment_bytes\":"<<after.host_readback_cache.peak_bytes-before.host_readback_cache.current_bytes<<",\"ids\":[";
  for(std::size_t r=0;r<out.size();++r){if(r)std::cout<<',';std::cout<<'[';for(std::size_t c=0;c<out[r].size();++c)std::cout<<(c?",":"")<<out[r][c];std::cout<<']';}std::cout<<"]}\n";
 }
 return 0;
}catch(const std::exception&e){std::cerr<<e.what()<<'\n';return 1;}}