"""Lab sinh ảnh 4×4 bằng flow matching; chỉ cần NumPy, không tải dữ liệu.
Tiền đề: biết MSE, quy tắc chuỗi và AdamW từ LLM Training (Text).
Đây là mạng truyền thẳng thu nhỏ, không phải DiT hay VAE đã huấn luyện.
"""
import numpy as np

SEED = 17
STEPS = 1800
LR = 0.003
DROP_PROB = 0.2
EMA_DECAY = 0.98
LORA_RANK = 2
LORA_STEPS = 450
SAMPLE_STEPS = 8
CFG_SCALE = 1.5
BATCH = 64
HIDDEN = 48
D = 16


def section(n, title, observe, conclusion, limit):
    print("\n[%d · %s]" % (n, title))
    print("Quan sát: " + observe)
    print("Kết luận: " + conclusion)
    print("Không chứng minh: " + limit)


def templates(thick=False):
    x = -np.ones((2, 4, 4))
    x[0, 1, :] = 1
    x[1, :, 1] = 1
    if thick:
        x[0, 2, :] = 1
        x[1, :, 2] = 1
    return x.reshape(2, D)


def batch(rng, n=BATCH, thick=False):
    labels = rng.integers(0, 2, n)
    clean = templates(thick)[labels] + rng.normal(0, 0.03, (n, D))
    noise = rng.normal(size=(n, D))
    t = rng.uniform(0, 1, (n, 1))
    c = np.eye(2)[labels]
    mask = rng.random((n, 1))
    return clean, noise, t, c, mask


def interpolate(b, drop=DROP_PROB):
    clean, noise, t, c, mask = b
    cond = c * (mask >= drop)
    zt = (1-t)*noise + t*clean  # t=0: nhiễu; t=1: ảnh sạch.
    return np.concatenate([zt, t, cond], axis=1), clean-noise


def init():
    rng = np.random.default_rng(SEED)
    return dict(W1=rng.normal(0, 0.16, (D+3, HIDDEN)), b1=np.zeros(HIDDEN),
                W2=rng.normal(0, 0.1, (HIDDEN, D)), b2=np.zeros(D))


def copy(p):
    return {k: v.copy() for k, v in p.items()}


def mm(a, b):
    # Tổng tích tường minh: tránh cảnh báo BLAS giả trên một số bản NumPy macOS.
    return np.einsum("ij,jk->ik", a, b, optimize=False)


def forward(p, x):
    h = np.tanh(mm(x, p['W1']) + p['b1'])
    return mm(h, p['W2']) + p['b2'], h


def loss_grad(p, x, y):
    pred, h = forward(p, x)
    delta = 2*(pred-y)/y.size  # MSE trung bình cả batch và 16 pixel.
    dh = (mm(delta, p['W2'].T))*(1-h*h)
    g = dict(W1=mm(x.T, dh), b1=dh.sum(axis=0), W2=mm(h.T, delta), b2=delta.sum(axis=0))
    return float(np.mean((pred-y)**2)), g


def adamw(p, g, m, v, step, lr=LR):
    for k in p:
        m[k] = 0.9*m[k] + 0.1*g[k]
        v[k] = 0.999*v[k] + 0.001*g[k]*g[k]
        update = (m[k]/(1-0.9**step))/(np.sqrt(v[k]/(1-0.999**step))+1e-8)
        decay = 0.01*p[k] if k.startswith('W') else 0
        p[k] -= lr*(update+decay)


def mse(p, x, y):
    return float(np.mean((forward(p, x)[0]-y)**2))


def train(drop):
    p = init()
    ema = copy(p)
    m = {k: np.zeros_like(v) for k, v in p.items()}
    v = copy(m)
    rng = np.random.default_rng(SEED+1)
    nulls = 0
    trace = []
    for step in range(1, STEPS+1):
        b = batch(rng)
        nulls += int((b[4] < drop).sum())
        x, y = interpolate(b, drop)
        loss, g = loss_grad(p, x, y)
        adamw(p, g, m, v, step)
        for k in p:
            ema[k] = EMA_DECAY*ema[k] + (1-EMA_DECAY)*p[k]
        if step in [1, STEPS//2, STEPS]:
            trace.append((step, loss))
    return p, ema, trace, nulls


def sample(p, initial, labels, steps=SAMPLE_STEPS, scale=CFG_SCALE):
    z = initial.copy()
    c = np.eye(2)[labels]
    for j in range(steps):
        t = np.full((len(z), 1), j/steps)
        v_c = forward(p, np.concatenate([z, t, c], axis=1))[0]
        v_u = forward(p, np.concatenate([z, t, np.zeros_like(c)], axis=1))[0]
        velocity = v_u + scale*(v_c-v_u)
        z += velocity/steps
    return z


def main():
    print("LLM Training (Image) · Lab NumPy · seed=%d" % SEED)
    print("Ảnh 4×4, hai nhãn 'ngang'/'dọc'; 16 số mỗi ảnh. Không tải mô hình hay ảnh thật.")
    print("MLP = mạng truyền thẳng: [16 pixel + 1 thời gian + 2 điều kiện] → %d → 16 vận tốc." % HIDDEN)
    print("Giữ kiến thức loss/gradient/AdamW; VAE và bộ mã hóa chữ ở đây là ánh xạ cố định, không được học.")
    print("Cách đọc số: MSE là trung bình bình phương sai số; RMS là căn của trung bình bình phương độ lệch giữa hai kết quả đang so.")
    print("Mọi số Quan sát được tính trong lần chạy này; loss thấp chưa nói ảnh đẹp hay mô hình hiểu tiếng Việt.")
    b = batch(np.random.default_rng(SEED+2), 128)
    clean, noise, t, cond, mask = b
    section(1, "Cặp ảnh và điều kiện", "ảnh=%s; điều kiện=%s; tổng pixel mỗi mẫu=%d; nhãn đầu=%s" %
            (clean.shape, cond.shape, D, 'ngang' if cond[0, 0] else 'dọc'),
            "Mỗi hàng điều kiện có hai cột: [1,0] chỉ nét ngang, [0,1] chỉ nét dọc; [0,0] là điều kiện rỗng.",
            "Bảng tra nhãn không phải CLIP/VLM; lab không dạy bộ mã hóa hiểu ảnh hay tiếng Việt.")
    identity_error = float(np.max(np.abs(clean.reshape(-1, 4, 4).reshape(-1, D)-clean)))
    x, y = interpolate(b)
    endpoint = max(float(np.max(np.abs((1-0)*noise+0*clean-noise))),
                   float(np.max(np.abs((1-1)*noise+1*clean-clean))))
    section(2, "Latent và đường nội suy", "ánh xạ thay VAE lệch=%.1e; hai đầu đường lệch=%.1e; t đầu=%.6f; z_t[0,0]=%.6f; u[0,0]=%.6f" %
            (identity_error, endpoint, t[0, 0], x[0, 0], y[0, 0]),
            "e là nhiễu Gaussian chuẩn (trung bình 0, độ lệch chuẩn 1); t=0 là nhiễu, t=1 là sạch. z_t=(1-t)e+t·z có đạo hàm u=z-e; ánh xạ đồng nhất giữ nguyên pixel.",
            "Phép reshape không nén và không phải VAE thật; encoder VAE thật cần được huấn luyện trước rồi đóng băng.")
    p0 = init()
    pred, h = forward(p0, x)
    nparam = sum(v.size for v in p0.values())
    section(3, "Forward denoiser", "đầu vào=%s; tầng ẩn=%s; đầu ra=%s; tham số=%d; vận tốc đầu=%.6f" %
            (x.shape, h.shape, pred.shape, nparam, pred[0, 0]),
            "Denoiser nhận ảnh nhiễu, thời gian và điều kiện cùng lúc, trả 16 vận tốc để dịch từng pixel.",
            "Mạng tanh này không có attention; chỉ thay vai trò denoiser của ô DiT, không kiểm chứng DiT.")
    initial_loss, g = loss_grad(p0, x, y)
    section(4, "MSE và gradient", "MSE ban đầu=%.6f; MSE dự đoán zero=%.6f; chuẩn gradient W1=%.6f" %
            (initial_loss, float(np.mean(y*y)), np.linalg.norm(g['W1'])),
            "Lấy trung bình bình phương sai số trên mọi mẫu và pixel; zero chỉ là baseline cùng mục tiêu vận tốc.",
            "MSE vận tốc không phải điểm thẩm mỹ; ε-prediction và v-prediction có mục tiêu và trọng số thời gian khác.")
    errors = []
    for k in p0:
        idx = np.unravel_index(np.argmax(np.abs(g[k])), g[k].shape)
        old = p0[k][idx]
        p0[k][idx] = old+1e-5
        plus = mse(p0, x, y)
        p0[k][idx] = old-1e-5
        minus = mse(p0, x, y)
        p0[k][idx] = old
        errors.append(abs((plus-minus)/2e-5-g[k][idx]))
    q = copy(p0)
    moments = {k: np.zeros_like(v) for k, v in q.items()}
    adamw(q, g, moments, copy(moments), 1)
    max_error = max(errors)
    section(5, "Backward và AdamW", "sai phân hữu hạn lớn nhất=%.3e; MSE sau một bước=%.6f; đổi W1 lớn nhất=%.6f" %
            (max_error, mse(q, x, y), np.max(np.abs(q['W1']-p0['W1']))),
            "Kiểm tra phần tử gradient lớn nhất trong từng tensor bằng dịch ±1e-5; AdamW dùng trung bình trượt và suy giảm trọng số tách riêng.",
            "Chỉ kiểm tra bốn phần tử, không phải chứng minh mọi gradient; một bước giảm không bảo đảm mọi bước đều giảm.")
    p, ema, trace, nulls = train(DROP_PROB)
    final_loss = mse(p, x, y)
    ema_loss = mse(ema, x, y)
    snapshot = copy(ema)
    reload_error = float(np.max(np.abs(forward(snapshot, x)[0]-forward(ema, x)[0])))
    section(6, "Vòng học và EMA", "trace (bước, loss minibatch)=%s; kiểm định cố định %.6f → %.6f; EMA=%.6f; nạp lại lệch=%.1e" %
            ([(s, round(l, 6)) for s, l in trace], initial_loss, final_loss, ema_loss, reload_error),
            "So trước/sau trên cùng 128 cặp ảnh–nhiễu–t; EMA là trung bình trọng số qua các bước, không phải trung bình các ảnh đầu ra.",
            "EMA không bảo đảm loss tốt hơn; checkpoint trong RAM chỉ mô phỏng weights-only, không lưu optimizer hay RNG để tiếp tục chính xác.")
    p_no, ema_no, _, zero_nulls = train(0)
    x_null, y_null = interpolate(b, 1)
    x_cond, _ = interpolate(b, 0)
    delta = float(np.sqrt(np.mean((forward(ema, x_cond)[0]-forward(ema, x_null)[0])**2)))
    null_with = mse(ema, x_null, y_null)
    null_without = mse(ema_no, x_null, y_null)
    section(7, "Bỏ điều kiện để học CFG", "null=%d/%d (%.4f); đối chứng p=0 có %d null; MSE null có dropout=%.6f, không dropout=%.6f; RMS có/rỗng=%.6f" %
            (nulls, BATCH*STEPS, nulls/(BATCH*STEPS), zero_nulls, null_with, null_without, delta),
            "Hai lần học dùng cùng khởi tạo và các batch. Dropout cho mô hình gặp [0,0], tạo nhánh không điều kiện dùng trong CFG (hướng dẫn không cần classifier).",
            "So sánh chỉ trên một seed; p=0 vẫn trả số cho null nhưng chưa được huấn luyện trên null. Không suy ra p tối ưu.")
    frozen = copy(ema)
    adapter_rng = np.random.default_rng(SEED+4)
    adapter = dict(A=adapter_rng.normal(0, 0.05, (HIDDEN, LORA_RANK)), B=np.zeros((LORA_RANK, D)))
    am = {k: np.zeros_like(v) for k, v in adapter.items()}
    av = copy(am)
    adapt_eval = batch(np.random.default_rng(SEED+5), 128, True)
    ax, ay = interpolate(adapt_eval, 0)
    baseline_adapt = mse(frozen, ax, ay)
    for step in range(1, LORA_STEPS+1):
        lx, ly = interpolate(batch(adapter_rng, thick=True), 0)
        base_pred, lh = forward(frozen, lx)
        residual = base_pred+mm(mm(lh, adapter['A']), adapter['B'])-ly
        dl = 2*residual/ly.size
        ag = dict(A=mm(lh.T, mm(dl, adapter['B'].T)), B=mm(mm(lh, adapter['A']).T, dl))
        adamw(adapter, ag, am, av, step, lr=0.006)
    merged = copy(frozen)
    merged['W2'] += mm(adapter['A'], adapter['B'])
    after_adapt = mse(merged, ax, ay)
    unmerged = forward(frozen, ax)[0]+mm(mm(forward(frozen, ax)[1], adapter['A']), adapter['B'])
    merge_error = float(np.max(np.abs(unmerged-forward(merged, ax)[0])))
    frozen_error = max(float(np.max(np.abs(frozen[k]-ema[k]))) for k in frozen)
    section(8, "LoRA trên denoiser", "hạng=%d; adapter=%d so với %d tham số gốc; MSE nét dày %.6f → %.6f; gốc lệch=%.1e; gộp lệch=%.3e" %
            (LORA_RANK, sum(a.size for a in adapter.values()), nparam, baseline_adapt, after_adapt, frozen_error, merge_error),
            "Chỉ A,B được cập nhật; hạng r là chiều trung gian của A·B, W2 mới = W2 + A·B. Dữ liệu mới có hai hàng/cột sáng, nên đây là thích nghi miền bằng ít tham số.",
            "Không phải thí nghiệm DreamBooth trên ảnh cá nhân; LoRA là cách chọn tham số, DreamBooth là cách tổ chức mục tiêu cá nhân hóa.")
    initial_noise = np.random.default_rng(SEED+6).normal(size=(2, D))
    labels = np.array([0, 1])
    generated = sample(snapshot, initial_noise, labels)
    no_cfg = sample(snapshot, initial_noise, labels, scale=1)
    null_generated = sample(snapshot, initial_noise, labels, scale=0)
    dense = sample(snapshot, initial_noise, labels, steps=64)
    pixels = generated.reshape(2, 4, 4)  # Bộ giải mã đồng nhất.
    section(9, "Checkpoint sang Euler và CFG", "bước=%d; scale=%.2f; MSE với hình mẫu=%.6f; scale=1: %.6f; scale=0: %.6f; RMS so 64 bước=%.6f" %
            (SAMPLE_STEPS, CFG_SCALE, np.mean((generated-templates())**2), np.mean((no_cfg-templates())**2),
             np.mean((null_generated-templates())**2), np.sqrt(np.mean((generated-dense)**2))),
            "Euler cộng Δt·v với Δt=1/số bước. CFG dùng v_rỗng + scale·(v_có-v_rỗng); scale=1 trả nhánh có điều kiện. Đây là sampler ít bước chưa chưng cất.",
            "Ít bước không đồng nghĩa distillation; chưa đo tốc độ GPU, FID hay chất lượng ảnh thật; không dùng vòng token/KV cache.")
    for label, image in zip(['ngang', 'dọc'], pixels):
        print("Ảnh lấy mẫu '%s', số thô (chưa cắt về [-1,1]):" % label)
        print(np.array2string(image, precision=3, suppress_small=True))
    weights_bytes = sum(a.nbytes for a in p.values())
    section(10, "Quy mô và ranh giới phương pháp", "trọng số float64=%d byte; cùng số phần tử float32=%d, float16=%d byte (tính từ dtype)" %
            (weights_bytes, nparam*4, nparam*2),
            "Đếm trọng số không gồm activation, gradient hay AdamW. Pretraining, fine-tuning, ControlNet, distillation và Diffusion-DPO đổi dữ liệu, nhánh học hoặc loss của vòng này.",
            "Lab không thực thi mixed precision, checkpointing, đổi độ phân giải, ControlNet, distillation hay DPO; đọc bài tương ứng để thiết kế thí nghiệm thật.")
    assert endpoint < 1e-12 and identity_error == 0
    assert max_error < 1e-7
    assert np.isfinite(final_loss) and final_loss < initial_loss
    assert reload_error == 0 and frozen_error == 0 and merge_error < 1e-12
    assert after_adapt < baseline_adapt
    assert np.isfinite(generated).all()
    print("PASS chỉ xác nhận các điều kiện vừa nêu trên dữ liệu hình học, không chứng nhận chất lượng sinh ảnh hay hiệu năng.")
    print("\nPASS · Hai đầu nội suy đúng; 4 gradient khớp sai phân <1e-7; loss kiểm định giảm; EMA nạp lại đúng; LoRA giảm loss và giữ gốc; gộp lệch <1e-12; Euler hữu hạn.")


if __name__ == '__main__':
    main()
