#include "ling3/linear.h"
#include "ling3/cpu_kernels.h"
#include "ling3/quantization.h"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <stdexcept>

int main(int argc,char**argv)try{
    if(argc!=2)throw std::runtime_error("usage: linear-bridge-check OFFICIAL.l3r");
    setenv("LING3_OFFICIAL_EXECUTION","w8",1);
    setenv("LING3_KEEP_SOURCE_WEIGHTS","1",1);
    ling3::ModelPackage package(argv[1]);
    for(int lanes:{1,3}){
    for(const std::string base:{"model.layers.1.mlp.experts.0.gate_up","model.layers.1.mlp.shared_experts.gate_up"}){
        const auto&t=package.tensor(base+".weight");int k=t.entry->dims[0],n=t.entry->dims[1];
        std::vector<int> cores=lanes==1?std::vector<int>{0}:std::vector<int>{0,1,2};
        auto target_cores=lanes==1?std::vector<int>{1}:cores;
        ling3::Linear owner(package,base,{k,n,1,cores,3}), target(package,base,{k,n,1,target_cores,3});
        std::vector<float> x(128*k),y(128*n),scales(128),row(k);
        std::vector<int8_t> q(x.size()),wq(k);
        for(size_t i=0;i<x.size();++i)x[i]=std::sin(float(i)*.137f)*1.3f;
        for(int r=0;r<128;++r)scales[r]=ling3::QuantizeSymmetricInt8(std::span<const float>(x).subspan(r*k,k),std::span<int8_t>(q).subspan(r*k,k)).scale;
        auto verify=[&](int rows,std::span<const size_t>indices){
            for(int col:{0,1,15,31,63,n/2,n-2,n-1}){
                float mx=0;
                for(int j=0;j<k;++j){
                    if(t.entry->dtype==unsigned(ling3::DataType::kBFloat16))row[j]=ling3::BFloat16ToFloat(reinterpret_cast<const uint16_t*>(t.data)[col*k+j]);
                    else{
                        const auto&s=package.tensor(base+".scales");
                        row[j]=ling3::DecodeInt4LowFirst({t.data,size_t(t.entry->data_bytes)},size_t(j)*n+col)*
                            ling3::BFloat16ToFloat(reinterpret_cast<const uint16_t*>(s.data)[col*(k/32)+j/32]);
                    }
                    mx=std::max(mx,std::abs(row[j]));
                }
                float ws=mx>0?mx/127:1;for(int j=0;j<k;++j)wq[j]=std::clamp(int(std::nearbyint(row[j]/ws)),-127,127);
                for(int r=0;r<rows;++r){int ir=indices.empty()?r:indices[r];int32_t v=0;for(int j=0;j<k;++j)v+=int32_t(q[ir*k+j])*wq[j];
                    float expected=float(v)*scales[ir]*ws;
                    if(!std::isfinite(y[r*n+col])||std::abs(y[r*n+col]-expected)>1e-5f*std::max(1.f,std::abs(expected)))throw std::runtime_error("bridge reference mismatch");}
            }
        };
        for(int rows:{1,3,16,128,3}){
            auto out=std::span<float>(y).first(rows*n);
            owner.RunBatchWithWeights(std::span<const float>(x).first(rows*k),rows,target,out);verify(rows,{});
            std::vector<size_t>indices(rows);for(int r=0;r<rows;++r)indices[r]=(r*17+5)%128;
            owner.RunBatchQuantizedRows(q,scales,indices,target,out);verify(rows,indices);
        }
        if(lanes==1){
        target.ShareInputFrom(owner);const float s=owner.PrepareInput(std::span<const float>(x).first(k));
        target.RunPrepared(s,std::span<float>(y).first(n));verify(1,{});
        target.SetSingleCore(2);target.Run(std::span<const float>(x).first(k),std::span<float>(y).first(n));verify(1,{});
        }
    }
    }
    std::cout<<"PASS: group32/raw weights, dynamic rows, rebound weights, indexed inputs, shared A and core rebinding\n";
}catch(const std::exception&e){std::cerr<<e.what()<<"\n";return 1;}