// Coverage grammars and multi-component routing, including time-major streams.
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <algorithm>
#include <cstdint>
#include <memory>
#include <stdexcept>
#include <vector>
#include "parallel.hpp"
#include "token_streams.hpp"

namespace py=pybind11;
namespace {
using State=std::vector<int64_t>;
using States=std::vector<State>;
class Grammar {
 public:
  int64_t horizon,vocabulary;
  virtual ~Grammar()=default;
  virtual State initial()const=0;
  virtual State advance(const State&,int64_t)const=0;
  virtual int64_t covered(const State&)const=0;
  virtual int64_t next_dimension(const State&)const=0;
  virtual bool allowed(const State&,int64_t)const=0;
  virtual std::shared_ptr<Grammar> clone(int64_t)const=0;
  py::array_t<bool> mask(const State& state)const {
    py::array_t<bool> result(vocabulary);
    py::gil_scoped_release release;
    for(int64_t token=0;token<vocabulary;++token)result.mutable_data()[token]=allowed(state,token);
    return result;
  }
};
class SetGrammar:public Grammar {
  uint64_t full_;
  using Patterns=std::vector<std::vector<std::pair<int64_t,uint64_t>>>;
  std::shared_ptr<const std::vector<int64_t>> roots_;
  std::shared_ptr<const Patterns> patterns_;
  std::shared_ptr<const std::vector<bool>> scope_;
 public:
  SetGrammar(int64_t steps,uint64_t full,std::vector<int64_t> roots,
             std::vector<std::vector<std::pair<int64_t,uint64_t>>> patterns)
      :full_(full),roots_(std::make_shared<const std::vector<int64_t>>(std::move(roots))),
       patterns_(std::make_shared<const Patterns>(std::move(patterns))) {
    if(steps<1 || !full || roots_->size()!=patterns_->size())throw std::invalid_argument("invalid grammar metadata");
    horizon=steps;vocabulary=roots_->size();
    std::vector<bool> scopes;
    for(const auto& pattern:*patterns_) {
      bool scope=true;
      for(auto[offset,support]:pattern) {
        if(offset<0)throw std::invalid_argument("negative grammar pattern offset");
        scope=scope && !(support & ~full_);
      }
      scopes.push_back(scope);
    }
    scope_=std::make_shared<const std::vector<bool>>(std::move(scopes));
  }
  std::shared_ptr<Grammar> clone(int64_t steps)const override {
    if(steps<1)throw std::invalid_argument("horizon must be positive");
    auto result=std::make_shared<SetGrammar>(*this);result->horizon=steps;return result;
  }
  State initial()const override{return State(horizon,0);}
  int64_t covered(const State& state)const override {
    if(static_cast<int64_t>(state.size())!=horizon)throw std::invalid_argument("invalid grammar state");
    for(int64_t i=0;i<horizon;++i)if(static_cast<uint64_t>(state[i])!=full_)return i;
    return horizon;
  }
  int64_t next_dimension(const State& state)const override {
    int64_t front=covered(state);if(front==horizon)return -1;
    uint64_t missing=full_ & ~static_cast<uint64_t>(state[front]);
    int64_t d=0;while(!(missing & 1)){missing>>=1;++d;}return d;
  }
  bool allowed(const State& state,int64_t token)const override {
    if(token<0 || token>=vocabulary || !(*scope_)[token])return false;
    const int64_t front=covered(state);
    if(front==horizon || (*roots_)[token]!=next_dimension(state))return false;
    for(auto[offset,support]:(*patterns_)[token]) {
      if(front+offset>=horizon || (static_cast<uint64_t>(state[front+offset]) & support))return false;
    }
    return true;
  }
  State advance(const State& state,int64_t token)const override {
    if(!allowed(state,token))return {};
    State result=state;int64_t front=covered(state);
    for(auto[offset,support]:(*patterns_)[token])result[front+offset]|=static_cast<int64_t>(support);
    return result;
  }
};
class PlainGrammar:public Grammar {
  std::shared_ptr<const std::vector<int64_t>> lengths_,dims_;
  std::shared_ptr<const std::vector<bool>> phase_;
 public:
  PlainGrammar(int64_t steps,std::vector<int64_t> lengths,std::vector<int64_t> dims,
                py::array_t<bool,py::array::c_style|py::array::forcecast> phase)
      :lengths_(std::make_shared<const std::vector<int64_t>>(std::move(lengths))),
       dims_(std::make_shared<const std::vector<int64_t>>(std::move(dims))) {
    horizon=steps;vocabulary=lengths_->size();
    if(steps<1 || dims_->empty() || phase.ndim()!=2 || phase.shape(0)!=static_cast<int64_t>(dims_->size()) || phase.shape(1)!=vocabulary)
      throw std::invalid_argument("invalid plain-BPE grammar metadata");
    for(auto value:*lengths_)if(value<1)throw std::invalid_argument("invalid token expansion");
    phase_=std::make_shared<const std::vector<bool>>(phase.data(),phase.data()+phase.size());
  }
  std::shared_ptr<Grammar> clone(int64_t steps)const override {
    if(steps<1)throw std::invalid_argument("horizon must be positive");
    auto result=std::make_shared<PlainGrammar>(*this);result->horizon=steps;return result;
  }
  State initial()const override{return {0};}
  int64_t covered(const State& state)const override {
    if(state.size()!=1 || state[0]<0)throw std::invalid_argument("invalid grammar state");
    return state[0]/dims_->size();
  }
  int64_t next_dimension(const State& state)const override{return (*dims_)[state[0]%dims_->size()];}
  bool allowed(const State& state,int64_t token)const override {
    return token>=0 && token<vocabulary && state[0]+(*lengths_)[token]<=horizon*static_cast<int64_t>(dims_->size()) && (*phase_)[(state[0]%dims_->size())*vocabulary+token];
  }
  State advance(const State& state,int64_t token)const override {
    if(!allowed(state,token))return {};
    return {state[0]+(*lengths_)[token]};
  }
};
using Grammars=std::vector<std::shared_ptr<Grammar>>;
int64_t next_component(const Grammars& grammars,const States& states,bool time_major) {
  int64_t selected=grammars.size(),minimum=INT64_MAX;
  for(size_t i=0;i<grammars.size();++i) {
    const int64_t front=grammars[i]->covered(states[i]);
    if(front>=grammars[i]->horizon)continue;
    if(!time_major)return i;
    if(front<minimum){minimum=front;selected=i;}
  }
  return selected;
}
States initial(const Grammars& grammars) {
  States states;for(const auto& grammar:grammars)states.push_back(grammar->initial());return states;
}
bool advance(const Grammars& grammars,const std::vector<int64_t>& offsets,States& states,int64_t index,int64_t token) {
  if(index<0 || index>=static_cast<int64_t>(grammars.size()))return false;
  State next=grammars[index]->advance(states[index],token-offsets[index]);
  if(next.empty())return false;
  states[index]=std::move(next);
  return true;
}
py::tuple state_tuple(const State& state) {
  py::tuple result(state.size());for(size_t i=0;i<state.size();++i)result[i]=state[i];return result;
}
py::object route_advance(const Grammars& grammars,const std::vector<int64_t>& offsets,
                         bool time_major,int64_t index,States states,int64_t token) {
  if(states.size()!=grammars.size() || offsets.size()!=grammars.size())throw std::invalid_argument("invalid routed metadata");
  if(!advance(grammars,offsets,states,index,token))return py::none();
  py::tuple local(states.size());for(size_t i=0;i<states.size();++i)local[i]=state_tuple(states[i]);
  return py::make_tuple(next_component(grammars,states,time_major),local);
}
std::vector<int64_t> interleave(const Grammars& grammars,const std::vector<int64_t>& offsets,
                                bool time_major,const std::vector<std::vector<int64_t>>& streams) {
  if(streams.size()!=grammars.size() || offsets.size()!=grammars.size())throw std::invalid_argument("one stream per component is required");
  States states=initial(grammars);std::vector<size_t> cursor(streams.size());std::vector<int64_t> output;
  int64_t index;
  while((index=next_component(grammars,states,time_major))<static_cast<int64_t>(grammars.size())) {
    if(cursor[index]==streams[index].size())throw std::invalid_argument("component stream ends before complete coverage");
    const int64_t token=streams[index][cursor[index]++]+offsets[index];
    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed component stream");
    output.push_back(token);
  }
  for(size_t i=0;i<cursor.size();++i)if(cursor[i]!=streams[i].size())throw std::invalid_argument("extra component tokens");
  return output;
}
std::vector<std::vector<int64_t>> split(const Grammars& grammars,const std::vector<int64_t>& offsets,
                                       bool time_major,const std::vector<int64_t>& tokens) {
  if(offsets.size()!=grammars.size())throw std::invalid_argument("invalid routed offsets");
  States states=initial(grammars);std::vector<std::vector<int64_t>> output(grammars.size());
  for(int64_t token:tokens) {
    int64_t index=next_component(grammars,states,time_major);
    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed token coverage");
    output[index].push_back(token-offsets[index]);
  }
  if(next_component(grammars,states,time_major)!=static_cast<int64_t>(grammars.size()))
    throw std::invalid_argument("Incomplete action token stream");
  return output;
}

using Streams=std::vector<std::vector<int64_t>>;
using ComponentStreams=std::vector<Streams>;

Streams interleave_batch(const Grammars& grammars,const std::vector<int64_t>& offsets,
                        bool time_major,const ComponentStreams& streams) {
  if(streams.size()!=grammars.size() || streams.empty() || streams.front().empty())
    throw std::invalid_argument("one non-empty batch per component is required");
  const int64_t batch=streams.front().size();
  for(const auto& rows:streams)if(static_cast<int64_t>(rows.size())!=batch)
    throw std::invalid_argument("component batches must have the same size");
  Streams output(batch);
  py::gil_scoped_release release;
  ac2::parallel_for(batch,0,[&](int64_t b) {
    Streams local;for(const auto& rows:streams)local.push_back(rows[b]);
    output[b]=interleave(grammars,offsets,time_major,local);
  });
  return output;
}

py::tuple decode_batch(const Grammars& prototypes,const std::vector<int64_t>& offsets,
                       bool time_major,const py::iterable& streams,
                       py::array_t<int64_t,py::array::c_style|py::array::forcecast> counts,
                       int64_t cells_per_step) {
  if(prototypes.empty() || counts.ndim()!=1 || cells_per_step<1)
    throw std::invalid_argument("invalid routed batch metadata");
  const auto rows=ac2::checked_token_streams(streams,counts.size(),true);
  int64_t steps=-1;
  ComponentStreams output(prototypes.size(),Streams(rows.size()));
  {
    py::gil_scoped_release release;
    for(const auto& row:rows) {
      int64_t cells=0;
      for(int64_t token:row) {
        if(counts.data()[token]<1)throw std::invalid_argument("token is outside the routed profile vocabularies");
        cells+=counts.data()[token];
      }
      if(cells%cells_per_step || cells/cells_per_step<1)
        throw std::invalid_argument("Incomplete action token stream: " +
            std::to_string(cells_per_step-cells%cells_per_step) +
            " more coordinate values are needed to finish a step. "
            "Constrain generation with codec.grammar(H, fps=...).next_token_mask(prefix) "
            "and check is_complete(tokens) for the intended H before decoding. "
            "Use codec.batch_decode(..., errors='return') to isolate invalid rows; "
            "do not pad truncated tokens with zeros.");
      if(steps>=0 && steps!=cells/cells_per_step)
        throw std::invalid_argument("all token streams must decode to the same length; use codec.batch_decode(...) for independent variable-length outputs");
      steps=cells/cells_per_step;
    }
    Grammars grammars;for(const auto& prototype:prototypes)grammars.push_back(prototype->clone(steps));
    ac2::parallel_for(rows.size(),0,[&](int64_t b) {
      auto local=split(grammars,offsets,time_major,rows[b]);
      for(size_t i=0;i<local.size();++i)output[i][b]=std::move(local[i]);
    });
  }
  return py::make_tuple(steps,output);
}

std::vector<int64_t> complete(const Grammars& grammars,const std::vector<int64_t>& offsets,
                              bool time_major,std::vector<int64_t> tokens,int64_t steps,
                              const std::vector<std::vector<int64_t>>& zero_atoms) {
  if(grammars.empty() || offsets.size()!=grammars.size() || zero_atoms.size()!=grammars.size() ||
     steps<1 || steps>grammars.front()->horizon)throw std::invalid_argument("invalid prefix completion metadata");
  py::gil_scoped_release release;
  States states=initial(grammars);
  for(int64_t token:tokens) {
    int64_t index=next_component(grammars,states,time_major);
    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid routed token prefix");
  }
  for(size_t i=0;i<grammars.size();++i)if(grammars[i]->covered(states[i])<steps)
    throw std::invalid_argument("action prefix does not cover the requested frames");
  int64_t index;
  while((index=next_component(grammars,states,time_major))<static_cast<int64_t>(grammars.size())) {
    int64_t dimension=grammars[index]->next_dimension(states[index]);
    if(dimension<0 || dimension>=static_cast<int64_t>(zero_atoms[index].size()))
      throw std::invalid_argument("missing zero-motion primitive atom");
    const int64_t token=offsets[index]+zero_atoms[index][dimension];
    if(!advance(grammars,offsets,states,index,token))throw std::invalid_argument("invalid zero-motion suffix");
    tokens.push_back(token);
  }
  return tokens;
}
}  // namespace

void register_grammar(py::module_& module) {
  py::class_<Grammar,std::shared_ptr<Grammar>>(module,"CoverageGrammar")
      .def("advance",[](const Grammar& g,const State& s,int64_t token)->py::object {
        State next=g.advance(s,token);return next.empty()?py::object(py::none()):py::object(state_tuple(next));
      })
      .def("mask",&Grammar::mask).def("covered",&Grammar::covered).def("clone",&Grammar::clone);
  py::class_<SetGrammar,Grammar,std::shared_ptr<SetGrammar>>(module,"SetGrammar")
      .def(py::init<int64_t,uint64_t,std::vector<int64_t>,std::vector<std::vector<std::pair<int64_t,uint64_t>>>>());
  py::class_<PlainGrammar,Grammar,std::shared_ptr<PlainGrammar>>(module,"PlainGrammar")
      .def(py::init<int64_t,std::vector<int64_t>,std::vector<int64_t>,py::array_t<bool,py::array::c_style|py::array::forcecast>>());
  module.def("route_advance",&route_advance);
  module.def("route_interleave",&interleave,py::call_guard<py::gil_scoped_release>());
  module.def("route_split",&split,py::call_guard<py::gil_scoped_release>());
  module.def("route_interleave_batch",&interleave_batch);
  module.def("route_decode_batch",&decode_batch);
  module.def("route_complete",&complete);
}