// SPDX-License-Identifier: MPL-2.0
#pragma once
#include "vibeqc/tensor/core.hpp"
#include <cmath>

namespace vibeqc::tensor {

// Conservative constant control envelope for the two registered kernels,
// their live Views/IndexPrograms/leases and a bounded cursor. No numerical
// payload is hidden here. The enclosing method must admit this per worker.
inline constexpr std::size_t stream_control_bytes = 65536;
using Complex = std::complex<double>;

// Tokens are minted by a method only AFTER authenticating its actual native
// source, metric and coordinate-frame owners. Text labels are not certificates.
struct SourceStamp {
    Space source, metric, frame;
    std::size_t partition = 0;
    bool operator==(const SourceStamp& other) const {
        return source == other.source && metric == other.metric && frame == other.frame && partition == other.partition;
    }
    void validate() const {
        if (!source.valid() || !metric.valid() || !frame.valid()) throw std::invalid_argument("missing tensor source binding");
    }
};
struct Tile { View values; SourceStamp source; };

// Count-only synchronous stream: no vector of tile keys or retained tile cache.
// The producer's work/read ceilings remain independently authoritative.
class StreamCursor {
    Budget::Reservation control_;
    SourceStamp source_;
    std::size_t maximum_, visits_ = 0;
    bool failed_ = false;
public:
    StreamCursor(const Budget& b, SourceStamp source, std::size_t maximum)
        : control_(b.reserve(sizeof(StreamCursor))), source_(source), maximum_(maximum) { source_.validate(); }
    template<class F> void consume(SourceStamp source, std::size_t sequence, F operation) {
        if (failed_) throw std::logic_error("failed tensor stream cannot resume");
        try {
            if (!(source == source_) || sequence != visits_ || visits_ == maximum_)
                throw std::invalid_argument("tensor stream source, order or visit cap mismatch");
            operation(); ++visits_;
        } catch (...) { failed_ = true; throw; }
    }
    void finish() const {
        if (failed_ || visits_ != maximum_) throw std::logic_error("tensor stream incomplete");
    }
    std::size_t visits() const { return visits_; }
};

inline Complex finite_complex(Complex value) {
    if (!std::isfinite(value.real()) || !std::isfinite(value.imag()))
        throw std::overflow_error("tensor compensated contraction has a non-finite lane");
    return value;
}
// Same independent-lane compensated update as the existing periodic kernels.
// Preserve each multiply association, visit order and final sum+correction.
inline void compensated_add(Complex term, Complex& sum, Complex& correction) {
    finite_complex(term); const auto next = finite_complex(sum + term);
    const auto re = std::abs(sum.real()) >= std::abs(term.real())
        ? (sum.real() - next.real()) + term.real() : (term.real() - next.real()) + sum.real();
    const auto im = std::abs(sum.imag()) >= std::abs(term.imag())
        ? (sum.imag() - next.imag()) + term.imag() : (term.imag() - next.imag()) + sum.imag();
    correction = finite_complex(correction + Complex(re, im)); sum = next;
}
inline void same_source(const SourceStamp& expected, const Tile& tile) {
    expected.validate(); tile.source.validate();
    if (!(expected == tile.source)) throw std::invalid_argument("tensor tile source/metric/frame revision mismatch");
}
inline void complex_access(const View& v, const Budget& b, bool output) {
    if (!b.same(v.allocation()->budget) || v.allocation()->type != DType::complex_fp64 ||
        (output ? v.access()!=Access::read_write || v.conjugate() : v.access()==Access::write))
        throw std::invalid_argument("complex tensor kernel access/dtype/ledger mismatch");
}
inline void separate_output(const View& output, const View& other) {
    if (output.allocation()==other.allocation() || output.overlaps(other))
        throw std::invalid_argument("streamed tensor output aliases another operand");
}
inline const Complex* read_base(const View& v) { return static_cast<const Complex*>(v.allocation()->base); }
inline Complex* write_base(const View& v) { return static_cast<Complex*>(v.allocation()->base); }
inline Complex maybe_conjugate(Complex x, bool conjugate) { return conjugate ? std::conj(x) : x; }

// Registered [p,a] x [p,b] -> [a,b] update. Explicit conjugation is on
// each input View. Used by the all-q VO/OV selected-pair integral leaf.
// Sun doi:10.1063/1.4998644 Eqs.3,4,13,16,21; Nejad
// doi:10.1063/5.0290816 Eqs.2-10,26-28. Chemistry stays in its adapter.
class CompensatedPairPlan {
    Budget::Reservation control_;
    Tile left_, right_;
    View sum_, correction_;
    SourceStamp source_;
    IndexProgram program_;
    std::uint64_t left_version_, right_version_;
public:
    CompensatedPairPlan(Tile left, Tile right, View sum, View correction, SourceStamp source)
        : control_(left.values.allocation()->budget.reserve(sizeof(CompensatedPairPlan))),
          left_(std::move(left)), right_(std::move(right)), sum_(std::move(sum)), correction_(std::move(correction)), source_(source),
          program_(left_.values.allocation()->budget,left_.values,{0,1},right_.values,{0,2},sum_,{1,2}),
          left_version_(left_.values.allocation()->version), right_version_(right_.values.allocation()->version) {
        const auto& b=left_.values.allocation()->budget;
        same_source(source_,left_); same_source(source_,right_);
        complex_access(left_.values,b,false); complex_access(right_.values,b,false);
        complex_access(sum_,b,true); complex_access(correction_,b,true);
        if(correction_.rank()!=2) throw std::invalid_argument("tensor compensation rank mismatch");
        for(std::size_t i=0;i<2;++i) if(!(sum_.axis(i).domain==correction_.axis(i).domain))
            throw std::invalid_argument("tensor compensation domain mismatch");
        for(const auto* output:{&sum_,&correction_}) {
            separate_output(*output,left_.values); separate_output(*output,right_.values);
        }
        separate_output(sum_,correction_);
    }
    const char* kernel() const { return "cpu.complex.compensated_pair.v1"; }
    const IndexProgram& index_program() const { return program_; }
    void execute() const {
        Lease l(left_.values.allocation(),false), r(right_.values.allocation(),false);
        if(left_.values.allocation()->version!=left_version_ || right_.values.allocation()->version!=right_version_)
            throw std::runtime_error("stale complex tensor plan input");
        Lease s(sum_.allocation(),true), c(correction_.allocation(),true);
        const auto& lv=left_.values; const auto& rv=right_.values;
        const auto* lp=read_base(lv); const auto* rp=read_base(rv);
        auto* sp=write_base(sum_); auto* cp=write_base(correction_);
        const auto np=lv.axis(0).domain.extent,na=lv.axis(1).domain.extent,nb=rv.axis(1).domain.extent;
        if(!np || !na || !nb) return;
        // Registered predecessor layout: checked above, no packing or changed
        // reduction tree. Keep this path when general strided lowering loses.
        if(lv.conjugate() && !rv.conjugate() && lv.axis(1).stride==1 && rv.axis(1).stride==1 &&
            lv.axis(0).stride==signed_size(na) && rv.axis(0).stride==signed_size(nb) &&
            sum_.axis(1).stride==1 && correction_.axis(1).stride==1 &&
            sum_.axis(0).stride>0 && correction_.axis(0).stride>0) {
            lp+=lv.offset();rp+=rv.offset();sp+=sum_.offset();cp+=correction_.offset();
            const auto ss=std::size_t(sum_.axis(0).stride),cs=std::size_t(correction_.axis(0).stride);
            for(std::size_t p=0;p<np;++p) for(std::size_t a=0;a<na;++a) for(std::size_t b=0;b<nb;++b)
                compensated_add(std::conj(lp[p*na+a])*rp[p*nb+b],sp[a*ss+b],cp[a*cs+b]);
            return;
        }
        for(std::size_t p=0;p<lv.axis(0).domain.extent;++p)
            for(std::size_t a=0;a<lv.axis(1).domain.extent;++a)
                for(std::size_t b=0;b<rv.axis(1).domain.extent;++b) {
                    const auto li=lv.offset()+stride_span(p,lv.axis(0).stride)+stride_span(a,lv.axis(1).stride);
                    const auto ri=rv.offset()+stride_span(p,rv.axis(0).stride)+stride_span(b,rv.axis(1).stride);
                    const auto si=sum_.offset()+stride_span(a,sum_.axis(0).stride)+stride_span(b,sum_.axis(1).stride);
                    const auto ci=correction_.offset()+stride_span(a,correction_.axis(0).stride)+stride_span(b,correction_.axis(1).stride);
                    compensated_add(maybe_conjugate(lp[li],lv.conjugate())*maybe_conjugate(rp[ri],rv.conjugate()),sp[si],cp[ci]);
                }
    }
};

// Fused logical X[l,r,mu,nu]=conj(C[l,mu])*C[r,nu], then
// B[p,l,r]+=X[l,r,mu,nu]*A[p,mu,nu]. X is never allocated. The method
// supplies a contiguous, ordered flattening of the AO-pair domain for each
// tile. All arithmetic remains complex; normalization is left to the method.
// This is the actual streamed contraction feeding Gaussian pair-MP2 panels.
class LocalPanelPlan {
    Budget::Reservation control_;
    Tile factors_, left_, right_;
    View sum_, correction_;
    SourceStamp source_;
    Schema<4> product_;
    Schema<3> logical_factors_;
    IndexProgram product_ir_, reduction_ir_;
    std::uint64_t fv_, lv_, rv_;
    static Schema<4> product(const View& l,const View& r) {
        if(l.rank()!=2 || r.rank()!=2) throw std::invalid_argument("local coefficient rank mismatch");
        return {{Axis{l.axis(0).domain,0},Axis{r.axis(0).domain,0},Axis{l.axis(1).domain,0},Axis{r.axis(1).domain,0}}};
    }
    static Schema<3> logical(const View& f,const View& l,const View& r) {
        if(f.rank()!=2) throw std::invalid_argument("local factor tile rank mismatch");
        return {{Axis{f.axis(0).domain,0},Axis{l.axis(1).domain,0},Axis{r.axis(1).domain,0}}};
    }
public:
    LocalPanelPlan(Tile factors,Tile left,Tile right,View sum,View correction,SourceStamp source)
        : control_(factors.values.allocation()->budget.reserve(sizeof(LocalPanelPlan))),
          factors_(std::move(factors)),left_(std::move(left)),right_(std::move(right)),sum_(std::move(sum)),correction_(std::move(correction)),source_(source),
          product_(product(left_.values,right_.values)), logical_factors_(logical(factors_.values,left_.values,right_.values)),
          product_ir_(sum_.allocation()->budget,left_.values,{1,3},right_.values,{2,4},product_,{1,2,3,4}),
          reduction_ir_(sum_.allocation()->budget,product_,{1,2,3,4},logical_factors_,{0,3,4},sum_,{0,1,2}),
          fv_(factors_.values.allocation()->version),lv_(left_.values.allocation()->version),rv_(right_.values.allocation()->version) {
        const auto& b=sum_.allocation()->budget;
        for(const auto* tile:{&factors_,&left_,&right_}) { same_source(source_,*tile); complex_access(tile->values,b,false); }
        complex_access(sum_,b,true); complex_access(correction_,b,true);
        if(correction_.rank()!=3) throw std::invalid_argument("panel compensation rank mismatch");
        for(std::size_t i=0;i<3;++i) if(!(sum_.axis(i).domain==correction_.axis(i).domain))
            throw std::invalid_argument("panel compensation domain mismatch");
        for(const auto* out:{&sum_,&correction_})
            for(const auto* in:{&factors_,&left_,&right_}) separate_output(*out,in->values);
        separate_output(sum_,correction_);
        const auto& pairs=factors_.values.axis(1).domain;
        const auto& ao=left_.values.axis(1).domain;
        if(!(ao==right_.values.axis(1).domain) || ao.selection || ao.begin!=0 || ao.step!=1 ||
            pairs.parent!=ao.space || pairs.parent_coordinate!=0 || pairs.selection || pairs.step!=1 ||
            pairs.begin>mul(ao.extent,ao.extent) || pairs.extent>mul(ao.extent,ao.extent)-pairs.begin)
            throw std::invalid_argument("local panel ordered AO-pair map mismatch");
    }
    const char* kernel() const { return "cpu.complex.local_panel.compensated.v1"; }
    const IndexProgram& product_program() const { return product_ir_; }
    const IndexProgram& reduction_program() const { return reduction_ir_; }
    void execute() const {
        Lease f(factors_.values.allocation(),false),l(left_.values.allocation(),false),r(right_.values.allocation(),false);
        if(factors_.values.allocation()->version!=fv_ || left_.values.allocation()->version!=lv_ || right_.values.allocation()->version!=rv_)
            throw std::runtime_error("stale local panel input");
        Lease s(sum_.allocation(),true),c(correction_.allocation(),true);
        const auto& fv=factors_.values; const auto& lv=left_.values; const auto& rv=right_.values;
        const auto* fp=read_base(fv);const auto* lp=read_base(lv);const auto* rp=read_base(rv);
        auto* sp=write_base(sum_);auto* cp=write_base(correction_);
        const auto n=lv.axis(1).domain.extent, first=fv.axis(1).domain.begin;
        const auto nt=fv.axis(1).domain.extent,np=fv.axis(0).domain.extent,nl=lv.axis(0).domain.extent,nr=rv.axis(0).domain.extent;
        if(!nt || !np || !nl || !nr) return;
        if(lv.conjugate() && !rv.conjugate() && !fv.conjugate() &&
            fv.axis(0).stride==signed_size(nt) && fv.axis(1).stride==1 &&
            lv.axis(0).stride==signed_size(n) && rv.axis(0).stride==signed_size(n) &&
            lv.axis(1).stride==1 && rv.axis(1).stride==1 &&
            sum_.axis(0).stride==signed_size(nl*nr) && correction_.axis(0).stride==signed_size(nl*nr) &&
            sum_.axis(1).stride==signed_size(nr) && correction_.axis(1).stride==signed_size(nr) &&
            sum_.axis(2).stride==1 && correction_.axis(2).stride==1) {
            fp+=fv.offset();lp+=lv.offset();rp+=rv.offset();sp+=sum_.offset();cp+=correction_.offset();
            for(std::size_t at=0;at<nt;++at) {
                const auto mu=(first+at)/n,nu=(first+at)%n;
                for(std::size_t p=0;p<np;++p) {
                    const auto factor=fp[p*nt+at];
                    for(std::size_t l=0;l<nl;++l) for(std::size_t r=0;r<nr;++r) {
                        const auto index=(p*nl+l)*nr+r;
                        compensated_add((std::conj(lp[l*n+mu])*rp[r*n+nu])*factor,sp[index],cp[index]);
                    }
                }
            }
            return;
        }
        for(std::size_t at=0;at<fv.axis(1).domain.extent;++at) {
            const auto mu=(first+at)/n,nu=(first+at)%n;
            for(std::size_t p=0;p<fv.axis(0).domain.extent;++p) {
                const auto factor=maybe_conjugate(fp[fv.offset()+stride_span(p,fv.axis(0).stride)+stride_span(at,fv.axis(1).stride)],fv.conjugate());
                for(std::size_t a=0;a<lv.axis(0).domain.extent;++a) for(std::size_t b=0;b<rv.axis(0).domain.extent;++b) {
                    const auto x=maybe_conjugate(lp[lv.offset()+stride_span(a,lv.axis(0).stride)+stride_span(mu,lv.axis(1).stride)],lv.conjugate());
                    const auto y=maybe_conjugate(rp[rv.offset()+stride_span(b,rv.axis(0).stride)+stride_span(nu,rv.axis(1).stride)],rv.conjugate());
                    const auto si=sum_.offset()+stride_span(p,sum_.axis(0).stride)+stride_span(a,sum_.axis(1).stride)+stride_span(b,sum_.axis(2).stride);
                    const auto ci=correction_.offset()+stride_span(p,correction_.axis(0).stride)+stride_span(a,correction_.axis(1).stride)+stride_span(b,correction_.axis(2).stride);
                    compensated_add((x*y)*factor,sp[si],cp[ci]);
                }
            }
        }
    }
};
} // namespace vibeqc::tensor