#!/usr/bin/env python3
"""Two-stage pipeline figure: ONE pipeline, two instantiations (T2M-GPT / NSLP-G).

The two systems share an identical two-step structure, so it is drawn once:

  Step 1  learn a pose representation from SKELETONS ALONE (no text)
          motion -> encoder -> latent/codebook -> decoder -> motion
  Step 2  learn to emit that representation FROM TEXT, stage 1 frozen
          text -> text encoder -> generator -> predicted latents -> stage-1 decoder -> motion

Each shared box names what the two systems put there (blue = T2M-GPT, purple = NSLP-G),
and a strip at the bottom lists only the places they actually diverge.

Visual language follows figs/Capture.PNG: numbered badges, outlined panels, and the
latent sequence drawn as little squares. All real shapes, so it stays editable.
"""
import argparse
import os

from pptx import Presentation
from pptx.dml.color import RGBColor
from pptx.enum.dml import MSO_LINE_DASH_STYLE
from pptx.enum.shapes import MSO_CONNECTOR, MSO_SHAPE
from pptx.enum.text import MSO_ANCHOR, PP_ALIGN
from pptx.oxml.ns import qn
from pptx.util import Emu, Inches, Pt

INK = RGBColor(0x1F, 0x24, 0x2E)
MUTED = RGBColor(0x60, 0x6A, 0x78)
PANEL_LN = RGBColor(0x7E, 0xB0, 0xE0)
WHITE = RGBColor(0xFF, 0xFF, 0xFF)
BADGE = RGBColor(0x2E, 0x75, 0xB6)
BOX_BG, BOX_LN = RGBColor(0xDE, 0xEB, 0xF7), RGBColor(0x9D, 0xC3, 0xE6)
TOK_BG, TOK_LN = RGBColor(0xD3, 0xC2, 0xE8), RGBColor(0x9B, 0x7E, 0xC8)
IO_BG, IO_LN = RGBColor(0xF2, 0xF2, 0xF2), RGBColor(0xC8, 0xC8, 0xC8)
GREY = RGBColor(0xBF, 0xBF, 0xBF)
RED = RGBColor(0xC0, 0x00, 0x00)
T2MC = RGBColor(0x1A, 0x4F, 0x8A)
NSLC = RGBColor(0x6B, 0x2E, 0x9E)
HDR_BG = RGBColor(0xF7, 0xF7, 0xF9)


def _fill(shape, lines, align=PP_ALIGN.CENTER):
    tf = shape.text_frame
    tf.word_wrap = True
    tf.vertical_anchor = MSO_ANCHOR.MIDDLE
    tf.margin_left = tf.margin_right = Emu(18000)
    tf.margin_top = tf.margin_bottom = 0
    for i, (t, sz, bd, col) in enumerate(lines):
        p = tf.paragraphs[0] if i == 0 else tf.add_paragraph()
        p.alignment = align
        r = p.add_run(); r.text = t
        r.font.size = Pt(sz); r.font.bold = bd
        r.font.color.rgb = col; r.font.name = 'Calibri'


def rbox(sl, x, y, w, h, lines, fill, line, *, radius=0.16, lw=1.25,
         shape=MSO_SHAPE.ROUNDED_RECTANGLE, align=PP_ALIGN.CENTER):
    s = sl.shapes.add_shape(shape, Inches(x), Inches(y), Inches(w), Inches(h))
    s.fill.solid(); s.fill.fore_color.rgb = fill
    s.line.color.rgb = line; s.line.width = Pt(lw)
    s.shadow.inherit = False
    if shape == MSO_SHAPE.ROUNDED_RECTANGLE:
        try:
            s.adjustments[0] = radius
        except Exception:
            pass
    _fill(s, lines, align)
    return s


def text(sl, x, y, w, lines, *, align=PP_ALIGN.LEFT, h=0.28):
    tb = sl.shapes.add_textbox(Inches(x), Inches(y), Inches(w), Inches(h))
    tf = tb.text_frame; tf.word_wrap = True
    tf.margin_left = tf.margin_right = tf.margin_top = tf.margin_bottom = 0
    for i, (t, sz, bd, col) in enumerate(lines):
        p = tf.paragraphs[0] if i == 0 else tf.add_paragraph()
        p.alignment = align
        r = p.add_run(); r.text = t
        r.font.size = Pt(sz); r.font.bold = bd
        r.font.color.rgb = col; r.font.name = 'Calibri'
    return tb


def arrow(sl, x1, y, x2, *, color=RGBColor(0x44, 0x44, 0x44), w=1.6):
    c = sl.shapes.add_connector(MSO_CONNECTOR.STRAIGHT, Inches(x1), Inches(y),
                                Inches(x2), Inches(y))
    c.line.color.rgb = color; c.line.width = Pt(w)
    ln = c.line._get_or_add_ln()
    ln.append(ln.makeelement(qn('a:tailEnd'),
                             {'type': 'triangle', 'w': 'med', 'len': 'med'}))
    return c


def badge(sl, x, y, n, d=0.30):
    s = sl.shapes.add_shape(MSO_SHAPE.OVAL, Inches(x), Inches(y), Inches(d), Inches(d))
    s.fill.solid(); s.fill.fore_color.rgb = BADGE
    s.line.fill.background(); s.shadow.inherit = False
    _fill(s, [(str(n), 12, True, WHITE)])
    return s


def cells(sl, x, y, n, *, s=0.20, gap=0.055, rows=1, ell=True):
    for r in range(rows):
        for i in range(n):
            sh = sl.shapes.add_shape(MSO_SHAPE.ROUNDED_RECTANGLE,
                                     Inches(x + i * (s + gap)), Inches(y + r * (s + gap)),
                                     Inches(s), Inches(s))
            sh.fill.solid(); sh.fill.fore_color.rgb = TOK_BG
            sh.line.color.rgb = TOK_LN; sh.line.width = Pt(0.9)
            sh.shadow.inherit = False
            try:
                sh.adjustments[0] = 0.15
            except Exception:
                pass
    endx = x + n * (s + gap)
    if ell:
        text(sl, endx - 0.01, y + (rows - 1) * (s + gap) / 2 - 0.02, 0.42,
             [('· · ·', 13, True, MUTED)], h=0.26)
        endx += 0.40
    return endx


def stickman(sl, x, y, w=0.56, h=0.84, col=RGBColor(0x33, 0x33, 0x33)):
    hd = sl.shapes.add_shape(MSO_SHAPE.OVAL, Inches(x + w / 2 - 0.085), Inches(y),
                             Inches(0.17), Inches(0.17))
    hd.fill.background(); hd.line.color.rgb = col; hd.line.width = Pt(1.4)
    hd.shadow.inherit = False

    def seg(x1, y1, x2, y2):
        c = sl.shapes.add_connector(MSO_CONNECTOR.STRAIGHT, Inches(x1), Inches(y1),
                                    Inches(x2), Inches(y2))
        c.line.color.rgb = col; c.line.width = Pt(1.5)
    cx = x + w / 2
    seg(cx, y + 0.17, cx, y + 0.50)
    seg(cx, y + 0.24, x + 0.02, y + 0.13)
    seg(cx, y + 0.24, x + w - 0.02, y + 0.36)
    seg(cx, y + 0.50, x + 0.09, y + h)
    seg(cx, y + 0.50, x + w - 0.09, y + h)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', default='figs/two_step_pipeline.pptx')
    args = ap.parse_args()

    prs = Presentation()
    prs.slide_width = Inches(13.333)
    prs.slide_height = Inches(7.5)
    sl = prs.slides.add_slide(prs.slide_layouts[6])

    text(sl, 0, 0.16, 13.333, [('Two-Step Pipeline — one structure, two instantiations',
                                23, True, INK)], align=PP_ALIGN.CENTER, h=0.4)
    text(sl, 0, 0.60, 13.333,
         [('T2M-GPT and NSLP-G share the identical two-step design; only the boxes marked '
           'below are filled differently.', 11.5, False, MUTED)],
         align=PP_ALIGN.CENTER, h=0.3)

    AY = 1.02          # panel 1 top
    BY = 3.42          # panel 2 top
    PH = 2.10

    # =============================================== STEP 1
    rbox(sl, 0.30, AY, 12.73, PH, [], WHITE, PANEL_LN, radius=0.05, lw=1.5,
         shape=MSO_SHAPE.ROUNDED_RECTANGLE)
    badge(sl, 0.48, AY + 0.13, 1)
    text(sl, 0.88, AY + 0.15, 6.4,
         [('Step 1: Latent Representation', 13, True, BADGE)], h=0.3)
    text(sl, 0.88, AY + 0.42, 6.4,
         [('trained on skeletons alone — no text', 9.5, False, MUTED)], h=0.24)

    y = AY + 0.80          # row top
    mid = y + 0.42
    stickman(sl, 0.62, y + 0.02)
    text(sl, 0.40, y + 0.92, 1.0, [('Input Motion', 9, False, INK)],
         align=PP_ALIGN.CENTER, h=0.24)
    arrow(sl, 1.24, mid, 1.56)
    rbox(sl, 1.62, y + 0.02, 2.62, 0.82,
         [('Encoder', 11.5, True, INK), ('VQ-VAE  (T2M-GPT)', 8.5, False, T2MC),
          ('Spatial VAE  (NSLP-G)', 8.5, False, NSLC)], BOX_BG, BOX_LN)
    arrow(sl, 4.28, mid, 4.60)
    ex = cells(sl, 4.66, y + 0.14, 4, rows=2)
    text(sl, 4.45, y + 0.92, 1.9,
         [('Latent / Codebook', 9, False, INK), ('(compressed)', 8.5, False, MUTED)],
         align=PP_ALIGN.CENTER, h=0.4)
    arrow(sl, ex + 0.02, mid, ex + 0.34)
    dx = ex + 0.40
    rbox(sl, dx, y + 0.02, 2.42, 0.82,
         [('Decoder', 11.5, True, INK), ('reconstructs the pose', 8.5, False, MUTED)],
         BOX_BG, BOX_LN)
    arrow(sl, dx + 2.44, mid, dx + 2.76)
    stickman(sl, dx + 2.82, y + 0.02)
    text(sl, dx + 2.60, y + 0.92, 1.0, [('Motion', 9, False, INK)],
         align=PP_ALIGN.CENTER, h=0.24)

    text(sl, 10.46, y + 0.02, 2.70,
         [('T2M-GPT  ', 9.5, True, T2MC),
          ('discrete: 512-entry codebook,', 9, False, INK),
          ('4× temporal downsample', 9, False, INK),
          ('NSLP-G  ', 9.5, True, NSLC),
          ('continuous: one Gaussian latent', 9, False, INK),
          ('per frame, no time compression', 9, False, INK)], h=1.3)

    # -------------------- linking arrow --------------------
    a = sl.shapes.add_shape(MSO_SHAPE.DOWN_ARROW, Inches(6.42), Inches(AY + PH + 0.06),
                            Inches(0.50), Inches(0.26))
    a.fill.solid(); a.fill.fore_color.rgb = GREY
    a.line.fill.background(); a.shadow.inherit = False
    text(sl, 7.00, AY + PH + 0.05, 5.6,
         [('stage 1 is frozen; its decoder is reused below', 9, False, MUTED)], h=0.24)

    # =============================================== STEP 2
    rbox(sl, 0.30, BY, 12.73, PH, [], WHITE, PANEL_LN, radius=0.05, lw=1.5)
    badge(sl, 0.48, BY + 0.13, 2)
    text(sl, 0.88, BY + 0.15, 6.4, [('Step 2: Frame Generation', 13, True, BADGE)], h=0.3)
    text(sl, 0.88, BY + 0.42, 6.4,
         [('text → latent, with step 1 frozen', 9.5, False, MUTED)], h=0.24)

    y = BY + 0.80
    mid = y + 0.42
    rbox(sl, 0.50, y + 0.02, 0.98, 0.82,
         [('Gloss /', 10, True, INK), ('sentence', 10, True, INK)], IO_BG, IO_LN)
    arrow(sl, 1.50, mid, 1.80)
    rbox(sl, 1.86, y + 0.02, 1.16, 0.82,
         [('Text', 10.5, True, INK), ('encoder', 10.5, True, INK)], BOX_BG, BOX_LN)
    arrow(sl, 3.04, mid, 3.34)
    rbox(sl, 3.40, y + 0.02, 2.14, 0.82,
         [('Generator', 11.5, True, INK),
          ('Autoregressive GPT  (T2M-GPT)', 8.5, False, T2MC),
          ('Non-AR decoder  (NSLP-G)', 8.5, False, NSLC)], BOX_BG, BOX_LN)
    arrow(sl, 5.56, mid, 5.86)
    ex = cells(sl, 5.92, y + 0.14, 3, rows=2)
    text(sl, 5.66, y + 0.92, 2.0,
         [('Predicted latents', 9, False, INK), ('/ tokens', 8.5, False, MUTED)],
         align=PP_ALIGN.CENTER, h=0.4)
    arrow(sl, ex + 0.02, mid, ex + 0.32)
    dx = ex + 0.38
    rbox(sl, dx, y + 0.02, 1.62, 0.82,
         [('Step-1 Decoder', 10, True, INK), ('frozen', 8.5, False, MUTED)],
         BOX_BG, BOX_LN, lw=1.25)
    d = sl.shapes.add_shape(MSO_SHAPE.ROUNDED_RECTANGLE, Inches(dx), Inches(y + 0.02),
                            Inches(1.62), Inches(0.82))
    d.fill.background(); d.line.color.rgb = BADGE; d.line.width = Pt(1.1)
    d.line.dash_style = MSO_LINE_DASH_STYLE.DASH
    d.shadow.inherit = False
    arrow(sl, dx + 1.64, mid, dx + 1.94)
    stickman(sl, dx + 2.00, y + 0.02)
    text(sl, dx + 1.70, y + 0.92, 1.16, [('Output Motion', 9, False, INK)],
         align=PP_ALIGN.CENTER, h=0.24)

    text(sl, 10.46, y - 0.02, 2.70,
         [('T2M-GPT  ', 9.5, True, T2MC),
          ('one token at a time; sampling', 9, False, INK),
          ('from the codebook; ', 9, False, INK),
          ('Cross-Entropy', 9, True, RED),
          ('NSLP-G  ', 9.5, True, NSLC),
          ('all frames at once; regression', 9, False, INK),
          ('to the latent; ', 9, False, INK), ('MSE', 9, True, BADGE)], h=1.5)

    # =============================================== bottom strip
    SY = BY + PH + 0.18
    rows = [('', 'T2M-GPT', 'NSLP-G'),
            ('Step-1 representation', 'discrete codebook (512)', 'continuous Gaussian, per frame'),
            ('Time compression', '4× downsample', 'none'),
            ('Step-2 decoding', 'autoregressive + sampling', 'non-autoregressive, parallel'),
            ('Step-2 objective', 'Cross-Entropy', 'MSE')]
    x0, wl, w1, hh = 0.30, 3.10, 4.81, 0.27
    for i, (a_, b_, c_) in enumerate(rows):
        yy = SY + i * hh
        head = i == 0
        for (xx, ww, tt, al) in ((x0, wl, a_, PP_ALIGN.LEFT),
                                 (x0 + wl, w1, b_, PP_ALIGN.CENTER),
                                 (x0 + wl + w1, w1, c_, PP_ALIGN.CENTER)):
            col = T2MC if (head and tt == 'T2M-GPT') else NSLC if (head and tt == 'NSLP-G') else INK
            s = rbox(sl, xx, yy, ww, hh, [(tt, 9.5, head, col)],
                     HDR_BG if head else WHITE, IO_LN, lw=0.75,
                     shape=MSO_SHAPE.RECTANGLE, align=al)

    os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True)
    prs.save(args.out)
    print(f'wrote {args.out}')


if __name__ == '__main__':
    main()