// =====================================================
// SIMPLE Law analytical model
// =====================================================


// Fitted parameters (thetas.csv)
const THETA0 = -1.86626365e-03;
const THETA1 = 3.30843561e-11;
const THETA2 = 1.04066816e-09;
const THETA3 = 2.02013691e-07;

const THETAS = [THETA0, THETA1, THETA2, THETA3];

const COMPONENT_LABELS =
[
    "Prefill Amortization",
    "Decode Accum. KV-cache",
    "Decode Accum. Next-token"
];

const COMPONENT_COLORS = ["#55A868", "#E2C044", "#DD8452"];
const COMPONENT_HATCHES = ["x", "/", "-"];

// bottom -> top of the stack (indices into the component arrays)
const STACK_ORDER = [2, 1, 0];

const GRID = "#e1e0d9";

// matplotlib RdYlGn_r
const RDYLGN_R =
[
    "#006837", "#1a9850", "#66bd63", "#a6d96a", "#d9ef8b", "#ffffbf",
    "#fee08b", "#fdae61", "#f46d43", "#d73027", "#a50026"
];

// Validated models, from models.json (up to ~9B)
const PRESET_MODELS =
{
    "Llama 3.2 1B":
    {
        params:1.24e+09,
        d:2048,
        L:16,
        nq:32,
        nkv:8
    },

    "Llama 3.2 3B":
    {
        params:3.21e+09,
        d:3072,
        L:28,
        nq:24,
        nkv:8
    },

    "Llama 3.1 8B":
    {
        params:8.03e+09,
        d:4096,
        L:32,
        nq:32,
        nkv:8
    },

    "Gemma 2 2B":
    {
        params:2.61e+09,
        d:2304,
        L:26,
        nq:8,
        nkv:4
    },

    "Gemma 2 9B":
    {
        params:9.24e+09,
        d:3584,
        L:42,
        nq:16,
        nkv:8
    },

    "Granite 3B":
    {
        params:3.48e+09,
        d:2560,
        L:32,
        nq:32,
        nkv:32
    },

    "Granite-8B":
    {
        params:8.05e+09,
        d:4096,
        L:36,
        nq:32,
        nkv:8
    },

    "Qwen2 1.5B":
    {
        params:1.54e+09,
        d:1536,
        L:28,
        nq:12,
        nkv:2
    },

    "Qwen2 7B":
    {
        params:7.62e+09,
        d:3584,
        L:28,
        nq:28,
        nkv:4
    },

    "OPT 1.3B":
    {
        params:1.32e+09,
        d:2048,
        L:24,
        nq:32,
        nkv:32
    },

    "OPT 2.7B":
    {
        params:2.75e+09,
        d:2560,
        L:32,
        nq:32,
        nkv:32
    },

    "OPT 6.7B":
    {
        params:6.71e+09,
        d:4096,
        L:32,
        nq:32,
        nkv:32
    },

    "Falcon 7B":
    {
        params:7.78e+09,
        d:4096,
        L:36,
        nq:64,
        nkv:64
    }
};

// Size classes for the preset dropdown, by parameter count
const PRESET_GROUPS =
[
    {label:"XS · ~1B", max:2e9},
    {label:"S · ~2B–3B", max:5e9},
    {label:"M · ~6B–9B", max:Infinity}
];

window.onload = function()
{

    let selector = document.getElementById("preset");


    let names =
        Object.keys(PRESET_MODELS)
        .sort((a,b)=>PRESET_MODELS[a].params-PRESET_MODELS[b].params);

    let min=0;

    for(let group of PRESET_GROUPS)
    {
        let optgroup = document.createElement("optgroup");

        optgroup.label=group.label;

        for(let name of names.filter(n=>PRESET_MODELS[n].params>=min && PRESET_MODELS[n].params<group.max))
        {
            let option = document.createElement("option");

            option.text=name;

            option.value=name;

            optgroup.appendChild(option);
        }

        selector.appendChild(optgroup);

        min=group.max;
    }

    selector.value = "Llama 3.1 8B";

    renderThetas();

    renderFormula();

    fitFormula();
    document.fonts.ready.then(fitFormula);

    // Refit only when the width changes (fitting itself changes the height)
    let formulaWidth=0;
    new ResizeObserver(entries=>{
        const width=entries[0].contentRect.width;
        if(width!==formulaWidth)
        {
            formulaWidth=width;
            fitFormula();
        }
    }).observe(document.querySelector(".formula"));

    loadPreset();

}

// kind: "ok" | "loading" | "error" | "custom"
function setStatus(kind,html)
{

    let status = document.getElementById("status");

    status.className = "status "+kind;

    status.innerHTML = html;

    // The plots keep showing the previous model while loading or on error
    if(kind==="ok" || kind==="custom")
        document.querySelectorAll(".plotModel").forEach(
            el=>el.innerHTML = html
        );

}

// n_kv / n_q reduced to lowest terms, e.g. 8/32 -> "1/4"
function ratioFraction(nkv,nq)
{

    let gcd=(a,b)=>b ? gcd(b,a%b) : a;

    let g=gcd(nkv,nq) || 1;

    let num=nkv/g, den=nq/g;

    return den===1 ? String(num) : `${num}/${den}`;

}

function attentionLabel(nq,nkv)
{

    if(nq===nkv) return "MHA";

    if(nkv===1) return "MQA";

    return `GQA (${+(nq/nkv).toFixed(1)}x)`;

}

// Derived quantities shown under the architecture fields
function updateSummary(d,L,nq,nkv)
{

    document.getElementById("archSummary").innerHTML =
    `
    <span class="chip">${attentionLabel(nq,nkv)}</span>
    <span class="chip">R = ${ratioFraction(nkv,nq)}</span>
    `;

}

function loadPreset()
{

    let name = document.getElementById("preset").value;


    let model = PRESET_MODELS[name];

    if (!model)
    {
        console.error("Unknown preset:", name);
        return;
    }

    document.getElementById("d").value=model.d;
    document.getElementById("L").value=model.L;
    document.getElementById("nq").value=model.nq;
    document.getElementById("nkv").value=model.nkv;


    setStatus("ok", `<b>${name}</b> · validated`);


    draw();

}

// Returns the three additive terms, as in the notebook's compute_components
function computeComponents(d,L,nq,nkv,nin,nout)
{

    const ratio = nkv/nq;

    let prefill=[];
    let kvCache=[];
    let nextTok=[];


    for(let n of nout)
    {

        prefill.push(
            THETA1*d*d*L*nin/n
        );

        kvCache.push(
            THETA2*d*L*ratio*(nin + n/2)
        );

        nextTok.push(
            THETA0 + THETA3*d*L
        );
    }


    return [
        prefill,kvCache,nextTok
    ];
}

function energyPerToken(d,L,nq,nkv,nin,nout)
{
    return THETA1*d*d*L*nin/nout
        + THETA2*d*L*(nkv/nq)*(nin + nout/2)
        + THETA0 + THETA3*d*L;
}



function logspace2(start,end,points)
{
    let arr=[];

    let a=Math.log2(start);
    let b=Math.log2(end);

    for(let i=0;i<points;i++)
    {
        let x=a+(b-a)*i/(points-1);
        arr.push(2**x);
    }

    return arr;
}



function draw()
{

    let d = Number(document.getElementById("d").value);

    let L = Number(document.getElementById("L").value);

    let nq = Number(document.getElementById("nq").value);

    let nkv = Number(document.getElementById("nkv").value);

    let nin = Number(document.getElementById("nin").value);


    updateSummary(d,L,nq,nkv);

    drawComponents(d,L,nq,nkv,nin);

    drawContour(d,L,nq,nkv,nin);

}



function drawComponents(d,L,nq,nkv,nin)
{

    let nout = logspace2(1,16384,400);


    let components =
        computeComponents(
            d,L,nq,nkv,nin,nout
        );


    let total =
        nout.map(
            (_,i)=>
            components.reduce(
                (a,b)=>a+b[i],
                0
            )
        );


    let traces=[];


    for(let i of STACK_ORDER)
    {

        traces.push({
            x:nout,
            y:components[i],
            mode:"lines",
            name:COMPONENT_LABELS[i],
            legendrank:i+1,
            stackgroup:"one",
            line:{
                color:"rgba(0,0,0,0.2)",
                width:0.5
            },
            fillcolor:COMPONENT_COLORS[i],
            fillpattern:{
                shape:COMPONENT_HATCHES[i],
                bgcolor:COMPONENT_COLORS[i],
                fgcolor:"rgba(0,0,0,0.6)",
                size:14,
                solidity:0.05
            },
            hovertemplate:"%{y:.3g} J/tok<extra>"+COMPONENT_LABELS[i]+"</extra>"
        });

    }

    traces.push({
        x:nout,
        y:total,
        mode:"lines",
        name:"Total E<sub>tok</sub>",
        legendrank:4,
        line:{
            color:"black",
            width:2.5
        },
        hovertemplate:"%{y:.3g} J/tok<extra>Total</extra>"

    });


    // Every other power of two, so the ticks fit the narrower panel
    let ticks=[];
    for(let e=0;e<=14;e+=2) ticks.push(2**e);


    let layout={
        height:450,
        font:{
            family:"sans-serif",
            size:13
        },
        xaxis:{
            title:
            `Output Tokens τ<sub>out</sub> (τ<sub>in</sub>=${nin}) [log<sub>2</sub>]`,
            type:"log",
            range:[0,Math.log10(16384)],
            tickmode:"array",
            tickvals:ticks,
            ticktext:ticks.map(String),
            tickangle:-45,
            gridcolor:GRID,
            showline:true,
            linecolor:"black",
            ticks:"outside"
        },


        yaxis:{
            title:
            "Energy-Per-Token E<sub>tok</sub> [J/tok]",
            rangemode:"tozero",
            gridcolor:GRID,
            showline:true,
            linecolor:"black",
            ticks:"outside",
            zeroline:false
        },

        hovermode:"x unified",

        legend:{
            x:0.98,
            xanchor:"right",
            y:1,
            traceorder:"normal",
            bgcolor:"rgba(255,255,255,0)"
        },

        plot_bgcolor:"white",

        margin:{
            l:70,
            r:20,
            t:20,
            b:80
        }

    };


    Plotly.newPlot(
        "plot",
        traces,
        layout,
        {
            responsive:true
        }
    );

}



// Same grid and level scheme as plot_E_tok_contour in the notebook.
// Axes are plotted in log2 units and E in log10, so evenly spaced
// Plotly contour levels match np.geomspace in the notebook.
const CONTOUR_NIN_RANGE = [16, 65536];
const CONTOUR_NOUT_RANGE = [4, 16384];
const CONTOUR_POINTS = 200;
const CONTOUR_LEVELS = 24;

function drawContour(d,L,nq,nkv,nin)
{

    let xs = logspace2(...CONTOUR_NIN_RANGE, CONTOUR_POINTS).map(Math.log2);

    let ys = logspace2(...CONTOUR_NOUT_RANGE, CONTOUR_POINTS).map(Math.log2);


    let z=[];
    let text=[];

    for(let y of ys)
    {
        let zrow=[];
        let trow=[];

        for(let x of xs)
        {
            let e = energyPerToken(d,L,nq,nkv,2**x,2**y);

            zrow.push(Math.log10(e));

            trow.push(
                `τ<sub>in</sub> = ${Math.round(2**x)}<br>`
                +`τ<sub>out</sub> = ${Math.round(2**y)}<br>`
                +`E<sub>tok</sub> = ${e.toPrecision(3)} J/tok`
            );
        }

        z.push(zrow);
        text.push(trow);
    }


    let flat = z.flat();
    let zmin = Math.min(...flat);
    let zmax = Math.max(...flat);
    let step = (zmax-zmin)/(CONTOUR_LEVELS-1);

    let levels=[];
    for(let i=0;i<CONTOUR_LEVELS;i++) levels.push(zmin+i*step);


    let colorscale =
        RDYLGN_R.map((c,i)=>[i/(RDYLGN_R.length-1),c]);


    let fill = {
        type:"contour",
        x:xs,
        y:ys,
        z:z,
        text:text,
        hovertemplate:"%{text}<extra></extra>",
        colorscale:colorscale,
        zmin:zmin,
        zmax:zmax,
        autocontour:false,
        contours:{
            start:zmin,
            end:zmax,
            size:step,
            coloring:"fill"
        },
        line:{
            width:0
        },
        colorbar:{
            title:{
                text:"Energy-Per-Token E<sub>tok</sub> [J/tok]",
                side:"right"
            },
            tickvals:levels.filter((_,i)=>i%3===0),
            ticktext:levels.filter((_,i)=>i%3===0).map(v=>formatEnergy(10**v)),
            outlinewidth:0.8
        }
    };


    let lines = {
        type:"contour",
        x:xs,
        y:ys,
        z:z,
        hoverinfo:"skip",
        showscale:false,
        autocontour:false,
        contours:{
            start:zmin,
            end:zmax,
            size:2*step,
            coloring:"none"
        },
        line:{
            color:"black",
            width:0.5
        }
    };


    let xticks=[];
    for(let e=Math.log2(CONTOUR_NIN_RANGE[0]);e<=Math.log2(CONTOUR_NIN_RANGE[1]);e++) xticks.push(e);

    let yticks=[];
    for(let e=Math.log2(CONTOUR_NOUT_RANGE[0]);e<=Math.log2(CONTOUR_NOUT_RANGE[1]);e++) yticks.push(e);


    let annotations =
        contourLabels(xs,ys,z,levels.filter((_,i)=>i%2===0),colorscale,zmin,zmax);


    // Mark the input-token count currently selected with the slider
    let shapes=[];
    let ninLog = Math.log2(nin);

    if(ninLog>=xs[0] && ninLog<=xs[xs.length-1])
    {
        shapes.push({
            type:"line",
            x0:ninLog,
            x1:ninLog,
            y0:ys[0],
            y1:ys[ys.length-1],
            line:{
                color:"black",
                width:1.5,
                dash:"dash"
            }
        });
    }


    let layout={
        height:450,
        font:{
            family:"sans-serif",
            size:13
        },
        xaxis:{
            title:"Input Tokens τ<sub>in</sub> [log<sub>2</sub>]",
            range:[xs[0],xs[xs.length-1]],
            tickmode:"array",
            tickvals:xticks,
            ticktext:xticks.map(e=>String(2**e)),
            tickangle:-45,
            showline:true,
            linecolor:"black",
            ticks:"outside",
            showgrid:false,
            zeroline:false
        },
        yaxis:{
            title:"Output Tokens τ<sub>out</sub> [log<sub>2</sub>]",
            range:[ys[0],ys[ys.length-1]],
            tickmode:"array",
            tickvals:yticks,
            ticktext:yticks.map(e=>String(2**e)),
            showline:true,
            linecolor:"black",
            ticks:"outside",
            showgrid:false,
            zeroline:false
        },
        annotations:annotations,
        shapes:shapes,
        margin:{
            l:80,
            r:20,
            t:20,
            b:90
        }
    };


    Plotly.newPlot(
        "contour",
        [fill,lines],
        layout,
        {
            responsive:true
        }
    );

}



// Place one label per contour line where it is farthest from every plot
// edge, skipping lines that never get far enough (as in the notebook).
function contourLabels(xs,ys,z,lineLevels,colorscale,zmin,zmax)
{

    const MARGIN_THRESHOLD = 0.09;

    let xmin=xs[0], xmax=xs[xs.length-1];
    let ymin=ys[0], ymax=ys[ys.length-1];

    let annotations=[];


    for(let level of lineLevels)
    {

        let best=null;
        let bestMargin=-1;

        let consider=(x,y)=>
        {
            let nx=(x-xmin)/(xmax-xmin);
            let ny=(y-ymin)/(ymax-ymin);
            let m=Math.min(nx,1-nx,ny,1-ny);

            if(m>bestMargin)
            {
                bestMargin=m;
                best=[x,y];
            }
        };


        for(let i=0;i<ys.length;i++)
        {
            for(let j=0;j<xs.length;j++)
            {
                let a=z[i][j]-level;

                if(j+1<xs.length)
                {
                    let b=z[i][j+1]-level;

                    if(a*b<0)
                        consider(xs[j]+(xs[j+1]-xs[j])*a/(a-b), ys[i]);
                }

                if(i+1<ys.length)
                {
                    let b=z[i+1][j]-level;

                    if(a*b<0)
                        consider(xs[j], ys[i]+(ys[i+1]-ys[i])*a/(a-b));
                }
            }
        }


        if(best && bestMargin>MARGIN_THRESHOLD)
        {
            annotations.push({
                x:best[0],
                y:best[1],
                text:formatEnergy(10**level),
                showarrow:false,
                font:{
                    size:11,
                    color:"black"
                },
                bgcolor:colorAt(colorscale,(level-zmin)/(zmax-zmin)),
                borderpad:1
            });
        }

    }


    return annotations;

}



function colorAt(colorscale,t)
{

    t=Math.min(1,Math.max(0,t));

    for(let i=1;i<colorscale.length;i++)
    {
        let [t1,c1]=colorscale[i];

        if(t<=t1)
        {
            let [t0,c0]=colorscale[i-1];
            let f=(t-t0)/(t1-t0);

            let rgb0=hexToRgb(c0);
            let rgb1=hexToRgb(c1);

            let mix=rgb0.map((v,k)=>Math.round(v+(rgb1[k]-v)*f));

            return `rgb(${mix.join(",")})`;
        }
    }

    return colorscale[colorscale.length-1][1];

}

function hexToRgb(hex)
{
    return [1,3,5].map(i=>parseInt(hex.slice(i,i+2),16));
}

function formatEnergy(v)
{
    return String(Number(v.toPrecision(2)));
}



// Shrinks the formula so it always fits on one line.
// All its inner sizes are in em, so the width scales linearly with font-size.
function fitFormula()
{

    const formula=document.querySelector(".formula");

    // Back to the CSS size, then scale down only if it overflows
    formula.style.fontSize="";

    const available=formula.clientWidth;
    const needed=formula.scrollWidth;

    if(needed>available)
    {
        const base=parseFloat(getComputedStyle(formula).fontSize);
        formula.style.fontSize=(base*available/needed*0.98)+"px";
    }

}

// Renders every element carrying a data-tex attribute
function renderFormula()
{

    document.querySelectorAll("[data-tex]").forEach(
        el=>katex.render(
            el.dataset.tex,
            el,
            {
                throwOnError:false
            }
        )
    );

}

// Component each theta belongs to (indices into COMPONENT_LABELS)
const THETA_COMPONENT = [2, 0, 1, 2];

const COMPONENT_CLASSES = ["prefill", "kvcache", "nexttok"];

function renderThetas()
{

    let body = document.getElementById("thetaBody");

    body.innerHTML = THETAS.map(
        (t,i)=>
        {
            let [mantissa,exponent] = t.toExponential(4).split("e");
            let c = THETA_COMPONENT[i];

            return `
            <tr class="${COMPONENT_CLASSES[c]}">
                <td><span class="dot"></span><span data-tex="\\theta_${i}"></span></td>
                <td data-tex="${mantissa}\\times10^{${Number(exponent)}}"></td>
                <td class="termName">${COMPONENT_LABELS[c]}</td>
            </tr>`;
        }
    ).join("");

}

// Reads d, L, n_q, n_kv from a config.json.
// Multimodal models (e.g. Gemma 3, Qwen-VL) nest the language model
// settings under text_config, so that is used when present.
function readArchitecture(config)
{

    let cfg = config.text_config || config;

    let d =
        cfg.hidden_size ||
        cfg.d_model ||
        cfg.n_embd;


    let L =
        cfg.num_hidden_layers ||
        cfg.num_layers ||
        cfg.n_layer;


    let nq =
        cfg.num_attention_heads ||
        cfg.n_head;


    let nkv =
        cfg.num_key_value_heads ||
        nq;


    return {cfg,d,L,nq,nkv};

}

function validateModel(config)
{

    let {cfg,d,L} = readArchitecture(config);

    // ------------------------------
    // Reject Mixture-of-Experts models
    // ------------------------------

    const moeKeys =
    [
        "num_experts",
        "num_local_experts",
        "num_experts_per_tok",
        "num_sparse_layers",
        "router_aux_loss_coef",
        "moe_layer_freq"
    ];


    for (let key of moeKeys)
    {
        if (cfg[key] !== undefined)
        {
            throw new Error(
                "MoE models are not supported. "
                +"Please select a dense Transformer model."
            );
        }
    }



    // ------------------------------
    // Reject hybrid / linear-attention models
    // (e.g. Qwen3.8: only 1 layer in 4 keeps a growing KV-cache)
    // ------------------------------

    const attentionTypes =
    [
        "full_attention",
        "sliding_attention"
    ];

    let layerTypes = cfg.layer_types || [];

    if(layerTypes.some(t=>!attentionTypes.includes(t)))
    {
        throw new Error(
            "Hybrid / linear-attention models are not supported. "
            +"Please select a standard Transformer model."
        );
    }



    // ------------------------------
    // Estimate parameter count
    // ------------------------------

    if(!d || !L)
    {
        throw new Error(
            "Cannot estimate model size from config.json"
        );
    }

    // Dense Transformer approximation
    let params = 12 * L * d * d;


    let paramsB = params / 1e9;

    // ------------------------------
    // Size restriction
    // ------------------------------

    const MAX_PARAMS_B = 20;


    if(paramsB > MAX_PARAMS_B)
    {
        throw new Error(
            `Model too large (~${paramsB.toFixed(1)}B parameters). `
            +
            `Only dense models below ${MAX_PARAMS_B}B are allowed.`
        );
    }


    return paramsB;

}

async function loadModel()
{

    let repo =
        document.getElementById("repo").value.trim();


    await loadRepo(repo);

}

async function loadRepo(repo)
{

    if(!repo)
    {
        setStatus("error", "Enter a repository id, e.g. deepseek-ai/DeepSeek-R1-Distill-Llama-8B");
        return;
    }

    let url =
    `https://huggingface.co/${repo}/resolve/main/config.json`;

    setStatus("loading", `Loading <b>${repo}</b>…`);


    try
    {

        let response = await fetch(url);

        if(!response.ok)
            throw Error(
                "Could not load config.json. Some models (e.g. meta-llama models) require licenses."
            );


        let cfg = await response.json();

        validateModel(cfg);

        let {d,L,nq,nkv} = readArchitecture(cfg);


        document.getElementById("d").value=d;
        document.getElementById("L").value=L;
        document.getElementById("nq").value=nq;
        document.getElementById("nkv").value=nkv;



        setStatus("ok", `<b>${repo}</b> · Hugging Face`);


        draw();


    }
    catch(e)
    {

        setStatus("error", e.message);

    }

}

// redraw when the architecture fields change

document
.querySelectorAll("input[type=number]")
.forEach(
    x=>x.addEventListener("change",()=>
    {
        setStatus("custom", "Custom architecture");
        draw();
    })
);

document
.getElementById("repo")
.addEventListener("keydown",e=>
{
    if(e.key==="Enter") loadModel();
});


draw();