math.randomseed(os.time())

-- Функции ANSI ESC для управления терминалом
local function clear()
  io.write("\27[2J")     -- очистить весь экран
  io.write("\27[H")      -- переместить курсор в левый верхний угол
end

local function gotoXY(x, y)
  io.write(string.format("\27[%d;%dH", y, x))
end

local function clearLine()
  io.write("\27[2K") -- очистить текущую строку
end

-- Гиперпараметры
local N = 53       -- размер словаря чисел (0..52)
local DIM = 16
local HIDDEN = 32
local LR = 0.01
local STEPS = 30000

-- Создаем тензор случайных чисел
local function randomMatrix(rows, cols)
  local m = {}
  for i=1, rows do
    m[i] = {}
    for j=1, cols do
      m[i][j] = (math.random() - 0.5) * 2
    end
  end
  return m
end

-- Ядро нейросети (все как раньше) —
-- ниже функции dot, relu, softmax, forward, backward, trainStep, predict —
-- для краткости не дублирую весь код здесь, используйте ваш.
-- Важно, что их интерфейс остаётся прежним.

-- Для полноты добавляю минимальный вариант (копируйте в ваш скрипт):

local function zeros(n)
  local t = {}
  for i=1,n do t[i]=0 end
  return t
end

local function dot(a,b)
  local s = 0
  for i=1,#a do s = s + a[i]*b[i] end
  return s
end

local function relu(x)
  local r = {}
  for i=1,#x do r[i] = math.max(0, x[i]) end
  return r
end

local function relu_derivative(x)
  local r = {}
  for i=1,#x do r[i] = x[i] > 0 and 1 or 0 end
  return r
end

local function softmax(x)
  local maxv = math.max(table.unpack(x))
  local exps = {}
  local sum = 0
  for i=1,#x do
    exps[i] = math.exp(x[i] - maxv)
    sum = sum + exps[i]
  end
  local r = {}
  for i=1,#x do r[i] = exps[i]/sum end
  return r
end

local function crossEntropyLoss(probs, target_idx)
  return -math.log(probs[target_idx] or 1e-15)
end

local function gradientCrossEntropy(probs, target_idx)
  local grad = {}
  for i=1,#probs do grad[i] = probs[i] end
  grad[target_idx] = grad[target_idx] - 1
  return grad
end

local function concat(a,b)
  local r = {}
  for i=1,#a do r[#r+1] = a[i] end
  for i=1,#b do r[#r+1] = b[i] end
  return r
end

-- Инициализация весов и эмбеддингов
local emb = randomMatrix(N, DIM)
local w1 = randomMatrix(2*DIM, HIDDEN)
local b1 = {}
for i=1,HIDDEN do b1[i] = 0 end
local w2 = randomMatrix(HIDDEN, N)
local b2 = {}
for i=1,N do b2[i] = 0 end

-- forward, backward, trainStep, predict (используйте ваши реализации из предыдущих шагов)

local function forward(a_idx, b_idx)
  local emb_a = emb[a_idx + 1]
  local emb_b = emb[b_idx + 1]
  local input = concat(emb_a, emb_b)

  local h_raw = {}
  for i=1,HIDDEN do
    local sum = b1[i]
    for j=1, 2*DIM do
      sum = sum + w1[j][i] * input[j]
    end
    h_raw[i] = sum
  end
  local h = relu(h_raw)

  local out_raw = {}
  for i=1,N do
    local sum = b2[i]
    for j=1,HIDDEN do
      sum = sum + w2[j][i] * h[j]
    end
    out_raw[i] = sum
  end
  local out = softmax(out_raw)

  return {
    input = input,
    emb_a = emb_a,
    emb_b = emb_b,
    h_raw = h_raw,
    h = h,
    out_raw = out_raw,
    out = out
  }
end

local function backward(cache, a_idx, b_idx, target_idx)
  -- аналогично вашему коду, обновляет веса
  local out = cache.out
  local h = cache.h
  local h_raw = cache.h_raw
  local input = cache.input

  local grad_out = gradientCrossEntropy(out, target_idx + 1)

  local grad_w2 = {}
  for j=1,HIDDEN do
    grad_w2[j] = {}
    for i=1,N do
      grad_w2[j][i] = grad_out[i] * h[j]
    end
  end
  local grad_b2 = grad_out

  local grad_h = zeros(HIDDEN)
  for j=1,HIDDEN do
    local s = 0
    for i=1,N do
      s = s + grad_out[i] * w2[j][i]
    end
    grad_h[j] = s
  end

  local relu_der = relu_derivative(h_raw)
  for j=1,HIDDEN do
    grad_h[j] = grad_h[j] * relu_der[j]
  end

  local grad_w1 = {}
  for i=1, 2*DIM do
    grad_w1[i] = {}
    for j=1,HIDDEN do
      grad_w1[i][j] = grad_h[j] * input[i]
    end
  end
  local grad_b1 = grad_h

  local grad_input = zeros(2*DIM)
  for i=1, 2*DIM do
    local s = 0
    for j=1,HIDDEN do
      s = s + grad_h[j] * w1[i][j]
    end
    grad_input[i] = s
  end

  for i=1, 2*DIM do
    for j=1,HIDDEN do
      w1[i][j] = w1[i][j] - LR * grad_w1[i][j]
    end
  end
  for j=1,HIDDEN do
    b1[j] = b1[j] - LR * grad_b1[j]
  end

  for j=1,HIDDEN do
    for i=1,N do
      w2[j][i] = w2[j][i] - LR * grad_w2[j][i]
    end
  end
  for i=1,N do
    b2[i] = b2[i] - LR * grad_b2[i]
  end

  for i=1,DIM do
    emb[a_idx + 1][i] = emb[a_idx + 1][i] - LR * grad_input[i]
    emb[b_idx + 1][i] = emb[b_idx + 1][i] - LR * grad_input[DIM + i]
  end
end

local function trainStep(a_idx, b_idx)
  local cache = forward(a_idx, b_idx)
  local target = (a_idx + b_idx) % N
  local loss = crossEntropyLoss(cache.out, target + 1)
  backward(cache, a_idx, b_idx, target)
  return loss
end

local function predict(a_idx, b_idx)
  local cache = forward(a_idx, b_idx)
  local out = cache.out
  local max_prob = -1
  local max_idx = -1
  for i=1,#out do
    if out[i] > max_prob then
      max_prob = out[i]
      max_idx = i-1
    end
  end
  return max_idx, max_prob
end

-- Визуализация эмбеддингов и график точности

local function meanCenter(matrix)
  local n = #matrix
  local dim = #matrix[1]
  local means = {}
  for d=1, dim do
    means[d] = 0
    for i=1, n do means[d] = means[d] + matrix[i][d] end
    means[d] = means[d] / n
  end
  local centered = {}
  for i=1, n do
    centered[i] = {}
    for d=1, dim do
      centered[i][d] = matrix[i][d] - means[d]
    end
  end
  return centered
end

local function project2D(matrix)
  local centered = meanCenter(matrix)
  local points = {}
  for i=1, #centered do
    points[i] = {x = centered[i][1], y = centered[i][2]}
  end
  return points
end

local function drawPoints(points)
  local width, height = 60, 20
  local xmin, xmax, ymin, ymax = points[1].x, points[1].x, points[1].y, points[1].y
  for _, p in ipairs(points) do
    if p.x < xmin then xmin = p.x end
    if p.x > xmax then xmax = p.x end
    if p.y < ymin then ymin = p.y end
    if p.y > ymax then ymax = p.y end
  end
  local scaleX = (xmax - xmin) > 0 and (width-1)/(xmax - xmin) or 1
  local scaleY = (ymax - ymin) > 0 and (height-1)/(ymax - ymin) or 1
  local grid = {}
  for y=1,height do
    grid[y] = {}
    for x=1,width do
      grid[y][x] = " "
    end
  end
  for _, p in ipairs(points) do
    local x = math.floor((p.x - xmin)*scaleX) + 1
    local y = height - math.floor((p.y - ymin)*scaleY)
    if x >= 1 and x <= width and y >= 1 and y <= height then
      grid[y][x] = "*"
    end
  end

  -- Очистка области и вывод
  gotoXY(1, 5)
  for y=1,height do
    clearLine()
    io.write(table.concat(grid[y]).."\n")
  end
end

local accHistory = {}

local function drawAccuracy()
  local width, height = 60, 10
  local len = #accHistory
  local start = len > width and (len - width + 1) or 1
  local slice = {}
  for i=start,len do
    slice[#slice+1] = accHistory[i]
  end
  local grid = {}
  for y=1,height do
    grid[y] = {}
    for x=1,#slice do
      grid[y][x] = " "
    end
  end
  for x=1,#slice do
    local value = slice[x]
    local barHeight = math.floor(value * height)
    for y=height, height - barHeight + 1, -1 do
      grid[y][x] = "#"
    end
  end

  gotoXY(1, 26)
  for y=1,height do
    clearLine()
    io.write(table.concat(grid[y]).."\n")
  end
  gotoXY(1, 36)
  clearLine()
  io.write(string.format("Точность сейчас: %.2f%%\n", (slice[#slice] or 0) * 100))
end

local function printProgress(step, maxSteps, loss, accuracy)
  gotoXY(1,1)
  clearLine()
  local barWidth = 30
  local progress = math.floor((step / maxSteps) * barWidth)
  local bar = "[" .. string.rep("=", progress) .. string.rep(" ", barWidth - progress) .. "]"
  io.write(string.format("Шаг %6d/%d %s  Потеря: %.4f  Точность: %.2f%%\n", step, maxSteps, bar, loss, accuracy * 100))
  io.flush()
end

local function computeAccuracy(samples)
  local correct = 0
  for _, sample in ipairs(samples) do
    local a,b,target = sample[1], sample[2], sample[3]
    local pred,_ = predict(a, b)
    if pred == target then correct = correct + 1 end
  end
  return correct / #samples
end

local function generateValidationSet(size)
  local valSet = {}
  for i=1,size do
    local a, b = math.random(0, N-1), math.random(0, N-1)
    table.insert(valSet, {a,b,(a+b)%N})
  end
  return valSet
end

-- Очистить экран перед запуском
clear()

print("Старт обучения...")
local valSet = generateValidationSet(500)
for step=1, STEPS do
  local a_idx = math.random(0, N-1)
  local b_idx = math.random(0, N-1)
  local loss = trainStep(a_idx, b_idx)

  if step % 100 == 0 then
    local acc = computeAccuracy(valSet)
    printProgress(step, STEPS, loss, acc)
    os.sleep(0.01)
  end

  if step % 1000 == 0 then
    accHistory[#accHistory + 1] = computeAccuracy(valSet)
    local points2D = project2D(emb)
    drawPoints(points2D)
    drawAccuracy()
    os.sleep(0.05)
  end
end

gotoXY(1, 40)
print("\nОбучение завершено.\nПримеры предсказаний после обучения:")
for a=0,10 do
  for b=0,10 do
    local pred, prob = predict(a,b)
    local target = (a+b) % N
    local res = (pred == target) and "OK" or "ERR"
    io.write(string.format("%2d+%2d=%2d(%s), p=%.2f  ", a,b,pred,res,prob))
  end
  print()
end