#pragma once
#include <charconv>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <string_view>

namespace ling3 {
inline std::set<std::string> ParseW8Families(const std::string& spec) {
    const std::set<std::string> valid{"attention","shared","dense","head","kda_in","kda_out","mla_in","mla_out"};
    std::set<std::string> result;
    if(spec.empty())return result;
    if(spec.back()==',')throw std::invalid_argument("empty W8 precision selector");
    std::istringstream in(spec);std::string item;
    while(std::getline(in,item,',')){
        if(!valid.contains(item))throw std::invalid_argument("unknown LING3_W8_FAMILIES item: "+item);
        result.insert(item);
    }
    return result;
}
inline std::string W8MatrixFamily(const std::string& base) {
    if(base=="lm_head")return "head";
    if(base.find(".attention.")!=std::string::npos){
        constexpr std::string_view prefix="model.layers.";
        if(!base.starts_with(prefix))throw std::invalid_argument("invalid attention matrix name");
        int layer=-1;auto [end,ec]=std::from_chars(base.data()+prefix.size(),base.data()+base.size(),layer);
#if LING3_EXPERIMENTAL_MTP
        constexpr int maximum_layer=24;
#else
        constexpr int maximum_layer=23;
#endif
        if(ec!=std::errc{}||end==base.data()+base.size()||*end!='.'||layer<0||layer>maximum_layer)
            throw std::invalid_argument("invalid attention layer");
#if LING3_EXPERIMENTAL_MTP
        // Optional layer 24 is MTP MLA, outside the main 24-layer trunk.
        return std::string((layer==24||(layer+1)%4==0)?"mla_":"kda_")+(base.ends_with(".o_proj")?"out":"in");
#else
        return std::string((layer+1)%4==0?"mla_":"kda_")+(base.ends_with(".o_proj")?"out":"in");
#endif
    }
    if(base.find(".mlp.shared_experts.")!=std::string::npos)return "shared";
    if(base=="model.layers.0.mlp.gate_up"||base=="model.layers.0.mlp.down_proj")return "dense";
    return ""; // Routed experts are never selected.
}
inline bool SelectW8(const std::set<std::string>& families,const std::string& base){
    const auto family=W8MatrixFamily(base);
    return !family.empty()&&(families.contains(family)||
        (families.contains("attention")&&(family.starts_with("kda_")||family.starts_with("mla_"))));
}
}