--!native
--!optimize 2
--!nolint LocalShadow

--[[ ## made by anexpia •…• ## ]]

-- [[ byte tags, if negative then its a dict key ]]
-- 0 -> at end: terminate, at start: deduplication table
-- 1 -> stop and start new table
-- 2 -> boolean: true
-- 3 -> boolean: false
-- 4 -> nil
-- 5 to 11 -> number
-- 12 to 14 -> buffer
-- 15 -> v3 (3 f32)
-- 16 -> v2 (2 f32)
-- 17 -> cframe
-- 18 -> ray (6 f32)
-- 19 -> enum < polymorphic - depends on setting 'rbxenum_behavior' >
-- 20 -> enumitem < polymorphic - depends on setting 'rbxenum_behavior' >
-- 21 -> numberrange (2 f32)
-- 22 -> numbersequence
-- 23 -> color3 (u8)
-- 24 -> color3 (f16)
-- 25 -> colorsequence
-- 26 -> udim (1 int32 1 f32)
-- 27 -> udim2 (2 int32 2 f32)
-- 28 -> rect (4 f32)
-- 29 -> physicalproperties (really long)
-- 30 -> brickcolor (int16)
-- 31 -> vector3int16 (3 i16)
-- 32 -> vector2int16 (2 i16)
-- 33 -> content (depends)
-- 34 -> font (depends)
-- 35 -> datetime (f64)

-- 51 -> table position (u8)
-- 52 -> table position (u16)
-- 53 -> table position (u32)
-- 54 -> u8 length string
-- 55 -> u16 length string
-- 56 -> u32 length string
-- 57 -> value deduplicate (u8)
-- 58 -> value deduplicate (u16)
-- 59 -> value deduplicate (u32)
-- 60 -> reference (u8)
-- 61 -- reference (u16)
-- 62 -- reference (u32)

-- 99 -> 2 byte custom value || <99> <-index>

-- 101 -> empty table
-- 102 -> empty string
-- 103 -> 0
-- 104 -> 1
-- 105 -> -1
-- 106 -> math.huge
-- 107 -> -math.huge
-- 108 -> nan

local uI16_max = (2 ^ 16) - 1

local Enums = require(script.Parent.Enums)
local Miscellaneous = require(script.Parent.Miscellaneous)
local RbxEnumEncoder = require(script.Parent.RbxEnumEncoder)
local Settings = require(script.Parent.Settings)
local Types = require(script.Parent.Types)

type encodeinfo = Types.encodeinfo
type WriteSettings = Types.WriteSettings

local valuetobyte = Enums.valuetobyte

local writef16 = Miscellaneous.writef16
local writeCFrame = Miscellaneous.writeCFrame
local CFrameSize = Miscellaneous.cframesize

-- shift the byte associated with a value/datatype by the number given
local tryshift = function(v: number, shiftseed: number?): number
	if shiftseed then return (v + math.random(1, 127)) % 128
	else return v end
end

-- write the byte associated with a given value or datatype
local function writebytesign(buff: buffer, cursor: number, value: number, isdict: boolean?, shiftseed: number?)
	if shiftseed then value = (value + math.random(1, 127)) % 128 end
	if isdict then value = 255 - value end

	buffer.writeu8(buff, cursor, value)
end

-- expand the size of the writingbuffer
local function expandbuffertosize(buff: buffer, newsize: number, info: encodeinfo): buffer
	if newsize > info.allocatedsize then
		newsize //= 1/1.375
		info.allocatedsize = newsize

		local newbuff = buffer.create(newsize)
		buffer.copy(newbuff, 0, buff)

		return newbuff
	end

	return buff
end

-- write every value into the buffer
local function write(
	buff: buffer,
	cursor: number,
	info: encodeinfo
): (buffer, number)

	local writesettings = info.settings

	local shiftseed: number?, dedupallowed = writesettings.shiftseed, writesettings.allowdeduplication
	local color3always6bytes = writesettings.color3always6bytes
	if color3always6bytes == nil then color3always6bytes = Settings.color3always6bytes end

	local enumbehavior, enumbehaviortype = RbxEnumEncoder.getBehaviorFromSetting(writesettings)
	local fullrbxenum: boolean = enumbehaviortype == 'full'

	local scancount: number = 1
	local deduplicatedcount: number = 0
	local referencecount: number = 0

	local valuepositionlookup = info.valuepositionlookup

	local scanlist = info.scanlist
	local referencelist = info.referencelist

	for _, tbl in scanlist do
		buff = expandbuffertosize(buff, cursor + 1, info)
		writebytesign(buff, cursor, 1, nil, shiftseed); cursor += 1

		local arraywritelength: number = rawlen(tbl)
		local isdoingdict: boolean = arraywritelength == 0

		local iternum: number = if isdoingdict then 2 else 1
		local lastkey: number = 0
		for k, v in next, tbl do
			if (not isdoingdict) then
				local diff: number = (k - lastkey) - 1

				if diff > 0 then
					buff = expandbuffertosize(buff, cursor + diff, info)
					if shiftseed then
						for i = 0, diff - 1 do
							buffer.writeu8(buff, cursor + i, (4 + math.random(1, 127)) % 128)
						end
					else buffer.fill(buff, cursor, 4, diff) end
					cursor += diff
				end

				lastkey = k
			end

			for i = 1, iternum do
				local isdict: boolean = isdoingdict and i == 1
				local value: any = if isdict then k else v

				local valuebyte: number? = valuetobyte[value]
				if valuebyte then
					if valuebyte <= 0 then
						buff = expandbuffertosize(buff, cursor + 2, info)
						writebytesign(buff, cursor, 99, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, -valuebyte); cursor += 2
					else
						buff = expandbuffertosize(buff, cursor + 1, info)
						writebytesign(buff, cursor, valuebyte, isdict, shiftseed); cursor += 1
					end

					continue
				end

				local t: string = typeof(value)
				local candeduplicate: boolean = dedupallowed and (t == 'string' or t == 'Vector3' or t == 'number' or (fullrbxenum and (t == 'EnumItem' or t == 'Enum')))

				if candeduplicate then
					if value ~= value then
						-- nan shouldnt be deduplicated, whether in a vector or as a number

						if t == 'number' then
							buff = expandbuffertosize(buff, cursor + 1, info)
							writebytesign(buff, cursor, 108, isdict, shiftseed); cursor += 1
						else
							local value: Vector3 = value
							-- vector in this case

							buff = expandbuffertosize(buff, cursor + 13, info)

							writebytesign(buff, cursor, 15, isdict, shiftseed); cursor += 1
							buffer.writef32(buff, cursor, value.X); cursor += 4
							buffer.writef32(buff, cursor, value.Y); cursor += 4
							buffer.writef32(buff, cursor, value.Z); cursor += 4
						end

						continue
					end

					if t == 'number' then
						local value: number = value
						-- numbers below 256 that arent float values shouldnt be deduplicated

						local abs = math.abs(value)
						if not (abs > 0xFF or abs % 1 ~= 0) then
							-- encoding here to avoid rechecking later
							buff = expandbuffertosize(buff, cursor + 5, info)

							writebytesign(buff, cursor, if value < 0 then 8 else 5, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, abs); cursor += 2

							continue
						end
					end

					local added: number? = valuepositionlookup[value]

					if added then
						if added <= 0xFF then -- u8
							buff = expandbuffertosize(buff, cursor + 2, info)
							writebytesign(buff, cursor, 57, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, added); cursor += 2
						elseif added <= 0xFFFF then -- u16
							buff = expandbuffertosize(buff, cursor + 3, info)
							writebytesign(buff, cursor, 58, isdict, shiftseed)
							buffer.writeu16(buff, cursor + 1, added); cursor += 3
						else -- u32
							buff = expandbuffertosize(buff, cursor + 5, info)
							writebytesign(buff, cursor, 59, isdict, shiftseed)
							buffer.writeu32(buff, cursor + 1, added); cursor += 5
						end

						continue
					else
						deduplicatedcount += 1
						valuepositionlookup[value] = deduplicatedcount
					end
				end

				if t == "nil" then
					buff = expandbuffertosize(buff, cursor + 1, info)
					writebytesign(buff, cursor, 4, isdict, shiftseed); cursor += 1
				elseif t == "number" then
					local value: number = value

					if value ~= value then
						buff = expandbuffertosize(buff, cursor + 1, info)
						writebytesign(buff, cursor, 108, isdict, shiftseed); cursor += 1
					else
						local abs = math.abs(value)
						if (value % 1) ~= 0 or abs > 0xFFFFFFFF then -- float64
							buff = expandbuffertosize(buff, cursor + 9, info)

							writebytesign(buff, cursor, 11, isdict, shiftseed)
							buffer.writef64(buff, cursor + 1, value); cursor += 9
						else
							local offset = if value < 0 then 3 else 0

							if abs <= 0xFF then -- uint8
								buff = expandbuffertosize(buff, cursor + 2, info)

								writebytesign(buff, cursor, 5 + offset, isdict, shiftseed)
								buffer.writeu8(buff, cursor + 1, abs); cursor += 2
							elseif abs <= 0xFFFF then -- uint16
								buff = expandbuffertosize(buff, cursor + 3, info)

								writebytesign(buff, cursor, 6 + offset, isdict, shiftseed)
								buffer.writeu16(buff, cursor + 1, abs); cursor += 3
							else
								buff = expandbuffertosize(buff, cursor + 5, info)

								writebytesign(buff, cursor, 7 + offset, isdict, shiftseed)
								buffer.writeu32(buff, cursor + 1, abs); cursor += 5
							end
						end
					end
				elseif t == "string" then
					local value: string = value
					local len: number = #value

					if len <= 0xFF then -- not zero-terminated since no benefit in making it zero-terminated
						buff = expandbuffertosize(buff, cursor + len + 2, info)

						writebytesign(buff, cursor, 54, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, len); cursor += 2
					elseif len <= 0xFFFF then
						buff = expandbuffertosize(buff, cursor + len + 3, info)

						writebytesign(buff, cursor, 55, isdict, shiftseed)
						buffer.writeu16(buff, cursor + 1, len); cursor += 3
					else
						buff = expandbuffertosize(buff, cursor + len + 5, info)

						writebytesign(buff, cursor, 56, isdict, shiftseed)
						buffer.writeu32(buff, cursor + 1, len); cursor += 5
					end

					buffer.writestring(buff, cursor, value); cursor += len
				elseif t == "table" then
					local value: {[any]: any} = value
					local scanpos: number? = valuepositionlookup[value]
					if (not scanpos) and (next(value) ~= nil) then
						scanpos = scancount + 1
						scancount = scanpos :: any
						valuepositionlookup[value] = scanpos
						scanlist[scanpos] = value
					end

					if scanpos then
						if scanpos <= 0xFF then
							buff = expandbuffertosize(buff, cursor + 2, info)

							writebytesign(buff, cursor, 51, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, scanpos); cursor += 2
						elseif scanpos <= 0xFFFF then
							buff = expandbuffertosize(buff, cursor + 3, info)

							writebytesign(buff, cursor, 52, isdict, shiftseed)
							buffer.writeu16(buff, cursor + 1, scanpos); cursor += 3
						else
							buff = expandbuffertosize(buff, cursor + 5, info)

							writebytesign(buff, cursor, 53, isdict, shiftseed)
							buffer.writeu32(buff, cursor + 1, scanpos); cursor += 5
						end
					else
						buff = expandbuffertosize(buff, cursor + 1, info)

						writebytesign(buff, cursor, 101, isdict, shiftseed); cursor += 1
					end
				elseif t == "buffer" then
					local value: buffer = value
					local length: number = buffer.len(value)

					if length <= 0xFF then
						buff = expandbuffertosize(buff, cursor + 2 + length, info)

						writebytesign(buff, cursor, 12, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, length)
						cursor += 2
					elseif length <= 0xFFFF then
						buff = expandbuffertosize(buff, cursor + 3 + length, info)

						writebytesign(buff, cursor, 13, isdict, shiftseed)
						buffer.writeu16(buff, cursor + 1, length)
						cursor += 3
					else
						buff = expandbuffertosize(buff, cursor + 5 + length, info)

						writebytesign(buff, cursor, 14, isdict, shiftseed)
						buffer.writeu32(buff, cursor + 1, length)
						cursor += 5
					end

					buffer.copy(buff, cursor, value)
					cursor += length
				elseif t == "Vector3" then
					local value: Vector3 = value
					buff = expandbuffertosize(buff, cursor + 13, info)

					writebytesign(buff, cursor, 15, isdict, shiftseed); cursor += 1
					buffer.writef32(buff, cursor, value.X); cursor += 4
					buffer.writef32(buff, cursor, value.Y); cursor += 4
					buffer.writef32(buff, cursor, value.Z); cursor += 4
				elseif t == "Vector2" then
					local value: Vector2 = value
					buff = expandbuffertosize(buff, cursor + 9, info)

					writebytesign(buff, cursor, 16, isdict, shiftseed); cursor += 1
					buffer.writef32(buff, cursor, value.X); cursor += 4
					buffer.writef32(buff, cursor, value.Y); cursor += 4
				elseif t == "CFrame" then
					local value: CFrame = value
					buff = expandbuffertosize(buff, cursor + CFrameSize + 1, info)

					writebytesign(buff, cursor, 17, isdict, shiftseed); cursor += 1
					writeCFrame(buff, cursor, value); cursor += CFrameSize
				elseif t == "Ray" then
					local value: Ray = value
					buff = expandbuffertosize(buff, cursor + 25, info)

					writebytesign(buff, cursor, 18, isdict, shiftseed); cursor += 1

					local pos: Vector3, dir: Vector3 = value.Origin, value.Direction
					buffer.writef32(buff, cursor, pos.X); cursor += 4
					buffer.writef32(buff, cursor, pos.Y); cursor += 4
					buffer.writef32(buff, cursor, pos.Z); cursor += 4
					buffer.writef32(buff, cursor, dir.X); cursor += 4
					buffer.writef32(buff, cursor, dir.Y); cursor += 4
					buffer.writef32(buff, cursor, dir.Z); cursor += 4
				elseif t == "Enum" then
					local value: Enum = value
					local size: number = 3 -- <u8> 19 | <u16> position
					if fullrbxenum then
						size = 2 + #RbxEnumEncoder.enumitem_to_type[value]
						-- <u8> byte | <u8> length | string
					end

					buff = expandbuffertosize(buff, cursor + size, info)
					writebytesign(buff, cursor, 19, isdict, shiftseed)
					cursor = enumbehavior.encodeEnum(buff, cursor + 1, value)
				elseif t == "EnumItem" then
					local value: EnumItem = value
					local size: number = 3 -- <u8> 20 | <u16> position
					if fullrbxenum then
						size = 4 + #RbxEnumEncoder.enumitem_to_type[value]
						-- <u8> 20 | <u8> length | string | <u16> value
					end

					buff = expandbuffertosize(buff, cursor + size, info)

					writebytesign(buff, cursor, 20, isdict, shiftseed)
					cursor = enumbehavior.encodeEnumItem(buff, cursor + 1, value)
				elseif t == "NumberRange" then
					local value: NumberRange = value
					buff = expandbuffertosize(buff, cursor + 9, info)
					writebytesign(buff, cursor, 21, isdict, shiftseed); cursor += 1

					buffer.writef32(buff, cursor, value.Min); cursor += 4
					buffer.writef32(buff, cursor, value.Max); cursor += 4
				elseif t == "NumberSequence" then
					local value: NumberSequence = value
					local keypoints: {NumberSequenceKeypoint} = value.Keypoints

					buff = expandbuffertosize(buff, cursor + 2 + #keypoints * 8, info)

					writebytesign(buff, cursor, 22, isdict, shiftseed);
					buffer.writeu8(buff, cursor + 1, #keypoints); cursor += 2

					for _, k in keypoints do
						buffer.writeu16(buff, cursor, math.round(k.Time * uI16_max))
						buffer.writef32(buff, cursor + 2, k.Value)

						writef16(buff, cursor + 6, k.Envelope); cursor += 8
					end
				elseif t == "Color3" then
					local value: Color3 = value
					local r: number, g: number, b: number = value.R, value.G, value.B

					if color3always6bytes or (math.max(r, g, b) > 1 or math.min(r, g, b) < 0) then
						buff = expandbuffertosize(buff, cursor + 7, info)

						-- one or more of values is out of 1 byte limit, so this encodes as 2 bytes each
						writebytesign(buff, cursor, 24, isdict, shiftseed)
						writef16(buff, cursor + 1, r)
						writef16(buff, cursor + 3, g)
						writef16(buff, cursor + 5, b)
						cursor += 7
					else
						buff = expandbuffertosize(buff, cursor + 4, info)

						writebytesign(buff, cursor, 23, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, math.round(r * 255))
						buffer.writeu8(buff, cursor + 2, math.round(g * 255))
						buffer.writeu8(buff, cursor + 3, math.round(b * 255))
						cursor += 4
					end
				elseif t == "ColorSequence" then
					local value: ColorSequence = value
					local keypoints: {ColorSequenceKeypoint} = value.Keypoints

					buff = expandbuffertosize(buff, cursor + 2 + #keypoints * 5, info)

					writebytesign(buff, cursor, 25, isdict, shiftseed)
					buffer.writeu8(buff, cursor + 1, #keypoints); cursor += 2

					for _, k in keypoints do
						local color = k.Value

						-- colors in colorsequences are always limited to 0-1
						buffer.writeu16(buff, cursor, math.round(k.Time * uI16_max))
						buffer.writeu8(buff, cursor + 2, math.round(color.R * 255))
						buffer.writeu8(buff, cursor + 3, math.round(color.G * 255))
						buffer.writeu8(buff, cursor + 4, math.round(color.B * 255))
						cursor += 5
					end
				elseif t == "UDim" then
					local value: UDim = value
					buff = expandbuffertosize(buff, cursor + 9, info)
					writebytesign(buff, cursor, 26, isdict, shiftseed)

					buffer.writef32(buff, cursor + 1, value.Scale)
					buffer.writei32(buff, cursor + 5, value.Offset)
					cursor += 9
				elseif t == "UDim2" then
					local value: UDim2 = value
					buff = expandbuffertosize(buff, cursor + 17, info)
					local x: UDim, y: UDim = value.X, value.Y

					writebytesign(buff, cursor, 27, isdict, shiftseed)

					buffer.writef32(buff, cursor + 1, x.Scale)
					buffer.writei32(buff, cursor + 5, x.Offset)
					buffer.writef32(buff, cursor + 9, y.Scale)
					buffer.writei32(buff, cursor + 13, y.Offset)
					cursor += 17
				elseif t == "Rect" then
					local value: Rect = value
					buff = expandbuffertosize(buff, cursor + 17, info)
					local min: Vector2, max: Vector2 = value.Min, value.Max

					writebytesign(buff, cursor, 28, isdict, shiftseed)

					buffer.writef32(buff, cursor + 1, min.X)
					buffer.writef32(buff, cursor + 5, min.Y)
					buffer.writef32(buff, cursor + 9, max.X)
					buffer.writef32(buff, cursor + 13, max.Y)
					cursor += 17
				elseif t == "PhysicalProperties" then
					local value: PhysicalProperties = value
					buff = expandbuffertosize(buff, cursor + 21, info)
					writebytesign(buff, cursor, 29, isdict, shiftseed)

					buffer.writef32(buff, cursor + 1, value.Density)
					buffer.writef32(buff, cursor + 5, value.Friction)
					buffer.writef32(buff, cursor + 9, value.Elasticity)
					buffer.writef32(buff, cursor + 13, value.FrictionWeight)
					buffer.writef32(buff, cursor + 17, value.ElasticityWeight)
					cursor += 21
				elseif t == "BrickColor" then
					local value: BrickColor = value
					buff = expandbuffertosize(buff, cursor + 3, info)
					writebytesign(buff, cursor, 30, isdict, shiftseed)
					buffer.writeu16(buff, cursor + 1, value.Number)
					cursor += 3
				elseif t == "Vector3int16" then
					local value: Vector3int16 = value
					buff = expandbuffertosize(buff, cursor + 7, info)
					writebytesign(buff, cursor, 31, isdict, shiftseed)
					buffer.writei16(buff, cursor + 1, value.X)
					buffer.writei16(buff, cursor + 3, value.Y)
					buffer.writei16(buff, cursor + 5, value.Z)
					cursor += 7
				elseif t == "Vector2int16" then
					local value: Vector2int16 = value
					buff = expandbuffertosize(buff, cursor + 5, info)
					writebytesign(buff, cursor, 32, isdict, shiftseed)
					buffer.writei16(buff, cursor + 1, value.X)
					buffer.writei16(buff, cursor + 3, value.Y)
					cursor += 5
				elseif t == "Content" then
					local value: Content = value

					if value.SourceType == Enum.ContentSourceType.Uri then
						local uri: string = value.Uri :: string
						local num: string? = string.match(uri, "%d+$")

						if num then
							buff = expandbuffertosize(buff, cursor + 10, info)
							writebytesign(buff, cursor, 33, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, 1)
							buffer.writef64(buff, cursor + 2, tonumber(num) :: number)
							cursor += 10
						else
							-- im going to make a guess and say that the content links are NOT going past 254 bytes in length 😨
							local len: number = #uri
							if len > 254 then
								error("content uri length cannot be more than 254 bytes.", 0)
							end

							buff = expandbuffertosize(buff, cursor + 2 + len, info)
							writebytesign(buff, cursor, 33, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, len + 1); cursor += 2
							buffer.writestring(buff, cursor , uri)

							cursor += len
						end

						continue
					else
						buff = expandbuffertosize(buff, cursor + 2, info)
						writebytesign(buff, cursor, 33, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, 0); cursor += 2
					end

				elseif t == "Font" then
					local value: Font = value

					local family: string = value.Family
					local num: string? = string.match(family, "%d+$")

					if num then
						buff = expandbuffertosize(buff, cursor + 13, info)

						writebytesign(buff, cursor, 34, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, 0)
						buffer.writef64(buff, cursor + 2, tonumber(num) :: number)

						cursor += 10
					else
						local familyname: string = string.match(family, "/(%w+).json$") or ""
						local length: number = #familyname

						-- same here
						if length > 254 then
							error("font name length cannot be more than 254 bytes.", 0)
						end

						buff = expandbuffertosize(buff, cursor + 5 + length, info)
						writebytesign(buff, cursor, 34, isdict, shiftseed)
						buffer.writeu8(buff, cursor + 1, length + 1); cursor += 2
						buffer.writestring(buff, cursor, familyname)

						cursor += length
					end

					buffer.writeu16(buff, cursor, RbxEnumEncoder.enumitem_to_value[value.Weight])
					buffer.writeu8(buff, cursor + 2, RbxEnumEncoder.enumitem_to_value[value.Style]); cursor += 3
				elseif t == "DateTime" then
					local value: DateTime = value

					buff = expandbuffertosize(buff, cursor + 9, info)
					writebytesign(buff, cursor, 35, isdict, shiftseed)
					buffer.writef64(buff, cursor + 1, value.UnixTimestampMillis); cursor += 9
				else
					if referencelist then
						local referenceposition: number = valuepositionlookup[value]

						if not referenceposition then
							referenceposition = referencecount + 1

							referencecount = referenceposition :: any
							valuepositionlookup[value] = referenceposition
							referencelist[referenceposition] = value
						end

						if referenceposition <= 0xFF then
							buff = expandbuffertosize(buff, cursor + 2, info)
							writebytesign(buff, cursor, 60, isdict, shiftseed)
							buffer.writeu8(buff, cursor + 1, referenceposition); cursor += 2
						elseif referenceposition <= 0xFFFF then
							buff = expandbuffertosize(buff, cursor + 3, info)
							writebytesign(buff, cursor, 61, isdict, shiftseed)
							buffer.writeu16(buff, cursor + 1, referenceposition); cursor += 3
						else
							buff = expandbuffertosize(buff, cursor + 5, info)
							writebytesign(buff, cursor, 62, isdict, shiftseed)
							buffer.writeu32(buff, cursor + 1, referenceposition); cursor += 5
						end
					elseif isdict then break
					else
						buff = expandbuffertosize(buff, cursor + 1, info)
						writebytesign(buff, cursor, 4, isdict, shiftseed); cursor += 1
					end
				end

			end

			if k == arraywritelength then iternum = 2; isdoingdict = true end
		end
	end

	return buff, cursor
end

local writingbuffer = buffer.create(256)
local writingbuffersize = 256

type BenchmarkerProfiler = {
	Begin: (name: string) -> nil;
	End: () -> nil;
}

--[[
returns the table given encoded as a buffer, and the size of buffer

<em>param</em> - value : the table to encode
<em>param</em> - writeoffset : writing starts from the offset given.
<em>param</em> - writesettings : {
. <strong>allowdeduplication: boolean</strong> -- attempt to deduplicate repeated values if enabled to reduce buffer size
. <strong>allowreferences: boolean</strong> -- if enabled, return a table containing the values it couldn't encode alongside the buffer.
. <strong>shiftseed: number</strong> -- the type bytes of values are shuffled using the seed.
. <strong>rbxenum_behavior: "full" | "compact"</strong> -- override for the default setting.
. <strong>color3always6bytes: boolean</strong> -- override for the default setting.
}
]]
return function(
	value: { [any]: any },
	writeoffset: number?,
	writesettings: WriteSettings?,
	Profiler: BenchmarkerProfiler?
): (buffer, { any }?)
	if type(value) == "table" then

		local writesettings: WriteSettings = writesettings or (Settings :: any)
		local shiftseed = writesettings.shiftseed

		if shiftseed then math.randomseed(shiftseed) end
		local referencelist = if writesettings.allowreferences then {} else nil

		if next(value) == nil then
			local b = buffer.create(1)
			buffer.writeu8(b, 0, tryshift(101, shiftseed))

			if shiftseed then math.randomseed(os.time()) end

			return b, referencelist
		end

		local cursor: number = writeoffset or 0

		local info: encodeinfo = {
			valuepositionlookup = {[value] = 1},
			scanlist = {value},

			referencelist = referencelist,
			settings = writesettings,
			allocatedsize = writingbuffersize;
		}

		local buff = writingbuffer

		if Profiler then
			Profiler.Begin("Write")
		end

		buff, cursor = write(buff, cursor, info)

		buff = expandbuffertosize(buff, cursor + 1, info)
		writebytesign(buff, cursor, 0, nil, shiftseed)

		if Profiler then
			Profiler.End()
			Profiler.Begin("Finish")
		end

		writingbuffer = buff
		writingbuffersize = info.allocatedsize

		local truncatedbuffer = buffer.create(cursor + 2)
		buffer.copy(truncatedbuffer, 0, buff, 0, cursor + 1)

		if shiftseed then math.randomseed(os.time()) end

		if Profiler then
			Profiler.End()
		end

		return truncatedbuffer, referencelist
	else
		error(`Expected a table to be encoded into a buffer, got {typeof(value)}`, 0)
	end
end


--[[
	## made by anexpia •…• ##

v---v---v---v---v---v currently supported types v---v---v---v---v---v

[x] -> datatype has deduplication
deduplication means that if the value exists in more than one table, it will be encoded only once
this reduces buffer size and time it takes to encode the value

luau types
	> string [x]
	> number [x]
	> boolean
	> nil
	> table
	> buffer
	> vector [x]

roblox types
	> Enum [x]? only if rbxenum_behavior is 'full'
	> EnumItem [x]? only if rbxenum_behavior is 'full'

	> BrickColor
	> Color3
	> ColorSequence
	> NumberSequence
	> NumberRange
	> CFrame
	> Ray

	> Vector3int16
	> Vector2
	> Vector2int16
	> Rect
	> UDim
	> UDim2
	> PhysicalProperties

	> Content
	> Font
	> DateTime

can also encode custom values into 2 bytes
	> this is done by using enums.register(name, value)
	> value is optional, itll give you a newproxy() if you dont give it a value
	> max 255 values

v---v---v---v---v---v---v---v settings v---v---v---v---v---v---v---v

> color3always6bytes: boolean >> default false
	this toggles if color3 components are always saved as float16 numbers

> rbxenum_behavior: string ( 'full' | 'compact' ) >> default 'compact'
	this has 2 behaviors currently.

	>  "full"
		encodes the enums as <19> <u8 length> <string>
		enumitems are <20> <u8 length> <string> <u16 value>

		this should be used for data saving as this guarantees the enumitems stay the same as long as roblox doesn't do anything weird

		* encoded values are deduplicated if allowed, so if the same enum & enumitem are encoded multiple times
		  it will not cost the same to encode it each time

	> "compact" *default behavior*
		this should ONLY BE USED for sending the enums & enumitems in temporary sessions!!! (like in remotes)
		the order of the enum & enumitem table is not guaranteed to be consistent between
		every session thus you will definitely encounter bugs if you use this for saving data

		* no deduplication since there's no benefit, as this saves both enums and enumitems as 3 bytes,

> serverclientsyncing: boolean
	determines whether to sync enumitems and custom values between client and server

> sanitize_nanandinf: boolean
	determines whether to turn NaN and math.huge into 0 when reading buffer content
	  this only sanitizes them for Numbers, Vectors, Rays, and CFrames
]]

return {
	read = require(script.Read);
	write = require(script.Write);
	enums = require(script.Enums)
}


--!optimize 2
-- jabby_util
local function convert_scale(unit: string, value: number): string
	local s = math.sign(value)
	value = math.abs(value)
	local prefixes = {
		[4] = "T",
		[3] = "G",
		[2] = "M",
		[1] = "k",
		[0] = "",
		[-1] = "m",
		[-2] = "μ",
		[-3] = "n",
		[-4] = "p",
	}

	local order = 0

	while value >= 1000 do
		order += 1
		value /= 1000
	end

	while value ~= 0 and value < 1 do
		order -= 1
		value *= 1000
	end

	local format = "%.01f"
	if value >= 100 then
		format = "%.01f"
		value = math.floor(value)
	elseif value >= 10 then
		format = "%.01f"
		value = math.floor(value * 1e1) / 1e1
	elseif value >= 1 then
		format = "%.02f"
		value = math.floor(value * 1e2) / 1e2
	end

	return string.format(format, value * s) .. prefixes[math.clamp(order, -4, 4)] .. unit
end

return convert_scale


--!optimize 2
-- jabby_util
local function convert_units(unit: string, value: number): string
	local s = math.sign(value)
	value = math.abs(value)
	local prefixes = {
		[4] = "T",
		[3] = "G",
		[2] = "M",
		[1] = "k",
		[0] = " ",
		[-1] = "m",
		[-2] = "u",
		[-3] = "n",
		[-4] = "p",
	}

	local order = 0

	while value >= 1000 do
		order += 1
		value /= 1000
	end

	while value ~= 0 and value < 1 do
		order -= 1
		value *= 1000
	end

	if value >= 100 then
		value = math.floor(value)
	elseif value >= 10 then
		value = math.floor(value * 1e1) / 1e1
	elseif value >= 1 then
		value = math.floor(value * 1e2) / 1e2
	end

	return value * s .. prefixes[order] .. unit
end

return convert_units


--!optimize 2
--!strict
--!native

--$Packages
local ThreadPool = require("@self/threadpool")

--$Types
export type Connection<T...> = {
	disconnect: (self: Connection<T...>) -> (),
	reconnect: (self: Connection<T...>) -> (),
	connected: boolean,
}

export type Event<T...> = {
	connect: (self: Event<T...>, fn: (T...) -> ()) -> Connection<T...>,
	connectOnce: (self: Event<T...>, fn: (T...) -> ()) -> Connection<T...>,
	disconnectAll: (self: Event<T...>) -> (),
	wait: (self: Event<T...>) -> T...,
	from: (emitter: (event: Event<T...>, T...) -> ()) -> Event<T...>,
}

export type Bindable<T...> = {
	new: () -> Bindable<T...>,
	wrap: (signal: RBXScriptSignal<T...>) -> Bindable<T...>,
	fire: (self: Bindable<T...>, T...) -> (),
	Destroy: (self: Bindable<T...>) -> (), -- why PascalCase? because it does call Destroy method on RBXScriptConnection if it exists
	RBXScriptConnection: RBXScriptConnection?,
	event: Event<T...>,
}

export type Signal<T...> = {
	new: () -> Signal<T...>,
	wrap: (signal: RBXScriptSignal<T...>) -> Signal<T...>,
	connect: (self: Signal<T...>, fn: (T...) -> ()) -> Connection<T...>,
	connectOnce: (self: Signal<T...>, fn: (T...) -> ()) -> Connection<T...>,
	disconnectAll: (self: Signal<T...>) -> (),
	wait: (self: Signal<T...>) -> T...,
	fire: (self: Signal<T...>, T...) -> (),
	Destroy: (self: Signal<T...>) -> (),
	RBXScriptConnection: RBXScriptConnection?,
}

local Connection = {} :: Connection<...any>
(Connection :: any).__index = Connection

local function disconnect(self)
	if not self.connected then
		return
	end
	self.connected = false

	local next = (self :: any)._next
	local prev = (self :: any)._prev

	if next then
		next._prev = prev
	end
	if prev then
		prev._next = next
	end

	local signal = (self :: any)._signal
	if signal._head == self then
		signal._head = next
	end
end
Connection.disconnect = disconnect -- i know this is annoying but this is only for micro optimization

function reconnect(self)
	if self.connected then
		return
	end
	self.connected = true

	local signal = (self :: any)._signal
	local head = signal._head
	if head then
		head._prev = self
	end
	signal._head = self;
	(self :: any)._next = head;
	(self :: any)._prev = false
end
Connection.reconnect = reconnect

--$Lime
local Bindable = {} :: Bindable<...any>
(Bindable :: any).__index = Bindable

local Event = {} :: Event<...any>
(Event :: any).__index = Event

local Signal = {} :: Signal<...any>
(Signal :: any).__index = Signal

local rbxConnect, rbxDisconnect
do
	local bindable = Instance.new("BindableEvent")
	rbxConnect = bindable.Event.Connect
	rbxDisconnect = bindable.Event:Connect(function() end).Disconnect
	bindable:Destroy()
end

local function connect<T...>(self, fn)
	local head = (self :: any)._head
	local cn = setmetatable({
		connected = true,
		_signal = self,
		_fn = fn,
		_next = head,
		_prev = false,
	}, Connection)

	if head then
		head._prev = cn
	end
	(self :: any)._head = cn :: any

	return cn :: any
end
Signal.connect = connect
Event.connect = connect :: any

local function once<T...>(self, fn)
	local cn
	cn = connect(self, function(...: any)
		disconnect(cn)
		fn(...)
	end)
	return cn
end
Signal.connectOnce = once
Event.connectOnce = once :: any

local function wait<T...>(self): T...
	local thread = coroutine.running()
	local cn
	cn = connect(self, function(...)
		disconnect(cn)
		task.spawn(thread, ...)
	end)
	return coroutine.yield()
end
Signal.wait = wait
Event.wait = wait :: any

local function fire<T...>(self, ...)
	local event = (self :: any).event or self
	local cn = event._head
	while cn do
		-- local thread
		-- if #freeThreads > 0 then
		-- 	thread = freeThreads[#freeThreads]
		-- 	freeThreads[#freeThreads] = nil
		-- else
		-- 	thread = coroutine.create(yielder)
		-- 	coroutine.resume(thread)
		-- end

		-- task.spawn(thread, cn._fn, thread, ...)

		-- now we spawn in shared threadpool!

		ThreadPool.spawn(cn._fn, ...)
		cn = cn._next
	end
end
Signal.fire = fire
Bindable.fire = fire :: any

local function disconnectAll<T...>(self)
	local cn = (self :: any)._head
	while cn do
		disconnect(cn)
		cn = cn._next
	end
end
Event.disconnectAll = disconnectAll
Signal.disconnectAll = disconnectAll

local function Destroy<T...>(self)
	disconnectAll((self :: any).event or self)
	local cn = self.RBXScriptConnection
	if cn then
		rbxDisconnect(cn)
		self.RBXScriptConnection = nil :: any
	end
end
Signal.Destroy = Destroy
Bindable.Destroy = Destroy :: any

function Event.from(emitter)
	return setmetatable({ _emitter = emitter, _head = false }, Event) :: any
end

function Bindable.new()
	local event = setmetatable({ _head = false }, Event)
	return setmetatable({ event = event }, Bindable) :: any
end

function Bindable.wrap(signal)
	local wrapper = Bindable.new()
	wrapper.RBXScriptConnection = rbxConnect(signal, function(...)
		fire(wrapper :: any, ...)
	end)
	return wrapper
end

function Signal.new()
	return setmetatable({ _head = false }, Signal) :: any
end

function Signal.wrap(signal)
	local wrapper = Signal.new()
	wrapper.RBXScriptConnection = rbxConnect(signal, function(...)
		fire(wrapper, ...)
	end)
	return wrapper
end

local function createEmitter(): <T...>(event: Event<T...>, T...) -> ()
	local function emitter(event, ...)
		if not event._emitter or event._emitter ~= emitter then
			error("Emitted an invalid event")
		end
		fire(event, ...)
	end
	return emitter :: any
end

return {
	Signal = Signal,
	Event = Event,
	Bindable = Bindable,
	createEmitter = createEmitter,
}


--!optimize 2
local Types = require("./Types")
local GradClip = require("./GradClip")
local Anomaly = require("./Util/Anomaly")

type Tensor = Types.Tensor

local Debug = {}

function Debug.checkTensor(t: Tensor, name: string?): boolean
	if Anomaly.hasBadValues(t) then
		warn(`Gradien.Debug: tensor '{name or "?"}' has NaN/Inf/invalid values`)
		return false
	end
	return true
end

function Debug.checkGradients(params: { Tensor }, label: string?): (boolean, number?)
	for i, p in params do
		local g = p._grad
		if g and Anomaly.hasBadValues(g) then
			warn(`Gradien.Debug: bad gradient in param {i} ({label or ""})`)
			return false, i
		end
	end
	return true, nil
end

function Debug.wrapOptimizer(
	opt,
	params: { Tensor },
	cfgs: { maxGradNorm: number?, clipValue: number?, warnOnNaN: boolean?, label: string? }?
)
	local cfg = cfgs or {}
	local maxGradNorm = cfg.maxGradNorm
	local clipValue = cfg.clipValue
	local warnOnNaN = cfg.warnOnNaN
	if type(warnOnNaN) ~= "boolean" then
		warnOnNaN = true
	end
	local label = cfg.label or "wrapped optimizer"

	local wrapped = {}

	function wrapped:zeroGrad()
		opt:zeroGrad()
	end

	function wrapped:step()
		if maxGradNorm then
			GradClip.clipNorm(params, maxGradNorm)
		end
		if clipValue then
			GradClip.clipValue(params, clipValue)
		end
		if warnOnNaN then
			Debug.checkGradients(params, label)
		end
		opt:step()
	end

	return wrapped
end

function Debug.checkModel(model, label: string?)
	local params = model:parameters()
	Debug.checkGradients(params, label)
	for i, p in params do
		Debug.checkTensor(p, `{label or "model"}.param{i}`)
	end
end

return Debug


--!native
--!optimize 2
local GC = {}
function GC.clipValue(params, clip)
	task.desynchronize()
	for _, p in params do
		if p._grad then
			for i = 1, #p._grad._storage do
				p._grad._storage[i] = math.clamp(p._grad._storage[i], -clip, clip)
			end
		end
	end
	task.synchronize()
end
function GC.clipNorm(params, maxNorm)
	task.desynchronize()
	local total = 0
	for _, p in params do
		if p._grad then
			for i = 1, #p._grad._storage do
				total += p._grad._storage[i] ^ 2
			end
		end
	end
	total = math.sqrt(total)
	if total > maxNorm then
		local scale = maxNorm / (total + 1e-9)
		for _, p in params do
			if p._grad then
				for i = 1, #p._grad._storage do
					p._grad._storage[i] *= scale
				end
			end
		end
	end
	task.synchronize()
end
return GC


--!native
--!optimize 2
local Types = require("./Types")
local Initializer = {}

local function _randn(): number
	local u1 = math.random()
	local u2 = math.random()
	if u1 <= 1e-7 then
		u1 = 1e-7
	end
	return math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
end

function Initializer.xavierUniform(W: Types.Tensor)
	task.desynchronize()
	local fanIn = W._shape[2]
	local fanOut = W._shape[1]
	local bound = math.sqrt(6 / (fanIn + fanOut))
	for i = 1, #W._storage do
		W._storage[i] = (math.random() * 2 - 1) * bound
	end
	task.synchronize()
end

function Initializer.kaimingUniform(W: Types.Tensor)
	task.desynchronize()
	local fanIn = W._shape[2]
	local bound = math.sqrt(6 / fanIn)
	for i = 1, #W._storage do
		W._storage[i] = (math.random() * 2 - 1) * bound
	end
	task.synchronize()
end

function Initializer.heUniform(W: Types.Tensor)
	Initializer.kaimingUniform(W)
end

-- N(0, 2 / fanIn)
function Initializer.heNormal(W: Types.Tensor)
	task.desynchronize()
	local fanIn = W._shape[2]
	local std = math.sqrt(2 / fanIn)
	for i = 1, #W._storage do
		W._storage[i] = _randn() * std
	end
	task.synchronize()
end

-- U(-sqrt(3 / fanIn), sqrt(3 / fanIn))
function Initializer.lecunUniform(W: Types.Tensor)
	task.desynchronize()
	local fanIn = W._shape[2]
	local bound = math.sqrt(3 / fanIn)
	for i = 1, #W._storage do
		W._storage[i] = (math.random() * 2 - 1) * bound
	end
	task.synchronize()
end

-- N(0, 1 / fanIn)
function Initializer.lecunNormal(W: Types.Tensor)
	task.desynchronize()
	local fanIn = W._shape[2]
	local std = math.sqrt(1 / fanIn)
	for i = 1, #W._storage do
		W._storage[i] = _randn() * std
	end
	task.synchronize()
end

function Initializer.zeros(B: Types.Tensor)
	task.desynchronize()
	for i = 1, #B._storage do
		B._storage[i] = 0
	end
	task.synchronize()
end

return Initializer


--!native
--!optimize 2
local Types = require("./Types")
local M = {}

function M.accuracy(logits: Types.Tensor, targets: Types.Tensor): number
	task.desynchronize()
	local C, B = logits._shape[1], logits._shape[2]
	local correct = 0
	for j = 1, B do
		local best = -1 / 0
		local arg = 1
		for i = 1, C do
			local v = logits._storage[(i - 1) * B + j]
			if v > best then
				best = v
				arg = i
			end
		end
		if arg == targets[j] then
			correct += 1
		end
	end
	task.synchronize()
	return correct / B
end

function M.topk(logits: Types.Tensor, targets: Types.Tensor, k: number): number
	task.desynchronize()
	local C, B = logits._shape[1], logits._shape[2]
	local hit = 0
	for j = 1, B do
		local arr = {}
		for i = 1, C do
			arr[i] = logits._storage[(i - 1) * B + j]
		end
		table.sort(arr, function(a, b)
			return a > b
		end)
		local thresh = arr[math.min(k, C)]
		local t = targets[j]
		local ok = false
		for i = 1, C do
			if logits._storage[(i - 1) * B + j] >= thresh and i == t then
				ok = true
			end
		end
		if ok then
			hit += 1
		end
	end
	task.synchronize()
	return hit / B
end

function M.mse(pred: Types.Tensor, target: Types.Tensor): number
	task.desynchronize()
	local s = 0
	for i = 1, #pred._storage do
		local d = pred._storage[i] - target._storage[i]
		s += d * d
	end
	task.synchronize()
	return s / pred._shape[#pred._shape]
end

function M.prf1(pred: Types.Tensor, target: Types.Tensor, C: number): (number, number, number)
	task.desynchronize()
	local yhat: { number }, B: number
	if type(pred) == "table" and pred._storage ~= nil then
		local Cx, Bb = pred._shape[1], pred._shape[2]
		B = Bb
		yhat = table.create(B, 1)
		for j = 1, B do
			local best, arg = -1 / 0, 1
			for i = 1, Cx do
				local v = pred._storage[(i - 1) * B + j]
				if v > best then
					best, arg = v, i
				end
			end
			yhat[j] = arg
		end
	else
		yhat = pred
		B = #yhat
	end
	local tpC: { number }, fpC: { number }, fnC: { number } = {}, {}, {}
	for c = 1, C do
		tpC[c] = 0
		fpC[c] = 0
		fnC[c] = 0
	end
	for j = 1, B do
		local p, t = yhat[j], target[j]
		if p and t and p >= 1 and p <= C and t >= 1 and t <= C then
			if p == t then
				tpC[p] += 1
			else
				fpC[p] += 1
				fnC[t] += 1
			end
		end
	end
	local sumP, sumR, sumF = 0, 0, 0
	for c = 1, C do
		local pc = tpC[c] / (tpC[c] + fpC[c] + 1e-9)
		local rc = tpC[c] / (tpC[c] + fnC[c] + 1e-9)
		local fc = 2 * pc * rc / (pc + rc + 1e-9)
		sumP += pc
		sumR += rc
		sumF += fc
	end
	task.synchronize()
	return sumP / C, sumR / C, sumF / C
end

function M.confusion(pred: Types.Tensor, target: Types.Tensor, C: number): { { number } }
	task.desynchronize()
	local yhat: { number }, B: number
	if type(pred) == "table" and pred._storage ~= nil then
		local Cx, Bb = pred._shape[1], pred._shape[2]
		B = Bb
		yhat = table.create(B, 1)
		for j = 1, B do
			local best, arg = -1 / 0, 1
			for i = 1, Cx do
				local v = pred._storage[(i - 1) * B + j]
				if v > best then
					best, arg = v, i
				end
			end
			yhat[j] = arg
		end
	else
		yhat = pred
		B = #yhat
	end
	local mat: { { number } } = {}
	for i = 1, C do
		mat[i] = table.create(C, 0)
		for j = 1, C do
			mat[i][j] = 0
		end
	end
	for j = 1, B do
		local t, p = target[j], yhat[j]
		if t >= 1 and t <= C and p >= 1 and p <= C then
			mat[t][p] += 1
		end
	end
	task.synchronize()
	return mat
end

return M


--!native
--!optimize 2
local Types = require("./Types")
local Util = require("./Util")
local State = {}

local assert = Util.Assert

local function mapParamsToIndices(params: { Types.Tensor }): { [Types.Tensor]: string }
	local map = {}
	for i, p in params do
		map[p] = `{i}`
	end
	return map
end

--@param params { Types.Tensor }
--@return Snapshot
function State.dump(params: { Types.Tensor }): Types.Snapshot
	task.desynchronize()
	local n = #params
	local out: Types.Snapshot = table.create(n)
	for i = 1, n do
		local t = params[i]
		assert(t and t._storage and t._shape, `State.dump: expected Tensor at params[{i}]`)
		local shp = table.clone(t._shape)
		local dat = table.create(#t._storage, 0)
		for k = 1, #t._storage do
			dat[k] = t._storage[k]
		end
		out[i] = { shape = shp, dtype = t._dtype or "f32", data = dat }
	end
	task.synchronize()
	return out
end

function State.load(params: { Types.Tensor }, snap: Types.Snapshot): boolean
	if snap == nil then
		return false
	end
	task.desynchronize()
	local changed = false
	if #params ~= #snap then
		warn(`Gradien: Model parameter count mismatch (Model: {#params}, Snapshot: {#snap})`)
	end
	for i = 1, math.min(#params, #snap) do
		local t = params[i]
		local s = snap[i]
		local data = s and (s.data or s.values)
		if t and data then
			if #t._storage ~= #data then
				warn(`Gradien: Layer {i} size mismatch (Model: {#t._storage}, Snapshot: {#data})`)
			end
			local n = math.min(#t._storage, #data)
			for k = 1, n do
				t._storage[k] = data[k] :: number
			end
			changed = true
		end
	end
	task.synchronize()
	return changed
end

function State.dumpModel(model: Types.Module): Types.Snapshot
	return State.dump(model:parameters())
end

function State.loadModel(model: Types.Module, snap: Types.Snapshot): boolean
	return State.load(model:parameters(), snap)
end

-- Serializes optimizer state by converting Tensor Keys -> Index Strings
function State.dumpOptimizer(optim: Types.Optimizer, modelParams: { Types.Tensor }): Types.OptimizerState?
	if optim and type(optim.getState) == "function" then
		local rawState = optim:getState()
		if not rawState then
			return nil
		end
		local serialized = table.clone(rawState)
		local paramMap = mapParamsToIndices(modelParams)
		for k, v in rawState do
			if type(v) == "table" then
				local isTensorMap = false
				local newTable = {}
				for tk, tv in v do
					if type(tk) == "table" and tk._storage then
						isTensorMap = true
						local idx = paramMap[tk]
						if idx then
							newTable[idx] = tv
						end
					else
						newTable[tk] = tv
					end
				end
				if isTensorMap then
					serialized[k] = newTable
				end
			end
		end
		return serialized
	end
	return nil
end

function State.loadOptimizer(optim: Types.Optimizer, snap: Types.OptimizerState?, modelParams: { Types.Tensor }): ()
	if not snap then
		return
	end
	if optim and type(optim.setState) == "function" then
		local restored = table.clone(snap)
		for k, v in snap do
			if type(v) == "table" then
				local reconstructed = {}
				local wasParamMap = false
				for idxStr, val in v do
					local idx = tonumber(idxStr)
					if idx and modelParams[idx] then
						wasParamMap = true
						reconstructed[modelParams[idx]] = val
					else
						reconstructed[idxStr] = val
					end
				end
				if wasParamMap then
					restored[k] = reconstructed
				end
			end
		end
		optim:setState(restored)
	end
end

function State.dumpTrainer(trainer: Types.Trainer): Types.TrainerSnapshot
	local params = trainer.model:parameters()
	return {
		model = State.dumpModel(trainer.model),
		optimizer = State.dumpOptimizer(trainer.optimizer, params),
		step = trainer.step or nil,
		epoch = trainer.epoch or nil,
		bestMetric = trainer.bestMetric,
	}
end

function State.loadTrainer(trainer: Types.Trainer, snap: Types.TrainerSnapshot)
	if not snap then
		error("State.loadTrainer: no snapshot provided")
	end
	State.loadModel(trainer.model, snap.model)
	if snap.optimizer then
		local params = trainer.model:parameters()
		State.loadOptimizer(trainer.optimizer, snap.optimizer, params)
	end
	trainer.step = snap.step or trainer.step or 0
	trainer.epoch = snap.epoch or trainer.epoch or 0
	trainer.bestMetric = snap.bestMetric
end

function State.toBuffer(snap: Types.Snapshot): buffer
	return Util.Buffer.write(snap)
end

function State.fromBuffer(b: buffer): Types.Snapshot?
	return Util.Buffer.read(b)
end

return State


--!native
--!optimize 2
local Tensor = {}

local Util = require("./Util")
local Types = require("./Types")
local Tape -- lazy-load

local TensorMethods = {}
local function attachMethods(t)
	for name, fn in TensorMethods do
		t[name] = fn
	end
end

export type Tensor = Types.Tensor

local sizeFromShape = Util.sizeFromShape
local assert = Util.Assert

local sqrt = math.sqrt
local log = math.log
local cos = math.cos
local sin = math.sin
local floor = math.floor
local pi = math.pi
local random = math.random

local function computeStrides(shape: { number }): { number }
	task.desynchronize()
	local rank = #shape
	local strides = table.create(rank)
	local acc: number = 1
	for i = rank, 1, -1 do
		strides[i] = acc
		acc *= shape[i]
	end
	task.synchronize()
	return strides
end

local function newTensor(
	storage: { number },
	shape: { number },
	dtype: Types.DType?,
	requiresGrad: boolean?,
	sizeOverride: number?
): Tensor
	local shapeClone = table.clone(shape)
	local strides = computeStrides(shapeClone)
	local t: Types.Tensor = {
		_storage = storage,
		_shape = shapeClone,
		_dtype = dtype or "f64",
		_requiresGrad = requiresGrad == true,
		_size = sizeOverride or #storage,
		_strides = strides,
		_parents = nil,
		_backprop = nil,
		_grad = nil,
	}
	attachMethods(t :: any)
	return t :: any
end

local function ensureSize(t: Types.Tensor): number
	local size = t._size
	if not size then
		size = #t._storage
		t._size = size
	end
	return size
end

local function is_contiguous(self: Tensor): boolean
	if not self._strides then
		return true
	end
	local computed = computeStrides(self._shape)
	for i = 1, #self._strides do
		if self._strides[i] ~= computed[i] then
			return false
		end
	end
	return true
end

local function contiguous(self: Tensor): Tensor
	if is_contiguous(self) then
		return self
	end
	local size = ensureSize(self)
	local out = newTensor(self._storage, self._shape, self._dtype, self._requiresGrad, size)
	return out
end

local function ensureStrides(t: Types.Tensor): { number }
	local strides = t._strides
	if not strides then
		strides = computeStrides(t._shape)
		t._strides = strides
	end
	return strides
end

function Tensor.zeros(shape: { number }, dtype: Types.DType?, requiresGrad: boolean?): Tensor
	dtype = dtype or "f64"
	local n = sizeFromShape(shape)
	return newTensor(table.create(n, 0), shape, dtype, requiresGrad, n)
end

function Tensor.ones(shape: { number }, dtype: Types.DType?, requiresGrad: boolean?): Tensor
	dtype = dtype or "f64"
	local n = sizeFromShape(shape)
	return newTensor(table.create(n, 1), shape, dtype, requiresGrad, n)
end

function Tensor.empty(shape: { number }, dtype: Types.DType?, requiresGrad: boolean?): Tensor
	dtype = dtype or "f64"
	local n = sizeFromShape(shape)
	return newTensor(table.create(n), shape, dtype, requiresGrad, n)
end

function Tensor.fromArray(data: { number }, shape: { number }, dtype: Types.DType?, requiresGrad: boolean?): Tensor
	local n = sizeFromShape(shape)
	assert(#data == n, ("Tensor.fromArray: got %d elements, expected %d"):format(#data, n))
	dtype = dtype or "f64"
	return newTensor(table.clone(data), shape, dtype, requiresGrad, n)
end

function Tensor.randn(shape: { number }, dtype: Types.DType?, requiresGrad: boolean?): Tensor
	dtype = dtype or "f64"
	local n = sizeFromShape(shape)
	local storage = table.create(n)
	task.desynchronize()
	for i = 1, n, 2 do
		local u1 = random()
		local u2 = random()
		local radius = sqrt(-2 * log(u1))
		local theta = 2 * pi * u2
		local z0 = radius * cos(theta)
		local z1 = radius * sin(theta)
		storage[i] = z0
		if i + 1 <= n then
			storage[i + 1] = z1
		end
	end
	task.synchronize()
	return newTensor(storage, shape, dtype, requiresGrad, n)
end

local function reshape(self: Tensor, newShape: { number }): Tensor
	local size = ensureSize(self)
	assert(sizeFromShape(newShape) == size, "reshape: element count mismatch")
	local out = newTensor(self._storage, newShape, self._dtype, self._requiresGrad, size)
	if not self._requiresGrad then
		return out
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	return Tape.makeNode(out, { self }, function(gy: Types.Tensor)
		if self._requiresGrad then
			self._grad = self._grad or Tensor.zeros(self._shape, self._dtype, false)
			task.desynchronize()
			local gs, ys = (self._grad :: Types.Tensor)._storage, gy._storage
			for i = 1, #ys do
				gs[i] += ys[i]
			end
			task.synchronize()
		end
	end)
end

local function expand(self: Tensor, newShape: { number }): Tensor
	local rank = #self._shape
	local newRank = #newShape
	assert(newRank >= rank, "expand: new rank must be >= old rank")
	local expansionStrides = table.create(newRank, 0)
	local oldShapeExpanded = table.create(newRank, 1)
	local oldStrides = ensureStrides(self)
	local oldStridesExpanded = table.create(newRank, 0)
	for i = 1, rank do
		oldShapeExpanded[newRank - rank + i] = self._shape[i]
		oldStridesExpanded[newRank - rank + i] = oldStrides[i]
	end
	for i = 1, newRank do
		if oldShapeExpanded[i] == 1 and newShape[i] > 1 then
			expansionStrides[i] = 0
		elseif oldShapeExpanded[i] == newShape[i] then
			expansionStrides[i] = oldStridesExpanded[i]
		else
			error("expand: incompatible shapes")
		end
	end
	local outSize = sizeFromShape(newShape)
	local outStorage = table.create(outSize)
	task.desynchronize()
	local src = self._storage
	local coord = table.create(newRank, 0)
	for flatOut = 1, outSize do
		local srcOffset = 0
		for d = 1, newRank do
			srcOffset += coord[d] * expansionStrides[d]
		end
		outStorage[flatOut] = src[srcOffset + 1]

		for d = newRank, 1, -1 do
			local val = coord[d] + 1
			if val < newShape[d] then
				coord[d] = val
				break
			end
			coord[d] = 0
		end
	end
	task.synchronize()
	local out = newTensor(outStorage, newShape, self._dtype, self._requiresGrad, outSize)
	if not self._requiresGrad then
		return out
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	local dtype = self._dtype
	local shapeCopy = table.clone(self._shape)
	return Tape.makeNode(out, { self }, function(gy: Types.Tensor)
		if not self._requiresGrad then
			return
		end
		self._grad = self._grad or Tensor.zeros(shapeCopy, dtype, false)
		task.desynchronize()
		local gs = (self._grad :: Types.Tensor)._storage
		local gys = gy._storage
		local coordBack = table.create(newRank, 0)
		for flatOut = 1, outSize do
			local srcOffset = 0
			for d = 1, newRank do
				srcOffset += coordBack[d] * expansionStrides[d]
			end
			gs[srcOffset + 1] += gys[flatOut]

			for d = newRank, 1, -1 do
				local val = coordBack[d] + 1
				if val < newShape[d] then
					coordBack[d] = val
					break
				end
				coordBack[d] = 0
			end
		end
		task.synchronize()
	end)
end

local function slice(self: Tensor, dim: number, startIdx: number, endIdx: number?, step: number?): Tensor
	local shape = self._shape
	local rank = #shape
	assert(dim >= 1 and dim <= rank, "slice: dim out of range")

	local sizeDim = shape[dim]
	local s = startIdx
	local e = endIdx or sizeDim
	local st = step or 1

	assert(st > 0, "slice: step must be > 0")
	assert(s >= 1 and s <= sizeDim, "slice: startIdx out of bounds")
	assert(e >= 1 and e <= sizeDim, "slice: endIdx out of bounds")
	assert(e >= s, "slice: endIdx must be >= startIdx")

	local length = floor((e - s) / st) + 1
	local outShape = table.clone(shape)
	outShape[dim] = length

	local outSize = sizeFromShape(outShape)
	local inStrides = ensureStrides(self)
	local outStorage = table.create(outSize)
	local indexMap = self._requiresGrad and table.create(outSize) or nil

	task.desynchronize()
	local coord = table.create(rank, 1)
	local srcStorage = self._storage

	for flatOut = 1, outSize do
		local posDim = s + (coord[dim] - 1) * st
		local flatIn0 = 0
		for d = 1, rank do
			local cd = coord[d]
			if d == dim then
				cd = posDim
			end
			flatIn0 += (cd - 1) * inStrides[d]
		end
		local storageIndex = flatIn0 + 1
		outStorage[flatOut] = srcStorage[storageIndex]
		if indexMap then
			indexMap[flatOut] = storageIndex
		end

		for d = rank, 1, -1 do
			local val = coord[d] + 1
			if val <= outShape[d] then
				coord[d] = val
				break
			end
			coord[d] = 1
		end
	end
	task.synchronize()

	local out = newTensor(outStorage, outShape, self._dtype, self._requiresGrad, outSize)
	if not self._requiresGrad then
		return out
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	local shapeCopy = table.clone(shape)
	local mapCopy = indexMap
	local dtype = self._dtype
	return Tape.makeNode(out, { self }, function(gy: Types.Tensor)
		if not self._requiresGrad then
			return
		end
		self._grad = self._grad or Tensor.zeros(shapeCopy, dtype, false)
		task.desynchronize()
		local gsrc = (self._grad :: Types.Tensor)._storage
		local gys = gy._storage
		if mapCopy then
			for i = 1, #gys do
				local dest = mapCopy[i]
				if dest then
					gsrc[dest] += gys[i]
				end
			end
		end
		task.synchronize()
	end)
end

local function transpose(self: Tensor, dim1: number?, dim2: number?): Tensor
	local shape = self._shape
	local rank = #shape
	local d1 = dim1
	local d2 = dim2
	if d1 == nil and d2 == nil then
		assert(rank == 2, "transpose: default dims only valid for rank-2 tensors")
		d1, d2 = 1, 2
	else
		assert(d1 ~= nil and d2 ~= nil, "transpose: both dim1 and dim2 must be provided")
	end
	assert(d1 >= 1 and d1 <= rank, "transpose: dim1 out of range")
	assert(d2 >= 1 and d2 <= rank, "transpose: dim2 out of range")
	assert(d1 ~= d2, "transpose: dim1 and dim2 must be different")
	local perm = table.create(rank)
	for i = 1, rank do
		perm[i] = i
	end
	perm[d1], perm[d2] = perm[d2], perm[d1]
	local outShape = table.create(rank)
	for i = 1, rank do
		outShape[i] = shape[perm[i]]
	end
	local size = ensureSize(self)
	local outSize = size
	local outStrides = computeStrides(outShape)
	local dst = table.create(outSize)
	local indexMap = self._requiresGrad and table.create(outSize) or nil
	task.desynchronize()
	local coordIn = table.create(rank, 1)
	for flatIn = 1, size do
		local flatOut0 = 0
		for outDim = 1, rank do
			local inDim = perm[outDim]
			local c = coordIn[inDim]
			flatOut0 += (c - 1) * outStrides[outDim]
		end
		local outIndex = flatOut0 + 1
		dst[outIndex] = self._storage[flatIn]
		if indexMap then
			indexMap[outIndex] = flatIn
		end
		for d = rank, 1, -1 do
			local val = coordIn[d] + 1
			if val <= shape[d] then
				coordIn[d] = val
				break
			end
			coordIn[d] = 1
		end
	end
	task.synchronize()
	local out = newTensor(dst, outShape, self._dtype, self._requiresGrad, outSize)
	if not self._requiresGrad then
		return out
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	local shapeCopy = table.clone(shape)
	local mapCopy = indexMap
	local dtype = self._dtype
	return Tape.makeNode(out, { self }, function(gy: Types.Tensor)
		if not self._requiresGrad then
			return
		end
		self._grad = self._grad or Tensor.zeros(shapeCopy, dtype, false)
		task.desynchronize()
		local gsrc = (self._grad :: Types.Tensor)._storage
		local gys = gy._storage
		if mapCopy then
			for i = 1, #gys do
				local srcIdx = mapCopy[i]
				if srcIdx then
					gsrc[srcIdx] += gys[i]
				end
			end
		end
		task.synchronize()
	end)
end

local function narrow(self: Tensor, dim: number, startIdx: number, length: number): Tensor
	assert(length >= 1, "narrow: length must be >= 1")
	local endIdx = startIdx + length - 1
	return slice(self, dim, startIdx, endIdx, 1)
end

local function sum(self: Tensor, dim: number?): Tensor
	if dim then
		error("Tensor.sum with dim not implemented yet")
	end
	local s = 0
	local storage = self._storage
	task.desynchronize()
	for i = 1, #storage do
		s += storage[i]
	end
	task.synchronize()
	local out = newTensor({ s }, { 1 }, self._dtype, self._requiresGrad, 1)
	if not self._requiresGrad then
		return out
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	return Tape.makeNode(out, { self }, function(gy: Types.Tensor)
		if self._requiresGrad then
			self._grad = self._grad or Tensor.zeros(self._shape, self._dtype, false)
			local g = gy._storage[1]
			task.desynchronize()
			local gs = (self._grad :: Types.Tensor)._storage
			for i = 1, #gs do
				gs[i] += g
			end
			task.synchronize()
		end
	end)
end

local function detach(self: Tensor): Tensor
	return newTensor(self._storage, self._shape, self._dtype, false, ensureSize(self))
end

local function noGrad(self: Tensor): ()
	self._requiresGrad = false
	self._grad = nil
end

local function backward(self: Tensor, grad: Tensor?)
	if not self._requiresGrad then
		error("Tensor does not require grad")
	end
	if not Tape then
		Tape = require("./autograd/Tape") :: any
	end
	if not grad then
		if self._size == 1 then
			grad = Tensor.ones(self._shape, self._dtype, false)
		else
			error("backward: grad must be provided for non-scalar tensors")
		end
	end
	Tape.backwardFrom(self, grad)
end

TensorMethods.reshape = reshape
Tensor.reshape = reshape

TensorMethods.expand = expand
Tensor.expand = expand

TensorMethods.is_contiguous = is_contiguous
Tensor.is_contiguous = is_contiguous

TensorMethods.contiguous = contiguous
Tensor.contiguous = contiguous

TensorMethods.slice = slice
Tensor.slice = slice

TensorMethods.transpose = transpose
Tensor.transpose = transpose

TensorMethods.narrow = narrow
Tensor.narrow = narrow

TensorMethods.sum = sum
Tensor.sum = sum

TensorMethods.detach = detach
Tensor.detach = detach

TensorMethods.noGrad = noGrad
Tensor.noGrad = noGrad

TensorMethods.backward = backward
Tensor.backward = backward

return Tensor


--!native
--!optimize 2
local Tape = require("./autograd/Tape")
local Types = require("./Types")
local Losses = require("./nn/Losses")
local Metrics = require("./Metrics")
local Util = require("./Util")

local Trainer = {}
Trainer.__index = Trainer

local assert = Util.Assert

local function asArrayCallbacks(raw: Types.Callbacks?, default: { Types.CallbackFn }?): { Types.CallbackFn }
	if type(raw) ~= "table" then
		return default or {}
	end
	if not next(raw) then
		return raw :: { Types.CallbackFn }
	end

	local tbl = raw :: Types.CallbackTable
	local out: { Types.CallbackFn } = {}
	if tbl.onStep then
		local fn = tbl.onStep
		table.insert(out, function(event: Types.TrainerEvent, ctx: Types.CallbackContext)
			if event == "step_end" then
				fn(ctx)
			end
		end)
	end
	if tbl.onEpochEnd then
		local fn = tbl.onEpochEnd
		table.insert(out, function(event: Types.TrainerEvent, ctx: Types.CallbackContext)
			if event == "epoch_end" then
				fn(ctx)
			end
		end)
	end
	if tbl.onBest then
		local fn = tbl.onBest
		table.insert(out, function(event: Types.TrainerEvent, ctx: Types.CallbackContext)
			if event == "best" then
				fn(ctx)
			end
		end)
	end
	return out
end

local function fireCallbacks(callbacks: { Types.CallbackFn }, event: Types.TrainerEvent, ctx: Types.CallbackContext)
	for _, cb in callbacks do
		cb(event, ctx)
	end
end

local function runSimple(cfg: Types.SimpleConfig): ()
	task.desynchronize()
	local model = cfg.model
	local opt = cfg.optimizer(model:parameters())
	local lossFn = cfg.loss
	local loader = cfg.loader
	local metricFn = cfg.metric
	local steps = cfg.steps or 100
	local reportEvery = cfg.reportEvery or 10
	local onMetric = cfg.onMetric
	local callbacks = asArrayCallbacks(cfg.callbacks, nil)

	local bestMetric: number? = nil

	for step = 1, steps do
		opt:zeroGrad()
		local X, Y = loader()
		local out = model:forward(X)
		local lossValue, grad = lossFn(out, Y)
		Tape.backwardFrom(out, grad)
		opt:step()

		local metricValue: number? = nil
		if metricFn and (step % reportEvery == 0) then
			local m = metricFn(out, Y)
			metricValue = m
			if onMetric then
				onMetric(m)
			end
			if bestMetric == nil or m > bestMetric then
				bestMetric = m
				if #callbacks > 0 then
					fireCallbacks(callbacks, "best", {
						epoch = 1,
						step = step,
						loss = lossValue,
						metric = m,
						model = model,
						optimizer = opt,
					})
				end
			end
		end

		if #callbacks > 0 then
			fireCallbacks(callbacks, "step_end", {
				epoch = 1,
				step = step,
				loss = lossValue,
				metric = metricValue,
				model = model,
				optimizer = opt,
			})
		end
	end
	task.synchronize()
end

--@param cfg { model: Module, optimizerFactory: (({ Tensor }) -> Optimizer)?, optimizer: Optimizer?, loss: LossFn, metric: ((Tensor, Tensor) -> number)?, reportEvery: number?, callbacks: Callbacks? }
--@return Trainer
function Trainer.new(cfg: Types.TrainerConfig)
	local model = cfg.model
	local params = model:parameters()

	local opt: Types.Optimizer
	if cfg.optimizerFactory then
		opt = cfg.optimizerFactory(params)
	else
		assert(cfg.optimizer, "Trainer.new: provide optimizerFactory OR optimizer instance")
		opt = cfg.optimizer :: Types.Optimizer
	end

	local self = setmetatable({
		model = model,
		optimizer = opt,
		loss = cfg.loss,
		metric = cfg.metric,
		reportEvery = cfg.reportEvery or 10,
		lrScheduler = cfg.lrScheduler,
		callbacks = asArrayCallbacks(cfg.callbacks, {}),
		bestMetric = nil,
		step = 0,
		epoch = 1,
	}, Trainer)

	return self
end

--@param makeLoader () -> () -> (Tensor?, Tensor?)
--@param opts { epochs: number?, stepsPerEpoch: number?, onMetric: ((number, number, number) -> ())? }
function Trainer:fit(makeLoader: () -> () -> (Types.Tensor?, Types.Tensor?), opts: Types.FitOptions?)
	local fitOpts = (opts or {}) :: Types.FitOptions
	local epochs = fitOpts.epochs or self.epoch or 1
	local stepsPerEpoch = fitOpts.stepsPerEpoch
	local onMetric = fitOpts.onMetric

	for epoch = 1, epochs do
		self.epoch = epoch
		local loader = makeLoader()
		local lastLoss: number? = nil
		local lastMetric: number? = nil
		local stepsThisEpoch = 0

		while true do
			local X, Y = loader()
			if not X then
				break
			end
			self.step += 1
			stepsThisEpoch += 1

			self.optimizer:zeroGrad()
			local out = self.model:forward(X)
			local lossValue, grad = self.loss(out, Y)
			Tape.backwardFrom(out, grad)
			self.optimizer:step()
			if self.lrScheduler then
				local newLr = self.lrScheduler(self.step)
				if type(newLr) == "number" then
					if type(self.optimizer.setLr) == "function" then
						self.optimizer:setLr(newLr)
					elseif type(self.optimizer.setLR) == "function" then
						self.optimizer:setLR(newLr)
					end
				end
			end

			lastLoss = lossValue
			local metricValue: number? = nil

			if self.metric and self.reportEvery > 0 and (self.step % self.reportEvery == 0) then
				local m = self.metric(out, Y)
				metricValue = m
				lastMetric = m

				if onMetric then
					onMetric(epoch, self.step, m)
				end
				if self.bestMetric == nil or m > (self.bestMetric :: number) then
					self.bestMetric = m
					fireCallbacks(self.callbacks, "best", {
						epoch = epoch,
						step = self.step,
						loss = lossValue,
						metric = m,
						model = self.model,
						optimizer = self.optimizer,
					})
				end
			end

			fireCallbacks(self.callbacks, "step_end", {
				epoch = epoch,
				step = self.step,
				loss = lossValue,
				metric = metricValue,
				model = self.model,
				optimizer = self.optimizer,
			})

			if stepsPerEpoch and stepsThisEpoch >= stepsPerEpoch then
				break
			end
		end

		fireCallbacks(self.callbacks, "epoch_end", {
			epoch = epoch,
			step = self.step,
			loss = lastLoss,
			metric = lastMetric,
			model = self.model,
			optimizer = self.optimizer,
		})
	end
end

--@param cfg Types.TrainerConfig
--@param opts { smoothing: number?, loss: Types.LossFn?, metric: ((Types.Tensor, any) -> number)? }?
function Trainer.newClassification(
	cfg: Types.TrainerConfig,
	opts: {
		smoothing: number?,
		loss: Types.LossFn?,
		metric: ((Types.Tensor, any) -> number)?,
	}?
)
	local cfgOpts = (opts or {}) :: {
		smoothing: number?,
		loss: Types.LossFn?,
		metric: ((Types.Tensor, any) -> number)?,
	}
	local smoothing = cfgOpts.smoothing
	local lossFn = cfgOpts.loss or cfg.loss
	if not lossFn then
		lossFn = function(logits, target)
			return Losses.cross_entropy_backward(logits, target, smoothing)
		end
	end
	local metricFn = cfgOpts.metric or cfg.metric
	if not metricFn then
		metricFn = function(logits, target)
			return Metrics.accuracy(logits, target)
		end
	end
	return Trainer.new({
		model = cfg.model,
		optimizerFactory = cfg.optimizerFactory,
		optimizer = cfg.optimizer,
		loss = lossFn,
		metric = metricFn,
		reportEvery = cfg.reportEvery,
		lrScheduler = cfg.lrScheduler,
		callbacks = cfg.callbacks,
	})
end

function Trainer:run(loader: () -> (Types.Tensor?, Types.Tensor?), opts: Types.FitOptions?)
	local function makeLoader()
		return loader
	end
	self:fit(makeLoader, opts)
end

setmetatable(Trainer, {
	__call = function(_, cfg: Types.SimpleConfig)
		return runSimple(cfg)
	end,
})

return Trainer


--!strict
export type DType = "f64" | "f32" | "i32" | "u8"
export type Shape = { number }
export type Storage = { number }

export type Tensor = {
	_storage: Storage,
	_shape: Shape,
	_dtype: DType,
	_requiresGrad: boolean,
	_size: number,
	_strides: { number }?,
	_parents: { Tensor }?,
	_backprop: ((grad: Tensor) -> ())?,
	_grad: Tensor?,
}

export type Module<TIn = Tensor, TOut = Tensor> = {
	forward: (self: Module<TIn, TOut>, x: TIn) -> TOut,
	parameters: (self: Module<TIn, TOut>) -> { Tensor },
}

export type OptimizerState = {
	lr: number,
	b1: number,
	b2: number,
	eps: number,
	t: number,
	m: { [Tensor]: { [number]: number } },
	v: { [Tensor]: { [number]: number } },
}
export type Optimizer = {
	step: (self: Optimizer) -> (),
	zeroGrad: (self: Optimizer) -> (),
	getState: (self: Optimizer) -> OptimizerState?,
	setState: (self: Optimizer, state: OptimizerState?) -> (),
}

export type LossFn = (pred: Tensor, target: Tensor) -> (number, Tensor)

export type Dataset = {
	size: (self: Dataset) -> number,
	at: (self: Dataset, i: number) -> (Tensor, Tensor),
	batch: (self: Dataset, indices: { number }) -> (Tensor, Tensor),
	slice: (self: Dataset, i: number, j: number) -> (Tensor, Tensor),
}

export type BatchIter = () -> ({ number }, { number }) | nil

export type Snapshot = { SnapshotTensor }
export type TrainerSnapshot = {
	model: Snapshot,
	optimizer: OptimizerState?,
	step: number?,
	epoch: number?,
	bestMetric: number?,
}

export type SnapshotTensor = {
	shape: { number },
	dtype: DType,
	data: { number },
}

export type TrainerEvent = "step_end" | "epoch_end" | "best"

export type CallbackContext = {
	epoch: number,
	step: number?,
	loss: number?,
	metric: number?,
	model: Module,
	optimizer: Optimizer,
}

export type CallbackFn = (TrainerEvent, CallbackContext) -> ()

export type CallbackTable = {
	onStep: ((CallbackContext) -> ())?,
	onEpochEnd: ((CallbackContext) -> ())?,
	onBest: ((CallbackContext) -> ())?,
}

export type Callbacks = { CallbackFn } | CallbackTable

export type SimpleConfig = {
	model: Module,
	optimizer: (params: { Tensor }) -> Optimizer,
	loss: LossFn,
	loader: () -> (Tensor, Tensor),
	metric: ((Tensor, Tensor) -> number)?,
	steps: number?,
	reportEvery: number?,
	onMetric: ((number) -> ())?,
	callbacks: Callbacks?,
}

export type TrainerConfig = {
	model: Module,
	optimizerFactory: (({ Tensor }) -> Optimizer)?,
	optimizer: Optimizer?,
	loss: LossFn,
	metric: ((Tensor, Tensor) -> number)?,
	reportEvery: number?,
	lrScheduler: ((number) -> number)?,
	callbacks: Callbacks?,
}

export type FitOptions = {
	epochs: number?,
	stepsPerEpoch: number?,
	onMetric: ((number, number, number) -> ())?, -- (epoch, step, metric)
}

export type Trainer = {
	step: number,
	epoch: number,
	model: Module,
	optimizer: Optimizer,
	loss: LossFn,
	metric: ((Tensor, Tensor) -> number)?,
	lrScheduler: ((number) -> number)?,
	reportEvery: number?,
	callbacks: { CallbackFn },
	bestMetric: number?,
}

export type Sample = {
	name: string,
	total: number,
	self: number,
	count: number,
	avg: number,
	selfAvg: number,
	min: number,
	max: number,
	last: number,
	percent: number,
}

export type Stat = {
	name: string,
	total: number,
	selfTime: number,
	count: number,
	min: number,
	max: number,
	last: number,
}

export type Marker = {
	name: string,
	start: number,
	child: number,
}

export type Callable = (...any) -> ...any

return nil


--!native
--!optimize 2
local Types = require("../Types")
local Anomaly = {}
function Anomaly.hasBadValues(t: Types.Tensor)
	for i = 1, #t._storage do
		local v = t._storage[i]
		if v ~= v or v == math.huge or v == -math.huge then
			return true
		end
	end
	return false
end
function Anomaly.hasNaN(t: Types.Tensor)
	for i = 1, #t._storage do
		local v = t._storage[i]
		if v ~= v then
			return true
		end
	end
	return false
end
function Anomaly.hasInf(t: Types.Tensor)
	for i = 1, #t._storage do
		local v = t._storage[i]
		if v == math.huge or v == -math.huge then
			return true
		end
	end
	return false
end
return Anomaly


--!optimize 2
--@Eternity_Devs
return function<T>(condition: T, errorMessage: string, level: number?): T
	if not (condition :: any) then
		error(errorMessage, level or 2)
	end
	return condition
end


--!native
--!native
--!optimize 2
--@Eternity_Devs
--from Warp lib
local DedicatedBuffer = {}

local create = buffer.create
local copy = buffer.copy
local writei8 = buffer.writei8
local writei16 = buffer.writei16
local writei32 = buffer.writei32
local writeu8 = buffer.writeu8
local writeu16 = buffer.writeu16
local writeu32 = buffer.writeu32
local writef32 = buffer.writef32
local writef64 = buffer.writef64
local writestring = buffer.writestring

local default: { [string]: number } = {
	point = 0,
	next = 0,
	size = 128,
	bufferSize = 128,
}

local function copyMethod(self: any, offset: number, b: buffer?, src: buffer?, srcOffset: number?, count: number?)
	if not b then
		copy(create(count or default.size), offset, src or self.buffer, srcOffset, count)
	else
		copy(b, offset, src or self.buffer, srcOffset, count)
	end
end

local function alloc(self: any, byte: number)
	local need = self.next + byte
	if need > self.size then
		local newSize = self.size
		repeat
			newSize = math.floor(newSize * 1.25 + 0.5)
		until newSize < need and not (newSize < need) or newSize >= need
		if newSize < need then
			newSize = need
		end

		local newBuffer: buffer = create(newSize)
		copy(newBuffer, 0, self.buffer, 0, self.next)
		self.buffer = newBuffer
		self.size = newSize
	end

	self.point = self.next
	self.next += byte
end

local function build(self: any): buffer
	local p: number = self.next > self.point and self.next or self.point
	local buildBuffer: buffer = create(p)

	copy(buildBuffer, 0, self.buffer, 0, p)
	return buildBuffer
end

local function buildAndRemove(self: any): (buffer, any?)
	local p: number = self.next > self.point and self.next or self.point
	local buildBuffer: buffer = create(p)
	local ref = #self.ref > 0 and table.clone(self.ref) or nil

	copy(buildBuffer, 0, self.buffer, 0, p)

	self:remove()
	return buildBuffer, ref
end

local function wi8(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 1)
	writei8(self.buffer, self.point, val)
end

local function wi16(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 2)
	writei16(self.buffer, self.point, val)
end

local function wi32(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 4)
	writei32(self.buffer, self.point, val)
end

local function wu8(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 1)
	writeu8(self.buffer, self.point, val)
end

local function wu16(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 2)
	writeu16(self.buffer, self.point, val)
end

local function wu32(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 4)
	writeu32(self.buffer, self.point, val)
end

local function wf32(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 4)
	writef32(self.buffer, self.point, val)
end

local function wf64(self: any, val: number, allocBytes: number?)
	if not val then
		return
	end
	self:alloc(allocBytes or 8)
	writef64(self.buffer, self.point, val)
end

local function wstring(self: any, val: string)
	if not val then
		return
	end
	self:alloc(#val)
	writestring(self.buffer, self.point, val)
end

local function wType(self: any, ref: number)
	writeu8(self.buffer, self.point, ref)
	self.point += 1
end

local function wRef(self: any, value: any)
	if not value then
		return
	end
	table.insert(self.ref, value)
	local index = #self.ref
	self:alloc(2)
	writeu16(self.buffer, self.point, index)
end

local function pack(self: any, data: any)
	local dataType: string = typeof(data)
	if dataType == "nil" then
		self:wi8(0)
	elseif dataType == "Instance" then
		self:wi8(-1)
		self:wRef(data)
	elseif dataType == "table" then
		local count = 0
		for _ in data do
			count += 1
		end
		local isArray: boolean = (#data == count)
		if isArray then
			self:wi8(-2)
			self:wu16(count)
			for i = 1, count do
				self:pack(data[i])
			end
		else
			self:wi8(-3)
			self:wu16(count)
			for k, v in data do
				self:pack(k)
				self:pack(v)
			end
		end
	elseif dataType == "EnumItem" then
		self:wi8(-4)
		local enumName: string = `{data.EnumType}`
		self:wi8(#enumName)
		self:wstring(enumName)
		self:wu16(data.Value)
	elseif dataType == "BrickColor" then
		self:wi8(-5)
		self:wi16(data.Number)
	elseif dataType == "Enum" then
		self:wi8(-6)
		local enumName: string = data.Name
		self:wi8(#enumName)
		self:wstring(enumName)
	elseif dataType == "number" then
		if math.floor(data) == data then
			if data >= 0 and data <= 255 then
				self:wi8(1)
				self:wu8(data)
			elseif data >= -32768 and data <= 32767 then
				self:wi8(2)
				self:wi16(data)
			elseif data >= -2147483648 and data <= 2147483647 then
				self:wi8(3)
				self:wi32(data)
			else
				self:wi8(4)
				self:wf64(data)
			end
		else
			self:wi8(4)
			self:wf64(data)
		end
	elseif dataType == "boolean" then
		self:wi8(data and 5 or 6)
	elseif dataType == "string" then
		local length = #data
		if length <= 255 then
			self:wi8(7)
			self:wu8(length)
		elseif length <= 65535 then
			self:wi8(8)
			self:wu16(length)
		else
			self:wi8(9)
			self:wi32(length)
		end
		self:wstring(data)
	elseif dataType == "Vector3" then
		self:wi8(10)
		self:wf32(data.X)
		self:wf32(data.Y)
		self:wf32(data.Z)
	elseif dataType == "Vector2" then
		self:wi8(11)
		self:wf32(data.X)
		self:wf32(data.Y)
	elseif dataType == "CFrame" then
		self:wi8(12)
		for _, v in { data:GetComponents() } do
			self:wf32(v)
		end
	elseif dataType == "Color3" then
		self:wi8(13)
		local r = math.clamp(math.floor(data.R * 255 + 0.5), 0, 255)
		local g = math.clamp(math.floor(data.G * 255 + 0.5), 0, 255)
		local b = math.clamp(math.floor(data.B * 255 + 0.5), 0, 255)
		self:wu8(r)
		self:wu8(g)
		self:wu8(b)
	else
		warn(`Unsupported data type: {dataType} value: {data}`)
	end
end

local function flush(self: any)
	self.point = default.point
	self.next = default.next
	self.size = default.size
	self.buffer = create(default.bufferSize)
	table.clear(self.ref)
end

local function remove(self: any)
	self:flush()
	table.clear(self)
end

local methods = {
	copy = copyMethod,
	alloc = alloc,
	build = build,
	buildAndRemove = buildAndRemove,
	wi8 = wi8,
	wi16 = wi16,
	wi32 = wi32,
	wu8 = wu8,
	wu16 = wu16,
	wu32 = wu32,
	wf32 = wf32,
	wf64 = wf64,
	wstring = wstring,
	wType = wType,
	wRef = wRef,
	pack = pack,
	flush = flush,
	remove = remove,
}

for name, fn in methods do
	DedicatedBuffer[name] = fn
end

function DedicatedBuffer.new()
	local obj = {
		point = default.point,
		next = default.next,
		size = default.size,
		buffer = create(default.bufferSize),
		ref = {},
	}
	for name, fn in methods do
		obj[name] = fn
	end
	return obj
end

export type DedicatedType = typeof(DedicatedBuffer.new())

return DedicatedBuffer.new :: typeof(DedicatedBuffer.new)


--!native
--!native
--!optimize 2
--@Eternity_Devs
--from Warp lib & extended
local assert = require("./Assert")
local Buffer = {}
Buffer.__index = Buffer

local Dedicated = require(script.Dedicated)

local tostring = buffer.tostring
local fromstring = buffer.fromstring
local readu8 = buffer.readu8
local readi8 = buffer.readi8
local readu16 = buffer.readu16
local readi16 = buffer.readi16
local readi32 = buffer.readi32
local readf32 = buffer.readf32
local readf64 = buffer.readf64
local readstring = buffer.readstring
local writeu8 = buffer.writeu8
local writei32 = buffer.writei32
local writef32 = buffer.writef32
local writef64 = buffer.writef64
local len = buffer.len
local create = buffer.create

local dtypeInfo = {
	f64 = {
		bytes = 8,
		read = readf64,
		write = writef64,
	},
	f32 = {
		bytes = 4,
		read = readf32,
		write = writef32,
	},
	i32 = {
		bytes = 4,
		read = readi32,
		write = writei32,
	},
	u8 = {
		bytes = 1,
		read = readu8,
		write = writeu8,
	},
}

local function ensureDType(dtype: string?)
	local info = dtypeInfo[dtype or "f64"]
	assert(info, `Buffer: unsupported dtype {dtype}`)
	return info
end

local function readValue(b: buffer, position: number, ref: { any }?): (any, number)
	local typeByte = readi8(b, position)
	position += 1
	if typeByte == 0 then -- nil
		return nil, position
	elseif typeByte == -1 then -- Instance
		if type(ref) ~= "table" then
			return nil, position + 1 -- still consume the index (compat if kept u8)
		end
		local idx = readu16(b, position)
		local value = ref[idx]
		if typeof(value) == "Instance" then
			return value, position + 2
		end
		return nil, position + 2
	elseif typeByte == -2 then -- array
		local length = readu16(b, position)
		position += 2
		local array = {}
		for _ = 1, length do
			local value
			value, position = readValue(b, position, ref)
			table.insert(array, value)
		end
		return array, position
	elseif typeByte == -3 then -- dictionary
		local length = readu16(b, position)
		position += 2
		local dict = {}
		for _ = 1, length do
			local key, value
			key, position = readValue(b, position, ref)
			value, position = readValue(b, position, ref)
			dict[key] = value
		end
		return dict, position
	elseif typeByte == -4 then -- EnumItem
		local length = readu8(b, position)
		local enumName = readstring(b, position + 1, length)
		local enumVal = readu16(b, position + 1 + length)
		return Enum[enumName]:FromValue(enumVal), position + 3 + length
	elseif typeByte == -5 then -- BrickColor
		local value = readi16(b, position)
		return BrickColor.new(value), position + 2
	elseif typeByte == -6 then -- Enum
		local length = readu8(b, position)
		local enumName = readstring(b, position + 1, length)
		return Enum[enumName], position + 1 + length
	elseif typeByte == 1 then -- int u8
		local value = readu8(b, position)
		return value, position + 1
	elseif typeByte == 2 then -- int i16
		local value = readi16(b, position)
		return value, position + 2
	elseif typeByte == 3 then -- int i32
		local value = readi32(b, position)
		return value, position + 4
	elseif typeByte == 4 then -- f64
		local value = readf64(b, position)
		return value, position + 8
	elseif typeByte == 5 then
		return true, position
	elseif typeByte == 6 then
		return false, position
	elseif typeByte == 7 then -- string u8
		local length = readu8(b, position)
		local value = readstring(b, position + 1, length)
		return value, position + length + 1
	elseif typeByte == 8 then -- string u16
		local length = readu16(b, position)
		local value = readstring(b, position + 2, length)
		return value, position + length + 2
	elseif typeByte == 9 then -- string i32
		local length = readi32(b, position)
		local value = readstring(b, position + 4, length)
		return value, position + length + 4
	elseif typeByte == 10 then -- Vector3
		local x = readf32(b, position)
		local y = readf32(b, position + 4)
		local z = readf32(b, position + 8)
		return Vector3.new(x, y, z), position + 12
	elseif typeByte == 11 then -- Vector2
		local x = readf32(b, position)
		local y = readf32(b, position + 4)
		return Vector2.new(x, y), position + 8
	elseif typeByte == 12 then -- CFrame
		local components = {}
		for i = 1, 12 do
			table.insert(components, readf32(b, position + (i - 1) * 4))
		end
		return CFrame.new(unpack(components)), position + 48
	elseif typeByte == 13 then -- Color3
		local r = readu8(b, position)
		local g = readu8(b, position + 1)
		local channelB = readu8(b, position + 2)
		return Color3.fromRGB(r, g, channelB), position + 3
	end
	error(`Unsupported type marker: {typeByte}`)
end

function Buffer.new(): Dedicated.DedicatedType
	return Dedicated()
end

function Buffer.convert(b: buffer): string
	return tostring(b)
end

function Buffer.revert(s: string): buffer
	return fromstring(s)
end

function Buffer.write(data: any): (buffer, any?)
	local newBuffer = Dedicated()
	newBuffer:pack(data)
	return newBuffer:buildAndRemove()
end

function Buffer.read(b: buffer, ref: { any }?): any?
	local position = 0
	local result = {}
	while position < len(b) do
		local value
		value, position = readValue(b, position, ref)
		table.insert(result, value)
	end
	ref = nil
	return result
end

function Buffer.writeDenseArray(data: { number }, dtype: string?, existing: buffer?): buffer
	local info = ensureDType(dtype)
	local count = #data
	local totalBytes = count * info.bytes
	local buf = existing
	if not buf or len(buf) ~= totalBytes then
		buf = create(totalBytes)
	end
	local offset = 0
	for i = 1, count do
		info.write(buf, offset, data[i])
		offset += info.bytes
	end
	return buf
end

function Buffer.readDenseArray(buf: buffer, dtype: string?, out: { number }?): { number }
	local info = ensureDType(dtype)
	local bytes = info.bytes
	local count = math.floor(len(buf) / bytes)
	local dest = out
	if not dest or #dest ~= count then
		dest = table.create(count)
	end
	local offset = 0
	for i = 1, count do
		dest[i] = info.read(buf, offset)
		offset += bytes
	end
	return dest
end

function Buffer.copyDenseToColumn(
	buf: buffer,
	dtype: string?,
	target: { number },
	column: number,
	batch: number,
	length: number?
)
	local info = ensureDType(dtype)
	local bytes = info.bytes
	local elements = length or math.floor(len(buf) / bytes)
	local offset = 0
	local idx = column
	for _ = 1, elements do
		target[idx] = info.read(buf, offset)
		offset += bytes
		idx += batch
	end
end

return Buffer :: typeof(Buffer)


--!native
--!optimize 2
local Types = require("../Types")
local GradStats = {}

local function stats(t: Types.Tensor): (number, number, number, number)
	local n = #t._storage
	if n == 0 then
		return 0, 0, 0, 0
	end
	task.desynchronize()
	local minv, maxv = math.huge, -math.huge
	local sum, sumsq = 0, 0
	for i = 1, n do
		local v = t._storage[i]
		if v < minv then
			minv = v
		end
		if v > maxv then
			maxv = v
		end
		sum += v
		sumsq += v * v
	end
	local mean = sum / n
	local var = math.max(sumsq / n - mean * mean, 0)
	local std = math.sqrt(var)
	task.synchronize()
	return minv, maxv, mean, std
end

function GradStats.forModel(model: Types.Module): { { min: number, max: number, mean: number, std: number, n: number } }
	local ps = model:parameters()
	local out = {}
	for i, p in ps do
		if p._grad then
			local minv, maxv, mean, std = stats(p._grad)
			out[i] = {
				min = minv,
				max = maxv,
				mean = mean,
				std = std,
				n = #p._grad._storage,
			}
		end
	end
	return out
end

return GradStats


--!optimize 2
local Hooks = {}

function Hooks.addForwardHook(module, fn)
	local orig = module.forward
	module.___fwd = orig
	module.forward = function(self, ...)
		local out = orig(self, ...)
		fn(self, out)
		return out
	end
end

function Hooks.removeForwardHook(module)
	if module.___fwd then
		module.forward = module.___fwd
		module.___fwd = nil
	end
end

return Hooks


--!native
--!optimize 2
local Parallel = {}
function Parallel.with(fn: () -> ...any): ...any
	task.desynchronize()
	local res = { fn() }
	task.synchronize()
	return table.unpack(res)
end
function Parallel.withSafe(fn: () -> ...any): (boolean, ...any)
	task.desynchronize()
	local res = { pcall(fn) }
	task.synchronize()
	return table.unpack(res)
end
return Parallel


--!native
--!optimize 2
local Types = require("../Types")
local assert = require("./Assert")

local Profiler = {}

local fmt = string.format
local stats: { [string]: Types.Stat } = {}
local stack: { Types.Marker } = {}
local enabled: boolean = true

local function ensureStat(name: string): Types.Stat
	local stat = stats[name]
	if not stat then
		stat = {
			name = name,
			total = 0,
			selfTime = 0,
			count = 0,
			min = math.huge,
			max = 0,
			last = 0,
		}
		stats[name] = stat
	end
	return stat
end

local function clamp(value: number): number
	if value < 0 or value ~= value then
		return 0
	end
	return value
end

local function formatTime(value: number, precision: number): string
	local absValue, unit, scaled = math.abs(value)
	if absValue >= 1 then
		unit = "s"
		scaled = value
	elseif absValue >= 1e-3 then
		unit = "ms"
		scaled = value * 1e3
	elseif absValue >= 1e-6 then
		unit = "us"
		scaled = value * 1e6
	else
		unit = "ns"
		scaled = value * 1e9
	end
	local prec = math.max(0, precision)
	local spec = `%.{prec}f{unit}`
	return fmt(spec, scaled, unit)
end

function Profiler.reset(): ()
	table.clear(stats)
	table.clear(stack)
end

function Profiler.flush(): ()
	table.clear(stack)
end

function Profiler.isEnabled(): boolean
	return enabled
end

function Profiler.setEnabled(flag: boolean): ()
	if flag == enabled then
		return
	end
	enabled = flag
	if not enabled then
		table.clear(stack)
	end
end

function Profiler.start(name: string): ()
	if not enabled then
		return
	end
	table.insert(stack, {
		name = name,
		start = os.clock(),
		child = 0,
	})
end

function Profiler.stop(name: string?): number?
	if not enabled then
		return nil
	end
	local depth = #stack
	if depth == 0 then
		-- error("Profiler.stop called without matching Profiler.start", 2)
		return nil
	end
	local marker = stack[depth]
	stack[depth] = nil
	if name and marker.name ~= name then
		error(`Profiler.stop mismatch (expected '{marker.name}', got '{name}')`, 2)
	end
	local elapsed = clamp(os.clock() - marker.start)
	local exclusive = clamp(elapsed - marker.child)
	local stat = ensureStat(marker.name)
	stat.total += elapsed
	stat.selfTime += exclusive
	stat.count += 1
	if elapsed < stat.min then
		stat.min = elapsed
	end
	if elapsed > stat.max then
		stat.max = elapsed
	end
	stat.last = elapsed
	if depth > 1 then
		stack[depth - 1].child += elapsed
	end
	return elapsed
end

local function callWithProfiling(name: string, fn: Types.Callable, ...: any): (number, ...any)
	assert(type(fn) == "function", "Profiler: expected callable")
	if not enabled then
		local results = table.pack(pcall(fn, ...))
		if not results[1] then
			error(results[2], 0)
		end
		return 0, table.unpack(results, 2, results.n)
	end
	local packed = table.pack(pcall(function(...)
		Profiler.start(name)
		return fn(...)
	end, ...))
	local elapsed = Profiler.stop(name) or 0
	if not packed[1] then
		error(packed[2], 0)
	end
	return elapsed, table.unpack(packed, 2, packed.n)
end

function Profiler.scope(name: string, fn: Types.Callable, ...: any): (number, ...any)
	return callWithProfiling(name, fn, ...)
end

function Profiler.wrap(fn: Types.Callable, name: string?)
	assert(type(fn) == "function", "Profiler.wrap expects a function")
	local label = name or "Profiler.wrap"
	return function(...)
		return select(2, callWithProfiling(label, fn, ...))
	end
end

function Profiler.instrument(target: any, methods: { string } | string, opts: { prefix: string? }?)
	assert(type(target) == "table", "Profiler.instrument expects a table target")
	local names
	if type(methods) == "string" then
		names = { methods }
	else
		names = methods
	end
	local prefix = if opts and opts.prefix then opts.prefix else `{target}`
	local originals = {}
	for _, key in names do
		local fn = target[key]
		assert(type(fn) == "function", `Profiler.instrument: '{key}' is not a function`)
		originals[key] = fn
		target[key] = function(...)
			return select(2, callWithProfiling(`${prefix}.{key}`, fn, ...))
		end
	end
	return function()
		for key, original in originals do
			target[key] = original
		end
	end
end

local function copyStat(stat: Types.Stat, total: number): Types.Sample
	local minv = if stat.min == math.huge then 0 else stat.min
	local count = stat.count
	local avg = if count > 0 then stat.total / count else 0
	local selfAvg = if count > 0 then stat.selfTime / count else 0
	local pct = if total > 0 then (stat.total / total) * 100 else 0
	return {
		name = stat.name,
		total = stat.total,
		self = stat.selfTime,
		count = count,
		avg = avg,
		selfAvg = selfAvg,
		min = minv,
		max = stat.max,
		last = stat.last,
		percent = pct,
	}
end

local function collectTotals(): (number, number)
	local total = 0
	local n = 0
	task.desynchronize()
	for _, stat in stats do
		total += stat.total
		n += 1
	end
	task.synchronize()
	return total, n
end

function Profiler.snapshot(): ({ Types.Sample }, number)
	local total, count = collectTotals()
	local out = table.create(count)
	local i = 1
	task.desynchronize()
	for _, stat in stats do
		out[i] = copyStat(stat, total)
		i += 1
	end
	task.synchronize()
	return out, total
end

function Profiler.get(name: string): Types.Sample?
	local stat = stats[name]
	if not stat then
		return nil
	end
	local total = 0
	task.desynchronize()
	for _, s in stats do
		total += s.total
	end
	task.synchronize()
	return copyStat(stat, total)
end

type ReportOptions = {
	sortBy: ("total" | "self" | "avg" | "selfAvg" | "max" | "count" | "last" | "min" | "percent")?,
	limit: number?,
	precision: number?,
	print: boolean?,
}

local function sortEntries(entries: { Types.Sample }, key: string)
	table.sort(entries, function(a, b)
		local av = (a :: any)[key] or 0
		local bv = (b :: any)[key] or 0
		if av == bv then
			return a.name < b.name
		end
		return av > bv
	end)
end

function Profiler.report(opts: ReportOptions?): ({ Types.Sample }, number)
	local options: ReportOptions = opts or {}
	local sortBy = options.sortBy or "total"
	local precision = options.precision or 4
	local entries, total = Profiler.snapshot()
	sortEntries(entries, sortBy)
	local limit = options.limit
	if limit and limit < #entries then
		for i = limit + 1, #entries do
			entries[i] = nil
		end
	end
	if options.print ~= false then
		print(fmt("[Profiler] total %s (%d sections)", formatTime(total, precision), #entries))
		print(fmt("%-24s %12s %12s %12s %8s %12s %12s", "section", "total", "self", "avg", "count", "min", "max"))
		for _, entry in entries do
			local totalStr = formatTime(entry.total, precision)
			local selfStr = formatTime(entry.self, precision)
			local avgStr = formatTime(entry.avg, precision)
			local minStr = formatTime(entry.min, precision)
			local maxStr = formatTime(entry.max, precision)
			print(
				fmt(
					"%-24s %12s %12s %12s %8d %12s %12s (%5.1f%%)",
					entry.name,
					totalStr,
					selfStr,
					avgStr,
					entry.count,
					minStr,
					maxStr,
					entry.percent
				)
			)
		end
	end
	return entries, total
end

function Profiler.withEnabled(flag: boolean, fn: Types.Callable, ...: any): ...any
	assert(type(fn) == "function", "Profiler.withEnabled expects a function")
	local prev = enabled
	Profiler.setEnabled(flag)
	local packed = table.pack(pcall(fn, ...))
	Profiler.setEnabled(prev)
	if not packed[1] then
		error(packed[2], 0)
	end
	return table.unpack(packed, 2, packed.n)
end

Profiler.push = Profiler.start
Profiler.pop = Profiler.stop
Profiler.time = Profiler.scope

return Profiler


--!native
--!optimize 2
local RNG = {}
local _rng = Random.new()

function RNG.seed(s: number)
	_rng = Random.new(math.floor(s))
	return s
end

function RNG.uniform(a: number?, b: number?, rng: Random?)
	a = a or 0
	b = b or 1
	return (rng or _rng):NextNumber(a, b)
end

function RNG.randint(lo: number, hi: number, rng: Random?)
	-- inclusive
	return lo + math.floor((rng or _rng):NextNumber(0, 1) * (hi - lo + 1))
end

function RNG.generator()
	return {
		randint = function(_, lo: number, hi: number, rng: Random?): number
			return RNG.randint(lo, hi, rng)
		end,
		random = function(_, rng: Random?): number
			return RNG.uniform(0, 1, rng)
		end,
		uniform = function(_, a: number?, b: number?, rng: Random?): number
			return RNG.uniform(a, b, rng)
		end,
	}
end

return RNG


--!native
--!optimize 2
local function update(self, x: number)
	self.n += 1
	local d = x - self.mean
	self.mean += d / self.n
	self.m2 += d * (x - self.mean)
end

local function var(self): number
	if self.n < 2 then
		return 1.0
	end
	return self.m2 / (self.n - 1)
end

local function std(self): number
	return math.sqrt(var(self)) + self.eps
end

local function normalize(self, x: number): number
	return (x - self.mean) / std(self)
end

local function new(eps: number?)
	return {
		n = 0,
		mean = 0.0,
		m2 = 0.0,
		eps = eps or 1e-8,
		update = update,
		var = var,
		std = std,
		normalize = normalize,
	}
end

return new