const std = @import("std");
const miara = @import("miara");
const symspell = miara.symspell;
const print = std.debug.print;

const SEED = 0x3C38E88DF7E69E67;
const REPEATS = 7;

const Config = struct {
    distance: usize,
    prefix: ?usize = null,
    compress_edits: bool = true,
    bitpack_refs: bool = true,
    exact_index: bool = true,
    fingerprint_bits: u8 = 4,
    report_text: bool = false,

    fn Index(comptime c: Config) type {
        const Func = struct {
            inline fn editDistanceBuffer(_: void, buffer: []u32, a: []const u8, b: []const u8, max: u32) ?u32 {
                return symspell.damerauLevenshteinDistanceLimited(u8, u32, buffer, a, b, max);
            }

            inline fn strLen(_: void, word: []const u8) ?usize {
                return word.len;
            }

            fn wordMaxDistance(_: void, _: []const u8) usize {
                return c.distance;
            }
        };

        const K = symspell.Keys([]const u8, void);
        const key = comptime if (c.prefix) |p| K.ends(p, 0) else K.whole;

        return symspell.GenericSymSpell(
            []const u8,
            void,
            Func.editDistanceBuffer,
            Func.strLen,
            Func.wordMaxDistance,
            key,
            c.compress_edits,
            c.bitpack_refs,
            c.exact_index,
            c.fingerprint_bits,
        );
    }

    fn name(comptime c: Config) []const u8 {
        const prefix = if (c.prefix) |p| std.fmt.comptimePrint(" prefix {d}", .{p}) else "";

        const variant = vr: {
            if (!c.compress_edits and !c.bitpack_refs) break :vr " raw";
            if (!c.compress_edits) break :vr " no ef";
            if (!c.exact_index and c.fingerprint_bits == 0) break :vr " smallest";
            if (!c.exact_index) break :vr " no exact";
            if (c.fingerprint_bits == 0) break :vr " no fp";
            if (c.fingerprint_bits != 4) break :vr std.fmt.comptimePrint(" fp {d}", .{c.fingerprint_bits});
            break :vr "";
        };

        return std.fmt.comptimePrint("miara d={d}{s}{s}", .{ c.distance, prefix, variant });
    }
};

const CONFIGS = [_]Config{
    .{ .distance = 1 },
    .{ .distance = 1, .prefix = 7 },
    .{ .distance = 1, .prefix = 7, .exact_index = false, .fingerprint_bits = 0 },
    .{ .distance = 2 },
    .{ .distance = 2, .prefix = 7, .report_text = true },
    .{ .distance = 2, .prefix = 7, .exact_index = false, .fingerprint_bits = 0 },
    .{ .distance = 2, .prefix = 7, .fingerprint_bits = 8 },
    .{ .distance = 2, .prefix = 7, .compress_edits = false },
    .{ .distance = 2, .compress_edits = false },
    .{ .distance = 2, .compress_edits = false, .bitpack_refs = false },
};

const Token = struct { word: []const u8, count: u32 };

fn loadDictionary(allocator: std.mem.Allocator, io: std.Io, path: []const u8) ![]Token {
    const content = try std.Io.Dir.cwd().readFileAlloc(io, path, allocator, .unlimited);
    var tokens: std.ArrayListUnmanaged(struct { word: []const u8, count: u64 }) = .empty;

    var count_max: u64 = 0;
    var lines = std.mem.tokenizeAny(u8, std.mem.trimStart(u8, content, "\xEF\xBB\xBF"), "\r\n");
    while (lines.next()) |line| {
        var fields = std.mem.tokenizeScalar(u8, line, ' ');
        const word = fields.next() orelse continue;
        const count = try std.fmt.parseInt(u64, fields.next() orelse continue, 10);
        count_max = @max(count_max, count);
        try tokens.append(allocator, .{ .word = word, .count = count });
    }

    var shift: u6 = 0;
    while ((count_max >> shift) > std.math.maxInt(u32)) shift += 1;

    const out = try allocator.alloc(Token, tokens.items.len);
    for (out, tokens.items) |*o, t| {
        o.* = .{ .word = t.word, .count = @intCast(@max(t.count >> shift, 1)) };
    }

    return out;
}

const Phrase = struct {
    input: []const u8,
    words: []const []const u8,
};

const Text = struct {
    error_percent: u32,
    words: []const []const u8,
    expect: []const []const u8,
};

const Queries = struct {
    hits: []const []const u8,
    typos: []const []const u8,
    misses: []const []const u8,
    phrases: []const Phrase,
    texts: []const Text,
};

fn readQueries(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !Queries {
    const content = try std.Io.Dir.cwd().readFileAlloc(io, path, allocator, .unlimited);
    const Section = enum { HITS, TYPOS, MISSES, PHRASES, TEXTS };

    var words: std.EnumArray(Section, std.ArrayListUnmanaged([]const u8)) = .initFill(.empty);
    var phrases: std.ArrayListUnmanaged(Phrase) = .empty;
    var texts: std.ArrayListUnmanaged(Text) = .empty;

    var section: Section = .HITS;
    var lines = std.mem.splitScalar(u8, content, '\n');
    while (lines.next()) |line| {
        if (line.len == 0 or line[0] == '#') continue;

        if (line[0] == '@') {
            section = std.meta.stringToEnum(Section, line[1..]) orelse return error.UnknownSection;
            continue;
        }

        var fields = std.mem.splitScalar(u8, line, '\t');
        const first = fields.next().?;
        switch (section) {
            .PHRASES => try phrases.append(allocator, .{
                .input = first,
                .words = try splitWords(allocator, fields.next() orelse return error.InvalidQueriesFile),
            }),
            .TEXTS => try texts.append(allocator, .{
                .error_percent = try std.fmt.parseInt(u32, first, 10),
                .words = try splitWords(allocator, fields.next() orelse return error.InvalidQueriesFile),
                .expect = try splitWords(allocator, fields.next() orelse return error.InvalidQueriesFile),
            }),
            else => try words.getPtr(section).append(allocator, first),
        }
    }

    return .{
        .hits = words.get(.HITS).items,
        .typos = words.get(.TYPOS).items,
        .misses = words.get(.MISSES).items,
        .phrases = phrases.items,
        .texts = texts.items,
    };
}

fn splitWords(allocator: std.mem.Allocator, line: []const u8) ![]const []const u8 {
    var out: std.ArrayListUnmanaged([]const u8) = .empty;
    var it = std.mem.tokenizeScalar(u8, line, ' ');
    while (it.next()) |word| try out.append(allocator, word);
    return out.toOwnedSlice(allocator);
}

fn elapsedNs(io: std.Io, start: std.Io.Timestamp) f64 {
    return @floatFromInt(start.durationTo(std.Io.Clock.awake.now(io)).toNanoseconds());
}

fn fastestNs(io: std.Io, ctx: anytype, comptime run: fn (@TypeOf(ctx)) anyerror!void) !f64 {
    var best = std.math.inf(f64);
    for (0..(REPEATS + 1)) |pass| {
        const start = std.Io.Clock.awake.now(io);
        try run(ctx);
        if (pass > 0) best = @min(best, elapsedNs(io, start));
    }
    return best;
}

fn percent(part: usize, whole: usize) f64 {
    return 100.0 * @as(f64, @floatFromInt(part)) / @as(f64, @floatFromInt(whole));
}

const TextResult = struct {
    error_percent: u32,
    us: f64,
    fixed: f64,
};

const Row = struct {
    name: []const u8,
    texts: []const TextResult = &.{},
    deletes: ?usize = null,
    build_s: ?f64 = null,
    bytes: ?usize = null,
    hit_ns: ?f64 = null,
    typo_ns: ?f64 = null,
    miss_ns: ?f64 = null,
    found: ?f64 = null,
    segment_us: ?f64 = null,
    exact: ?f64 = null,
    text5_us: ?f64 = null,
    text10_us: ?f64 = null,
    fixed: ?f64 = null,
};

const ROW = "{s:<28} {s:>9} {s:>8} {s:>9} | {s:>7} {s:>7} {s:>7} {s:>7} | {s:>10} {s:>6} | {s:>10} {s:>11} {s:>6}\n";

fn printRow(r: Row) void {
    var buf: [12][24]u8 = undefined;
    var hb = miara.util.HumanBytes{};
    print(ROW, .{
        r.name,
        cell(&buf[0], "{d}", r.deletes),
        cell(&buf[1], "{d:.2}", r.build_s),
        if (r.bytes) |b| hb.fmt(b) else "",
        cell(&buf[2], "{d:.0}", r.hit_ns),
        cell(&buf[3], "{d:.0}", r.typo_ns),
        cell(&buf[4], "{d:.0}", r.miss_ns),
        cell(&buf[5], "{d:.1}%", r.found),
        cell(&buf[6], "{d:.0}", r.segment_us),
        cell(&buf[7], "{d:.1}%", r.exact),
        cell(&buf[8], "{d:.0}", r.text5_us),
        cell(&buf[9], "{d:.0}", r.text10_us),
        cell(&buf[10], "{d:.1}%", r.fixed),
    });
}

fn cell(buf: []u8, comptime fmt: []const u8, value: anytype) []const u8 {
    const v = value orelse return "";
    return std.fmt.bufPrint(buf, fmt, .{v}) catch buf[0..0];
}

fn benchIndex(comptime SS: type, allocator: std.mem.Allocator, io: std.Io, name: []const u8, dict: []const SS.Token, queries: Queries) !Row {
    var row = Row{ .name = name };

    var seed: u64 = SEED;
    var ss = while (true) : (seed += 1) {
        const start = std.Io.Clock.awake.now(io);
        const ss = SS.init(allocator, seed, dict, {}, .{}) catch continue;
        row.build_s = elapsedNs(io, start) / 1e9;
        break ss;
    };
    defer ss.deinit(allocator);
    row.deletes = ss.pthash.size();

    var encoded = std.Io.Writer.Allocating.init(allocator);
    defer encoded.deinit();

    try ss.writeTo(&encoded.writer);
    row.bytes = encoded.written().len;

    var searcher = try SS.Searcher.init(&ss, allocator);
    defer searcher.deinit(allocator);

    const Lookup = struct {
        searcher: *SS.Searcher,
        set: []const []const u8,
        found: usize = 0,
        fn run(self: *@This()) !void {
            self.found = 0;
            for (self.set) |q| {
                try self.searcher.load(q, self.searcher.sym_spell.dict_stats.max_distance);
                const hit = try self.searcher.top();
                std.mem.doNotOptimizeAway(hit);
                if (hit != null) self.found += 1;
            }
        }
    };

    var hits = Lookup{ .searcher = &searcher, .set = queries.hits };
    row.hit_ns = try fastestNs(io, &hits, Lookup.run) / @as(f64, @floatFromInt(queries.hits.len));

    var typos = Lookup{ .searcher = &searcher, .set = queries.typos };
    row.typo_ns = try fastestNs(io, &typos, Lookup.run) / @as(f64, @floatFromInt(queries.typos.len));
    row.found = percent(typos.found, queries.typos.len);

    var misses = Lookup{ .searcher = &searcher, .set = queries.misses };
    row.miss_ns = try fastestNs(io, &misses, Lookup.run) / @as(f64, @floatFromInt(queries.misses.len));

    const Segment = struct {
        ss: *const SS,
        allocator: std.mem.Allocator,
        phrases: []const Phrase,
        exact: usize = 0,
        fn run(self: *@This()) !void {
            self.exact = 0;
            for (self.phrases) |phrase| {
                var s = try self.ss.wordSegmentation(self.allocator, phrase.input) orelse continue;
                defer s.deinit(self.allocator);
                if (s.parts.len != phrase.words.len) continue;
                for (s.parts, phrase.words) |part, word| {
                    if (!std.mem.eql(u8, part.corrected, word)) break;
                } else self.exact += 1;
            }
        }
    };

    var segment = Segment{ .ss = &ss, .allocator = allocator, .phrases = queries.phrases };
    row.segment_us = try fastestNs(io, &segment, Segment.run) / 1e3 / @as(f64, @floatFromInt(queries.phrases.len));
    row.exact = percent(segment.exact, queries.phrases.len);

    const Correct = struct {
        searcher: *SS.Searcher,
        text: Text,
        right: usize = 0,
        fn run(self: *@This()) !void {
            self.right = 0;
            for (self.text.words, self.text.expect) |word, expect| {
                try self.searcher.load(word, self.searcher.sym_spell.dict_stats.max_distance);
                const fixed = if (try self.searcher.top()) |hit| hit.word else word;
                if (std.mem.eql(u8, fixed, expect)) self.right += 1;
            }
        }
    };

    const texts = try allocator.alloc(TextResult, queries.texts.len);
    for (queries.texts, texts) |text, *result| {
        var correct = Correct{ .searcher = &searcher, .text = text };
        result.* = .{
            .error_percent = text.error_percent,
            .us = try fastestNs(io, &correct, Correct.run) / 1e3,
            .fixed = percent(correct.right, text.words.len),
        };
        if (text.error_percent == 5) row.text5_us = result.us;
        if (text.error_percent == 10) {
            row.text10_us = result.us;
            row.fixed = result.fixed;
        }
    }
    row.texts = texts;

    return row;
}

fn benchHashMap(allocator: std.mem.Allocator, io: std.Io, dict: []const Token, queries: Queries) !Row {
    var map = std.StringHashMap(u32).init(allocator);
    try map.ensureTotalCapacity(@intCast(dict.len));
    var word_bytes: usize = 0;
    for (dict) |t| {
        map.putAssumeCapacity(t.word, t.count);
        word_bytes += t.word.len;
    }

    const Probe = struct {
        map: *std.StringHashMap(u32),
        set: []const []const u8,
        fn run(self: *@This()) !void {
            for (self.set) |q| std.mem.doNotOptimizeAway(self.map.get(q));
        }
    };

    var times: [3]f64 = undefined;
    for ([_][]const []const u8{ queries.hits, queries.typos, queries.misses }, &times) |set, *time| {
        var probe = Probe{ .map = &map, .set = set };
        time.* = try fastestNs(io, &probe, Probe.run) / @as(f64, @floatFromInt(set.len));
    }

    return .{
        .name = "hash map, words only",
        .bytes = map.capacity() * (@sizeOf([]const u8) + @sizeOf(u32) + 1) + word_bytes,
        .hit_ns = times[0],
        .typo_ns = times[1],
        .miss_ns = times[2],
    };
}

fn comparePrefix(comptime prefix_length: usize, allocator: std.mem.Allocator, dict: []const Token, queries: Queries) !void {
    const Full = Config.Index(.{ .distance = 2 });
    const Prefixed = Config.Index(.{ .distance = 2, .prefix = prefix_length });

    var full = try Full.init(allocator, SEED, @ptrCast(dict), {}, .{});
    defer full.deinit(allocator);

    var prefixed = try Prefixed.init(allocator, SEED, @ptrCast(dict), {}, .{});
    defer prefixed.deinit(allocator);

    var s1 = try Full.Searcher.init(&full, allocator);
    defer s1.deinit(allocator);

    var s2 = try Prefixed.Searcher.init(&prefixed, allocator);
    defer s2.deinit(allocator);

    var total: usize = 0;
    var same: usize = 0;
    for ([_][]const []const u8{ queries.hits, queries.typos, queries.misses }) |set| {
        for (set) |q| {
            try s1.load(q, 2);
            try s2.load(q, 2);

            const a = try s1.top();
            const b = try s2.top();
            total += 1;
            if (a == null and b == null) {
                same += 1;
            } else if (a != null and b != null and std.mem.eql(u8, a.?.word, b.?.word) and a.?.edit_distance == b.?.edit_distance) {
                same += 1;
            }
        }
    }

    print("\nprefix {d} vs full at d=2: same word and distance on {d} of {d} queries\n", .{ prefix_length, same, total });
}

pub fn main(init: std.process.Init) !void {
    const io = init.io;

    var arena = std.heap.ArenaAllocator.init(init.gpa);
    defer arena.deinit();
    const allocator = arena.allocator();

    var args = try std.process.Args.Iterator.initAllocator(init.minimal.args, allocator);
    _ = args.next();

    const dict_path = args.next() orelse "";
    const queries_path = args.next() orelse "";
    if (dict_path.len == 0 or queries_path.len == 0) {
        print("usage: zig build bench-symspell -Doptimize=ReleaseFast -- <frequency dictionary> <queries file>\n", .{});
        return error.InvalidArgument;
    }

    const dict = try loadDictionary(allocator, io, dict_path);
    const queries = try readQueries(allocator, io, queries_path);

    print("{d} words, {d} queries per set, {d} phrases, text of {d} words\n\n", .{
        dict.len,
        queries.hits.len,
        queries.phrases.len,
        queries.texts[0].words.len,
    });

    print(ROW, .{
        "",
        "deletes",
        "build s",
        "size",
        "hit ns",
        "typo ns",
        "miss ns",
        "found",
        "segment us",
        "exact",
        "text 5% us",
        "text 10% us",
        "fixed",
    });

    var text_report: ?Row = null;
    inline for (CONFIGS) |c| {
        const row = try benchIndex(c.Index(), allocator, io, c.name(), @ptrCast(dict), queries);
        printRow(row);
        if (c.report_text) text_report = row;
    }
    printRow(try benchHashMap(allocator, io, dict, queries));

    if (text_report) |row| {
        print("\n{s}, {d} word document:", .{ row.name, queries.texts[0].words.len });
        for (row.texts, 0..) |t, i| {
            print("{s} {d}% errors {d:.0} us {d:.1}% right", .{ if (i == 0) "" else ",", t.error_percent, t.us, t.fixed });
        }
        print("\n", .{});
    }
    try comparePrefix(7, allocator, dict, queries);
}