const std = @import("std");

/// A data type to store connections of nodes in directional-graph.
pub fn Digraph(comptime T: type, comptime lessThanFn: LessThanFunc(T)) type {
    return struct {
        ///
        const Node = T;

        /// A connection between 2 nodes.
        /// e.g.) `Conn { .from = X, .to = Y }` =>  "X is connected to Y"
        const Conn = struct { from: T, to: T, };

        ///
        const ConnList = std.ArrayList(Conn);

        ///
        const Error = error {
            AlreadyConnected,
            NotConnected,
        };

        ///
        map: ConnList,

        ///
        pub fn init(alloc: std.mem.Allocator, mapUnsorted: []const Conn) !@This() {
            var mapSorted = ConnList.init(alloc);
            try mapSorted.ensureTotalCapacity(mapUnsorted.len);
            for (mapUnsorted) |conn| {
                try mapSorted.append(conn);
            }
            std.mem.sort(Conn, mapSorted.items, {}, compareConn);

            return .{
                .map  = mapSorted,
            };
        }
        ///
        pub fn deinit(self: *@This()) void {
            self.map.deinit();
        }

        /// Makes new connection between `from` to `to`.
        /// Returns false if they are already connected otherwise true.
        pub fn connectIf(self: *@This(), from: T, to: T) !bool {
            const begin, const end = self.findSegment(from);
            if (self.findConnectionInSegment(from, to, begin, end)) |_| {
                return false;
            } else {
                try self.map.insert(end, Conn { .from = from, .to = to, });
                return true;
            }
        }
        /// Same to `connectIf`, but returns an error if it's already connected.
        pub fn connect(self: *@This(), from: T, to: T) !void {
            if (!try self.connectIf(from, to)) {
                return Error.AlreadyConnected;
            }
        }

        /// Removes an existing connection between `from` to `to`.
        /// Returns false if they are not connected otherwise true.
        pub fn disconnectIf(self: *@This(), from: T, to: T) bool {
            const begin, const end = self.findSegment(from);
            if (self.findConnectionInSegment(from, to, begin, end)) |idx| {
                _ = self.map.orderedRemove(idx);
                return true;
            } else {
                return false;
            }
        }
        /// Same to `disconnectIf`, but returns an error if it's not connected.
        pub fn disconnect(self: *@This(), from: T, to: T) !void {
            if (!self.disconnectIf(from, to)) {
                return Error.NotConnected;
            }
        }

        ///
        pub fn isConnected(self: *const @This(), from: T, to: T) bool {
            const begin, const end = self.findSegment(from);
            return self.findConnectionInSegment(from, to, begin, end) != null;

        }

        ///
        pub fn getChildrenOf(self: *const @This(), from: T) []const Conn {
            const begin, const end = self.findSegment(from);
            return self.map.items[begin..end];
        }

        fn findConnectionInSegment(self: *const @This(), from: T, to: T, begin: usize, end: usize) ?usize {
            for (self.map.items[begin..end], begin..) |v, idx| {
                if ((v.from == from) and (v.to == to)) {
                    return idx;
                }
            }
            return null;
        }
        fn findSegment(self: *const @This(), from: T) struct { usize, usize } {
            const n = self.map.items.len;
            if (n == 0) {
                return .{ 0, 0, };
            }

            const baseIdx  = self.binsearch(from).?;
            const baseFrom = self.map.items[baseIdx].from;

            var begin: usize = undefined;
            var end  : usize = undefined;
            if (baseFrom < from) {
                begin = baseIdx;
                while ((begin < n) and (self.map.items[begin].from < from)) { begin += 1; }

                end = begin;
                while ((end < n) and (self.map.items[end].from == from)) { end += 1; }

            } else if (baseFrom > from) {
                end = baseIdx;
                while ((end > 0) and (self.map.items[end-1].from > from)) { end -= 1; }

                begin = end;
                while ((begin > 0) and (self.map.items[begin-1].from == from)) { begin -= 1; }

            } else {
                begin = baseIdx;
                while ((begin > 0) and (self.map.items[begin-1].from == from)) { begin -= 1; }

                end = baseIdx;
                while ((end < n) and (self.map.items[end].from == from)) { end += 1; }
            }
            return .{ begin, end, };
        }
        fn binsearch(self: *const @This(), from: T) ?usize {
            if (self.map.items.len == 0) {
                return null;
            }

            var left : usize = 0;
            var right: usize = self.map.items.len;

            var idx: usize = undefined;
            while (left < right) {
                idx = (left + right) / 2;

                const target = self.map.items[idx].from;
                if (target < from) {
                    left = idx + 1;
                } else if (target > from) {
                    right = idx -| 1;
                } else {
                    break;
                }
            }
            return idx;
        }
        fn compareConn(_: void, a: Conn, b: Conn) bool {
            return lessThanFn(a.from, b.from);
        }
    };
}

/// A type of comparator function for the type T, which is to be passed as an argument of `Digraph()`.
pub fn LessThanFunc(comptime T: type) type {
    return fn (lhs: T, rhs: T) bool;
}

/// Returns a lessThanFunc for the comparable type T.
pub fn lessThanFuncFor(comptime T: type) LessThanFunc(T) {
    return struct {
        fn inner(lhs: T, rhs: T) bool {
            if (@typeInfo(T) == .pointer) {
                return @intFromPtr(lhs) < @intFromPtr(rhs);
            } else {
                return lhs < rhs;
            }
        }
    }.inner;
}

test "compile check for various types" {
    _ = Digraph(u8,  lessThanFuncFor(u8));
    _ = Digraph(u16, lessThanFuncFor(u16));
    _ = Digraph(i8,  lessThanFuncFor(i8));
    _ = Digraph(i16, lessThanFuncFor(i16));
    _ = Digraph(*i8, lessThanFuncFor(*i8));
    _ = Digraph(*anyopaque, lessThanFuncFor(*anyopaque));
}
test "check if connected" {
    const Sut = Digraph(u8, lessThanFuncFor(u8));

    const map = [_]Sut.Conn {
        .{ .from = 3, .to = 0, },
        .{ .from = 0, .to = 1, },
        .{ .from = 1, .to = 3, },
    };

    var sut = try Sut.init(std.testing.allocator, map[0..]);
    defer sut.deinit();

    try std.testing.expect(sut.isConnected(0, 1));
    try std.testing.expect(!sut.isConnected(1, 0));

    try std.testing.expect(sut.isConnected(1, 3));
    try std.testing.expect(!sut.isConnected(3, 1));

    try std.testing.expect(sut.isConnected(3, 0));
    try std.testing.expect(!sut.isConnected(0, 3));

    try std.testing.expect(!sut.isConnected(0, 2));
    try std.testing.expect(!sut.isConnected(2, 0));

    try std.testing.expect(!sut.isConnected(1, 2));
    try std.testing.expect(!sut.isConnected(2, 1));
}
test "make new connection" {
    const Sut = Digraph(u8, lessThanFuncFor(u8));

    var sut = try Sut.init(std.testing.allocator, &.{});
    defer sut.deinit();

    try std.testing.expect(try sut.connectIf(2, 1));

    try std.testing.expect(sut.isConnected(2, 1));
    try std.testing.expect(!sut.isConnected(1, 2));

    try sut.connect(3, 1);

    try std.testing.expect(sut.isConnected(3, 1));
    try std.testing.expect(!sut.isConnected(1, 3));
}
test "making an existing connection fails" {
    const Sut = Digraph(u8, lessThanFuncFor(u8));

    const map = [_]Sut.Conn {
        .{ .from = 0, .to = 1, },
    };
    var sut = try Sut.init(std.testing.allocator, map[0..]);
    defer sut.deinit();

    try std.testing.expect(!try sut.connectIf(0, 1));
    try std.testing.expectError(Sut.Error.AlreadyConnected, sut.connect(0, 1));
}
test "disconnect an existing connection" {
    const Sut = Digraph(u8, lessThanFuncFor(u8));

    const map = [_]Sut.Conn {
        .{ .from = 0, .to = 1, },
        .{ .from = 2, .to = 3, },
    };
    var sut = try Sut.init(std.testing.allocator, map[0..]);
    defer sut.deinit();

    try std.testing.expect(sut.disconnectIf(0, 1));
    try std.testing.expect(!sut.isConnected(0, 1));

    try sut.disconnect(2, 3);
    try std.testing.expect(!sut.isConnected(2, 3));
}
test "disconnecting a missing connection fails" {
    const Sut = Digraph(u8, lessThanFuncFor(u8));

    var sut = try Sut.init(std.testing.allocator, &.{});
    defer sut.deinit();

    try std.testing.expect(!sut.disconnectIf(0, 1));
    try std.testing.expectError(Sut.Error.NotConnected, sut.disconnect(1, 0));
}
test "chaotic operation" {
    const Sut = Digraph(u16, lessThanFuncFor(u16));

    var sut = try Sut.init(std.testing.allocator, &.{});
    defer sut.deinit();

    const N = 100;
    const P = 109;
    const Q = 113;
    const R = 127;
    for (0..N) |v| {
        const x: Sut.Node = @intCast(v);
        try sut.connect((x*Q)%P, (x*R)%P);
    }
    for (N/2..N) |v| {
        const x: Sut.Node = @intCast(v);
        try sut.disconnect((x*Q)%P, (x*R)%P);
    }
    for (0..N/2) |v| {
        const x: Sut.Node = @intCast(v);
        try std.testing.expect(sut.isConnected((x*Q)%P, (x*R)%P));
    }
    for (N/2..N) |v| {
        const x: Sut.Node = @intCast(v);
        try std.testing.expect(!sut.isConnected((x*Q)%P, (x*R)%P));
    }
}