-- parstat.opm
-- Copyright 2025-2026 Petr Krajník
--
-- This work may be distributed and/or modified under the
-- conditions of the LaTeX Project Public License, either version 1.3c
-- of this license or (at your option) any later version.
-- The latest version of this license is in
--   https://www.latex-project.org/lppl.txt
-- and version 1.3c or later is part of all distributions of LaTeX
-- version 2008 or later.
--
-- This work has the LPPL maintenance status 'maintained'.
--
-- The Current Maintainer of this work is Petr Krajník.
--
-- This work consists of the files parstat.opm and parstat.lua.
--

local PKG_PREFIX = "_parstat_"

-- Register names
local REG_ENABLED = PKG_PREFIX .. "enabled"
local REG_VERBOSE = PKG_PREFIX .. "verbose"
local REG_FSKIP   = PKG_PREFIX .. "fskip"
local REG_BSKIP   = PKG_PREFIX .. "bskip"

-- Node types
local NODE_HLIST = node.id("hlist")
local NODE_VLIST = node.id("vlist")
local NODE_GLYPH = node.id("glyph")
local NODE_GLUE  = node.id("glue")

-- Node subtypes
local GLYPH_LIGATURE_FLAG = 2
local GLYPH_GHOST_FLAG = 4
local GLUE_SPACE = 13
local GLUE_XSPACE = 14

-- Line statistic outputs per console line
local STATS_PER_LINE = 5
assert(STATS_PER_LINE > 0, "STATS_PER_LINE must be greater than zero")

-- Line skips (only valid when processing)
local front_skip, back_skip

local function init_line_skips()
    front_skip = math.max(0, tex.count[REG_FSKIP])
    back_skip  = math.max(0, tex.count[REG_BSKIP])
end

-- Statistic structure
local function stat_new(name)
    return {
        name = name,
        samples = 0, -- Equals to line count
        average = 0,
        M2 = 0, -- Sum of squared diffs for stddev
        min = math.huge,
        max = 0 }
end

local function make_stats()
    return stat_new("Glyphs"), stat_new("Spaces")
end

local function stat_add(stat, value)
    stat.samples = stat.samples + 1

    -- Average and stddev (Welford’s Algorithm)
    local delta_old = value - stat.average
    stat.average = stat.average + delta_old / stat.samples
    local delta_new = value - stat.average
    stat.M2 = stat.M2 + delta_old * delta_new

    stat.min = math.min(stat.min, value)
    stat.max = math.max(stat.max, value)
end

local function stat_stddev(stat)
    if stat.samples > 1 then
        return math.sqrt(stat.M2 / (stat.samples - 1))
    end
    return 0.0
end

local function stat_min(stat)
    if stat.min == math.huge then
        return 0
    end
    return stat.min
end

-- Document statistics
local function doc_stat_new()
    local g, s = make_stats()
    return {
        glyphs = g,
        spaces = s,
        par_count = 0
    }
end

local doc_stat = doc_stat_new()

-- Logging
local function log_write_nl(msg)
    texio.write_nl(msg)
end

local function log_writef_nl(fmt, ...)
    assert(select('#', ...) > 0)
    texio.write_nl(string.format(fmt, ...))
end

local function log_writef(fmt, ...)
    assert(select('#', ...) > 0)
    texio.write(string.format(fmt, ...))
end

local function log_header(what)
    log_writef_nl("Parstat %s from %s:%d",
        tostring(what), status.filename, status.linenumber)
end

local function log_stat(stat)
    log_writef_nl("  %s: ave %.2f, stddev %.2f, min %d, max %d",
        stat.name, stat.average, stat_stddev(stat), stat_min(stat), stat.max)
end

local function log_line_stat(line, line_ignored, glyph_count, space_count)
    if (line - 1) % STATS_PER_LINE == 0 then
        log_write_nl(" ") -- Limit line stats per line
    end
    log_writef(" %d%s:(%dg,%ds)",
        line, line_ignored and "X" or "", glyph_count, space_count)
end

local function log_newline()
    texio.write_nl("")
end

--------------------------------

local function enough_lines(line_count)
    return (front_skip + back_skip) < line_count
end

local function is_line_ignored(line, line_count)
    return line <= front_skip or line > (line_count - back_skip)
end

local function is_real_glyph(n)
    return n.id == NODE_GLYPH and (n.subtype & GLYPH_GHOST_FLAG) == 0
end

local function is_ligature(n)
    return (n.subtype & GLYPH_LIGATURE_FLAG) ~= 0
end

local function is_space(n)
    return n.id == NODE_GLUE and (n.subtype == GLUE_SPACE
        or n.subtype == GLUE_XSPACE)
end

local function is_nested_node(n)
    return n.id == NODE_HLIST or n.id == NODE_VLIST
end

local function count_glyphs(head)
    local glyph_count = 0
    local space_count = 0

    for n in node.traverse(head) do
        if is_real_glyph(n) then
            if is_ligature(n) and n.components then
                -- Count Ligature components (see limitations)
                local gc, sc =  count_glyphs(n.components)
                glyph_count = glyph_count + gc
                space_count = space_count + sc
            else
                glyph_count = glyph_count + 1
            end
        elseif is_space(n) then
            glyph_count = glyph_count + 1
            space_count = space_count + 1
        elseif is_nested_node(n) then
            local gc, sc =  count_glyphs(n.head)
            glyph_count = glyph_count + gc
            space_count = space_count + sc
        end
    end
    return glyph_count, space_count
end

local function analyze_par(head, line_count, verbose)
    local glyph_stat, space_stat = make_stats()
    local curr_line = 0

    for line in node.traverse_id(NODE_HLIST, head) do
        curr_line = curr_line + 1
        local line_ignored = is_line_ignored(curr_line, line_count)

        -- Single line stat
        local glyph_count, space_count = count_glyphs(line.head)
        if verbose >= 2 then
            log_line_stat(curr_line, line_ignored, glyph_count, space_count)
        end

        if not line_ignored then
            -- Add the line results to the stats
            stat_add(glyph_stat, glyph_count)
            stat_add(space_stat, space_count)
            stat_add(doc_stat.glyphs, glyph_count)
            stat_add(doc_stat.spaces, space_count)
        end
    end
    assert(curr_line == line_count)

    if verbose > 0 then
        log_stat(glyph_stat)
        log_stat(space_stat)
        log_newline()
    end
end

--------------------------------

local parstat = _ENV.parstat or {}
_ENV.parstat = parstat

function parstat.run(head)
    if tex.count[REG_ENABLED] > 0 then
        init_line_skips()
        local line_count = node.count(NODE_HLIST, head)
        local verbose = tex.count[REG_VERBOSE]

        if enough_lines(line_count) then
            doc_stat.par_count = doc_stat.par_count + 1
            if verbose > 0 then
                log_header(doc_stat.par_count)
            end
            analyze_par(head, line_count, verbose)
        end
    end
    return head
end

function parstat.summary()
    log_header("summary")
    log_writef_nl("  Analyzed %d paragraphs, %d lines in total",
        doc_stat.par_count, doc_stat.glyphs.samples)
    log_stat(doc_stat.glyphs)
    log_stat(doc_stat.spaces)
    log_newline()
end

function parstat.reset()
    log_header("reset")
    log_newline()

    doc_stat = doc_stat_new()
end