File size: 59,569 Bytes
bf8ae1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33080a1
 
 
 
 
bf8ae1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33080a1
 
 
0dcf39e
33080a1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0dcf39e
 
 
 
 
33080a1
 
 
 
 
 
 
 
 
bf8ae1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33080a1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bf8ae1c
 
 
 
 
 
 
 
 
 
 
 
0dcf39e
 
 
 
 
 
 
 
 
 
 
bf8ae1c
 
 
 
 
 
 
0dcf39e
 
bf8ae1c
 
 
 
 
 
 
 
 
33080a1
 
 
bf8ae1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
"""AIFlow Math Ink 0.6의 shared trajectory와 raster→virtual-stroke 신경망이다."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import time
from typing import Any, Sequence

import numpy as np
import torch
from PIL import Image
from torch import Tensor, nn

from .ink06_canonical import FEATURE_NAMES_06, MAX_EVENTS, canonicalize_ink06


class ResidualTcnBlock06(nn.Module):
    """필요 변수: channel·dilation. 작동 원리: 모바일 호환 Conv1d residual로 타점 패턴을 인코딩한다."""

    def __init__(self, channels: int, dilation: int) -> None:
        super().__init__()
        groups = 8 if channels % 8 == 0 else 1
        self.network = nn.Sequential(
            nn.Conv1d(channels, channels, 5, padding=2 * dilation, dilation=dilation),
            nn.GroupNorm(groups, channels), nn.GELU(), nn.Dropout(0.10),
            nn.Conv1d(channels, channels, 1), nn.GroupNorm(groups, channels),
        )
        self.activation = nn.GELU()

    def forward(self, value: Tensor) -> Tensor:
        """필요 변수: B×C×T. 작동 원리: 동일 길이 residual feature를 반환한다."""

        return self.activation(value + self.network(value))


class SharedTrajectoryEncoder06(nn.Module):
    """필요 변수: 19채널·hidden. 작동 원리: padding을 제외한 attention 통계로 shared embedding을 만든다."""

    def __init__(self, input_size: int = len(FEATURE_NAMES_06), hidden_size: int = 128) -> None:
        super().__init__()
        groups = 8 if hidden_size % 8 == 0 else 1
        self.input_projection = nn.Sequential(
            nn.Conv1d(input_size, hidden_size, 1), nn.GroupNorm(groups, hidden_size), nn.GELU(),
        )
        self.blocks = nn.Sequential(*(ResidualTcnBlock06(hidden_size, dilation) for dilation in (1, 2, 4, 8)))
        self.attention = nn.Conv1d(hidden_size, 1, 1)

    def forward(self, sequence: Tensor) -> Tensor:
        """필요 변수: B×128×19. 작동 원리: stroke_progress=-1 padding을 attention/통계에서 제거한다."""

        mask = sequence[:, :, 8] >= 0
        encoded = self.blocks(self.input_projection(sequence.transpose(1, 2)))
        attention = self.attention(encoded).masked_fill(~mask.unsqueeze(1), -1e4)
        weights = attention.softmax(dim=2)
        mean = (encoded * weights).sum(dim=2)
        variance = ((encoded - mean.unsqueeze(2)).square() * weights).sum(dim=2)
        maximum = encoded.masked_fill(~mask.unsqueeze(1), -1e4).amax(dim=2)
        return torch.cat((mean, maximum, torch.sqrt(variance.clamp_min(1e-6))), dim=1)


class VirtualTrajectoryAdapter06(nn.Module):
    """필요 변수: virtual 19채널 feature. 작동 원리: 온라인 계약 채널을 보존하며 raster 전용 residual 보정을 학습한다."""

    def __init__(self, channels: int = len(FEATURE_NAMES_06), hidden_size: int = 48) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Conv1d(channels, hidden_size, 1), nn.GELU(),
            nn.Conv1d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size), nn.GELU(),
            nn.Conv1d(hidden_size, channels, 1),
        )
        nn.init.zeros_(self.network[-1].weight)
        nn.init.zeros_(self.network[-1].bias)
        # pen-up/progress/missing/source 계약은 adapter가 바꾸지 않고 관측 feature만 보정한다.
        mutable = torch.ones(channels)
        mutable[[7, 8, 17, 18]] = 0.0
        self.register_buffer("mutable_channels", mutable.view(1, 1, channels), persistent=False)

    def forward(self, sequence: Tensor) -> Tensor:
        """필요 변수: B×128×19 feature. 작동 원리: zero-init residual을 허용 채널에만 더한다."""

        delta = self.network(sequence.transpose(1, 2)).transpose(1, 2)
        return sequence + delta * self.mutable_channels


class DepthwiseRasterEncoder06(nn.Module):
    """필요 변수: 128×128 grayscale. 작동 원리: depthwise CNN의 8×8 공간 배치를 보존해 vectorizer에 전달한다."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        channels = (16, 32, 64, hidden_size)
        layers: list[nn.Module] = [nn.Conv2d(1, channels[0], 3, stride=2, padding=1), nn.GELU()]
        for source, target in zip(channels, channels[1:]):
            layers.extend([
                nn.Conv2d(source, source, 3, stride=2, padding=1, groups=source),
                nn.Conv2d(source, target, 1), nn.GroupNorm(8 if target % 8 == 0 else 1, target), nn.GELU(),
            ])
        self.network = nn.Sequential(*layers)
        self.spatial_projection = nn.Sequential(
            nn.Flatten(), nn.Linear(hidden_size * 8 * 8, hidden_size), nn.LayerNorm(hidden_size), nn.GELU(),
        )
        self.position_projection = nn.Linear(2, hidden_size, bias=False)
        axis = torch.linspace(-1.0, 1.0, 8)
        grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
        self.register_buffer("spatial_positions", torch.stack((grid_x, grid_y), dim=-1).view(64, 2), persistent=False)
        self.fine_projection = nn.Conv2d(64, hidden_size, 1)
        fine_axis = torch.linspace(-1.0, 1.0, 16)
        fine_y, fine_x = torch.meshgrid(fine_axis, fine_axis, indexing="ij")
        self.register_buffer(
            "fine_positions", torch.stack((fine_x, fine_y), dim=-1).view(256, 2), persistent=False,
        )
        self.pointer_projection = nn.Conv2d(32, hidden_size, 1)
        pointer_axis = (torch.arange(32, dtype=torch.float32) + 0.5) / 32.0
        pointer_y, pointer_x = torch.meshgrid(pointer_axis, pointer_axis, indexing="ij")
        self.register_buffer(
            "pointer_positions", torch.stack((pointer_x, pointer_y), dim=-1).view(1024, 2), persistent=False,
        )

    def forward(
        self, raster: Tensor, *, fine_tokens: bool = False, pointer_tokens: bool = False,
    ) -> tuple[Tensor, Tensor]:
        """필요 변수: B×1×128×128·해상도 선택. 작동 원리: 전역 요약과 8/16/32-grid 위치 token을 반환한다."""

        feature = raster
        fine_feature = None
        pointer_feature = None
        for index, layer in enumerate(self.network):
            feature = layer(feature)
            if index == 5:
                pointer_feature = feature
            if index == 9:
                fine_feature = feature
        if pointer_tokens:
            if pointer_feature is None:
                raise RuntimeError("32×32 raster feature가 생성되지 않았습니다.")
            tokens = self.pointer_projection(pointer_feature).flatten(2).transpose(1, 2)
            positions = self.pointer_positions * 2.0 - 1.0
            tokens = tokens + self.position_projection(positions).unsqueeze(0)
        elif fine_tokens:
            if fine_feature is None:
                raise RuntimeError("16×16 raster feature가 생성되지 않았습니다.")
            tokens = self.fine_projection(fine_feature).flatten(2).transpose(1, 2)
            tokens = tokens + self.position_projection(self.fine_positions).unsqueeze(0)
        else:
            tokens = feature.flatten(2).transpose(1, 2)
            tokens = tokens + self.position_projection(self.spatial_positions).unsqueeze(0)
        return self.spatial_projection(feature), tokens


class RasterCrossAttentionBlock06(nn.Module):
    """필요 변수: trajectory query·4×4 raster token. 작동 원리: 각 가상 타점이 대응할 이미지 위치를 직접 조회한다."""

    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        heads = 4 if hidden_size % 4 == 0 else 1
        self.query_norm = nn.LayerNorm(hidden_size)
        self.memory_norm = nn.LayerNorm(hidden_size)
        self.attention = nn.MultiheadAttention(hidden_size, heads, batch_first=True)
        self.output_norm = nn.LayerNorm(hidden_size)
        self.residual_gate = nn.Parameter(torch.zeros(()))

    def forward(self, query: Tensor, memory: Tensor, *, gated: bool = False) -> Tensor:
        """필요 변수: query·공간 memory·gate 여부. 작동 원리: 위치 증거를 직접 또는 zero-init residual로 합친다."""

        attended, _weights = self.attention(
            self.query_norm(query), self.memory_norm(memory), self.memory_norm(memory), need_weights=False,
        )
        if gated:
            return query + torch.tanh(self.residual_gate) * attended
        return self.output_norm(query + attended)


class VirtualStrokeDecoder06(nn.Module):
    """필요 변수: raster embedding·가설 수. 작동 원리: 4-layer causal Conv1d가 top-k 좌표와 pen state를 만든다."""

    def __init__(self, hidden_size: int = 128, hypotheses: int = 4, max_events: int = MAX_EVENTS) -> None:
        super().__init__()
        self.hypotheses = hypotheses
        self.max_events = max_events
        self.query = nn.Parameter(torch.randn(max_events, hidden_size) * 0.02)
        self.hypothesis = nn.Embedding(hypotheses, hidden_size)
        self.decoder = nn.ModuleList([
            nn.Sequential(
                nn.Conv1d(hidden_size, hidden_size, kernel_size=5),
                nn.GroupNorm(8 if hidden_size % 8 == 0 else 1, hidden_size), nn.GELU(),
            )
            for _ in range(4)
        ])
        self.cross_attention = RasterCrossAttentionBlock06(hidden_size)
        self.coordinate_head = nn.Linear(hidden_size, 2)
        self.state_head = nn.Linear(hidden_size, 3)
        self.progress_head = nn.Linear(hidden_size, 1)
        self.score_head = nn.Linear(hidden_size, 1)
        self.pointer_query = nn.Linear(hidden_size, hidden_size, bias=False)
        self.pointer_key = nn.Linear(hidden_size, hidden_size, bias=False)
        self.pointer_temperature = 0.5
        self.pointer_logits_for_loss: Tensor | None = None

    def forward(
        self, embedding: Tensor, spatial_tokens: Tensor | None = None, *, gated_attention: bool = False,
        pointer_positions: Tensor | None = None, ink_prior: Tensor | None = None,
    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        """필요 변수: B×H 요약·선택 spatial token/ink prior. 작동 원리: causal path와 선택적 ink-pointer로 top-4 궤적을 반환한다."""

        batch = embedding.shape[0]
        self.pointer_logits_for_loss = None
        query = self.query.view(1, 1, self.max_events, -1)
        hypothesis = self.hypothesis.weight.view(1, self.hypotheses, 1, -1)
        value = query + hypothesis + embedding.view(batch, 1, 1, -1)
        value = value.reshape(batch * self.hypotheses, self.max_events, -1).transpose(1, 2)
        memory = None
        if spatial_tokens is not None:
            memory = spatial_tokens.unsqueeze(1).expand(-1, self.hypotheses, -1, -1)
            memory = memory.reshape(batch * self.hypotheses, spatial_tokens.shape[1], spatial_tokens.shape[2])
            # Pointer mode는 아래 좌표 head 자체가 memory attention이므로 중복 MHA를 만들지 않는다.
            if pointer_positions is None:
                value = self.cross_attention(
                    value.transpose(1, 2), memory, gated=gated_attention,
                ).transpose(1, 2)
        for layer in self.decoder:
            value = value + layer(nn.functional.pad(value, (4, 0)))
        decoded = value.transpose(1, 2)
        if pointer_positions is not None:
            if memory is None or ink_prior is None:
                raise ValueError("ink pointer에는 spatial memory와 ink prior가 모두 필요합니다.")
            if pointer_positions.shape != (memory.shape[1], 2) or ink_prior.shape != (batch, memory.shape[1]):
                raise ValueError("ink pointer position/prior shape가 spatial token과 일치하지 않습니다.")
            query = self.pointer_query(decoded)
            key = self.pointer_key(memory)
            pointer_logits = torch.bmm(query, key.transpose(1, 2)) / (decoded.shape[-1] ** 0.5)
            expanded_prior = ink_prior[:, None].expand(-1, self.hypotheses, -1).reshape(
                batch * self.hypotheses, memory.shape[1],
            )
            # 빈 배경은 확률상 허용하되 강하게 억제해 모든 좌표가 관측 ink 주변에서만 학습되게 한다.
            pointer_logits = pointer_logits + 2.5 * (expanded_prior + 1e-4).log().unsqueeze(1)
            self.pointer_logits_for_loss = pointer_logits.view(
                batch, self.hypotheses, self.max_events, memory.shape[1],
            )
            soft_probability = (pointer_logits / self.pointer_temperature).softmax(dim=-1)
            hard_probability = nn.functional.one_hot(
                soft_probability.argmax(dim=-1), num_classes=soft_probability.shape[-1],
            ).to(dtype=soft_probability.dtype)
            # Forward는 실제 ink cell 하나만 선택하고 backward는 soft distribution gradient를 사용한다.
            pointer_probability = soft_probability if self.training else hard_probability
            coordinates = torch.matmul(pointer_probability, pointer_positions.to(decoded)).view(
                batch, self.hypotheses, self.max_events, 2,
            )
        else:
            coordinates = self.coordinate_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events, 2)
        states = self.state_head(decoded).view(batch, self.hypotheses, self.max_events, 3)
        progress = self.progress_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events)
        scores = self.score_head(decoded[:, -1]).view(batch, self.hypotheses)
        return coordinates, states, progress, scores


def virtual_features06(
    coordinates: Tensor, state_logits: Tensor, stroke_progress: Tensor | None = None,
    *, contract: str = "legacy_v1",
) -> Tensor:
    """필요 변수: 좌표·state·progress·계약. 작동 원리: virtual stroke를 19채널 shared encoder 입력으로 변환한다."""

    if contract not in {"legacy_v1", "canonical_v2"}:
        raise ValueError("지원하지 않는 virtual feature contract입니다.")

    batch, hypotheses, steps, _axis = coordinates.shape
    probability = state_logits.softmax(dim=-1)
    pen_start = probability[..., 1]
    minimum = coordinates.amin(dim=2, keepdim=True)
    span = (coordinates.amax(dim=2, keepdim=True) - minimum).clamp_min(1e-8 if contract == "canonical_v2" else 1e-5)
    shape = (coordinates - minimum) / span
    canvas_delta = torch.cat((torch.zeros_like(coordinates[:, :, :1]), coordinates[:, :, 1:] - coordinates[:, :, :-1]), dim=2)
    shape_delta = torch.cat((torch.zeros_like(shape[:, :, :1]), shape[:, :, 1:] - shape[:, :, :-1]), dim=2)
    delta = shape_delta if contract == "canonical_v2" else canvas_delta
    delta = delta * (1.0 - pen_start).unsqueeze(-1)
    distance = delta.square().sum(dim=-1, keepdim=True).clamp_min(1e-8).sqrt()
    direction = delta / distance
    previous = torch.cat((torch.zeros_like(direction[:, :, :1]), direction[:, :, :-1]), dim=2)
    curvature = previous[..., 0] * direction[..., 1] - previous[..., 1] * direction[..., 0]
    progress = stroke_progress
    if progress is None:
        progress = torch.linspace(0.0, 1.0, steps, device=coordinates.device).view(1, 1, steps).expand(batch, hypotheses, -1)
    aspect = (span[..., 0] / span[..., 1]).expand(-1, -1, steps)
    ones = torch.ones_like(progress)
    bbox_top = minimum[..., 1].expand(-1, -1, steps)
    bbox_bottom = (minimum[..., 1] + span[..., 1]).expand(-1, -1, steps)
    bbox_height = span[..., 1].expand(-1, -1, steps)
    center_y = ((bbox_top + bbox_bottom) * 0.5)
    if contract == "canonical_v2":
        canvas_distance = canvas_delta.square().sum(dim=-1).sqrt() * 128.0
        time_delta = canvas_distance / (8.0 * 6.0)
        speed = torch.where(canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0), torch.zeros_like(canvas_distance))
    else:
        time_delta = (1.0 / (6.0 * steps)) * ones
        speed = canvas_delta.square().sum(dim=-1).clamp_min(1e-8).sqrt() * 6.0
    features = torch.stack((
        shape[..., 0], shape[..., 1], coordinates[..., 0], coordinates[..., 1],
        direction[..., 0], direction[..., 1], curvature, pen_start, progress, aspect,
        bbox_top, bbox_bottom, bbox_height, center_y, ones, time_delta,
        speed, ones, ones,
    ), dim=-1)
    if contract == "legacy_v1":
        valid = 1.0 - probability[..., 2]
        features[..., 8] = torch.where(valid > 0.5, features[..., 8], -torch.ones_like(features[..., 8]))
    return features


def equivalent_trajectory_targets06(
    coordinates: Tensor, states: Tensor, hypotheses: int = 4,
) -> tuple[Tensor, Tensor, Tensor]:
    """필요 변수: B×T 좌표·state. 작동 원리: 같은 raster를 만드는 방향/획순서 대안 trajectory를 생성한다."""

    if hypotheses != 4:
        raise ValueError("현재 equivalent target 계약은 top-4 전용입니다.")
    coordinate_batches: list[Tensor] = []
    state_batches: list[Tensor] = []
    progress_batches: list[Tensor] = []
    for sample_coordinates, sample_states in zip(coordinates, states, strict=True):
        starts = torch.nonzero(sample_states == 1, as_tuple=False).flatten().tolist()
        if not starts or starts[0] != 0:
            starts.insert(0, 0)
        starts = sorted(set(int(value) for value in starts if int(value) < len(sample_states)))
        boundaries = starts + [len(sample_states)]
        strokes = [sample_coordinates[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))]
        variants = (
            strokes,
            [stroke.flip(0) for stroke in strokes],
            list(reversed(strokes)),
            [stroke.flip(0) for stroke in reversed(strokes)],
        )
        sample_coordinate_targets = []
        sample_state_targets = []
        sample_progress_targets = []
        for variant in variants:
            joined = torch.cat(variant, dim=0)
            target_states = torch.zeros(len(joined), dtype=states.dtype, device=states.device)
            target_progress = torch.zeros(len(joined), dtype=coordinates.dtype, device=coordinates.device)
            cursor = 0
            for stroke in variant:
                target_states[cursor] = 1
                target_progress[cursor:cursor + len(stroke)] = torch.linspace(
                    0.0, 1.0, len(stroke), dtype=coordinates.dtype, device=coordinates.device,
                )
                cursor += len(stroke)
            target_states[-1] = 2
            sample_coordinate_targets.append(joined)
            sample_state_targets.append(target_states)
            sample_progress_targets.append(target_progress)
        coordinate_batches.append(torch.stack(sample_coordinate_targets))
        state_batches.append(torch.stack(sample_state_targets))
        progress_batches.append(torch.stack(sample_progress_targets))
    return torch.stack(coordinate_batches), torch.stack(state_batches), torch.stack(progress_batches)


def equivalent_modality_features06(sequence: Tensor) -> Tensor:
    """필요 변수: B×128×19 online feature. 작동 원리: raster 모드용 방향/획순서 불변 variant 네 개를 만든다."""

    batches: list[Tensor] = []
    for sample in sequence:
        starts = torch.nonzero(sample[:, 7] > 0.5, as_tuple=False).flatten().tolist()
        if not starts or starts[0] != 0:
            starts.insert(0, 0)
        starts = sorted(set(int(value) for value in starts if int(value) < len(sample)))
        boundaries = starts + [len(sample)]
        strokes = [sample[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))]
        variants = (
            strokes, [stroke.flip(0) for stroke in strokes], list(reversed(strokes)),
            [stroke.flip(0) for stroke in reversed(strokes)],
        )
        rows = []
        for variant in variants:
            value = torch.cat(variant, dim=0).clone()
            value[:, 7] = 0.0
            cursor = 0
            for stroke in variant:
                value[cursor, 7] = 1.0
                value[cursor:cursor + len(stroke), 8] = torch.linspace(
                    0.0, 1.0, len(stroke), device=value.device, dtype=value.dtype,
                )
                cursor += len(stroke)
            delta_shape = torch.cat((torch.zeros_like(value[:1, :2]), value[1:, :2] - value[:-1, :2]), dim=0)
            delta_canvas = torch.cat((torch.zeros_like(value[:1, 2:4]), value[1:, 2:4] - value[:-1, 2:4]), dim=0)
            delta_shape[value[:, 7] > 0.5] = 0.0
            delta_canvas[value[:, 7] > 0.5] = 0.0
            distance_shape = delta_shape.square().sum(dim=-1).sqrt()
            direction = delta_shape / distance_shape.clamp_min(1e-6).unsqueeze(-1)
            previous = torch.cat((torch.zeros_like(direction[:1]), direction[:-1]), dim=0)
            value[:, 4:6] = direction
            value[:, 6] = previous[:, 0] * direction[:, 1] - previous[:, 1] * direction[:, 0]
            canvas_distance = delta_canvas.square().sum(dim=-1).sqrt() * 128.0
            value[:, 15] = canvas_distance / 48.0
            value[:, 16] = torch.where(
                canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0),
                torch.zeros_like(canvas_distance),
            )
            value[:, 17] = 1.0
            value[:, 18] = 1.0
            rows.append(value)
        batches.append(torch.stack(rows))
    return torch.stack(batches)


def soft_rasterize_virtual06(
    coordinates: Tensor, *, size: int = 32, sigma: float = 0.025, point_stride: int = 1,
    point_weights: Tensor | None = None,
) -> Tensor:
    """필요 변수: B×K×T 좌표·선택 weight. 작동 원리: END/padding을 제외한 대표 타점을 부드러운 raster로 변환한다."""

    if size <= 0 or sigma <= 0 or point_stride <= 0:
        raise ValueError("raster size·sigma·point_stride는 양수여야 합니다.")
    original_time_shape = coordinates.shape[:-1]
    if point_weights is not None:
        if point_weights.shape != original_time_shape:
            raise ValueError("point weight는 coordinate의 원본 B×K×T 축과 일치해야 합니다.")
    coordinates = coordinates[:, :, ::point_stride]
    if point_weights is not None:
        point_weights = point_weights[:, :, ::point_stride]
    axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size
    grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
    grid = torch.stack((grid_x, grid_y), dim=-1)
    minimum = torch.full((*coordinates.shape[:2], size, size), torch.inf, dtype=coordinates.dtype, device=coordinates.device)
    weighted_maximum = torch.zeros((*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device)
    # 전체 T×H×W tensor를 한 번에 만들지 않아 Colab/모바일 연구 메모리 사용을 제한한다.
    offset = 0
    for chunk in coordinates.split(32, dim=2):
        distance = (chunk[:, :, :, None, None] - grid).square().sum(dim=-1)
        if point_weights is None:
            minimum = torch.minimum(minimum, distance.amin(dim=2))
        else:
            weights = point_weights[:, :, offset:offset + chunk.shape[2], None, None]
            occupancy = torch.exp(-distance / (2.0 * sigma * sigma)) * weights
            weighted_maximum = torch.maximum(weighted_maximum, occupancy.amax(dim=2))
            offset += chunk.shape[2]
    if point_weights is not None:
        return weighted_maximum
    return torch.exp(-minimum / (2.0 * sigma * sigma))


def soft_rasterize_virtual_segments06(
    coordinates: Tensor, state_logits: Tensor, *, size: int = 32, sigma: float = 0.025,
    segment_stride: int = 2,
) -> Tensor:
    """필요 변수: 좌표·pen state·출력 크기. 작동 원리: pen-start 연결을 억제한 선분 거리로 differentiable raster를 만든다."""

    if coordinates.shape[:-1] != state_logits.shape[:-1] or state_logits.shape[-1] != 3:
        raise ValueError("coordinate와 state logit의 batch·가설·시간 축이 일치해야 합니다.")
    if size <= 0 or sigma <= 0 or segment_stride <= 0:
        raise ValueError("raster size·sigma·segment_stride는 양수여야 합니다.")
    axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size
    grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
    grid = torch.stack((grid_x, grid_y), dim=-1)
    starts = coordinates[:, :, :-segment_stride:segment_stride]
    ends = coordinates[:, :, segment_stride::segment_stride]
    segment_count = min(starts.shape[2], ends.shape[2])
    starts, ends = starts[:, :, :segment_count], ends[:, :, :segment_count]
    state_probability = state_logits.softmax(dim=-1)
    pen_start = state_probability[..., 1]
    pen_end = state_probability[..., 2]
    valid_rows = []
    for start in range(0, coordinates.shape[2] - segment_stride, segment_stride):
        boundary = pen_start[..., start + 1:start + segment_stride + 1].amax(dim=-1)
        # END는 legacy padding의 첫 좌표이기도 하므로 target 위치까지 포함해 연결을 차단한다.
        ended_before_target = pen_end[..., start:start + segment_stride + 1].amax(dim=-1)
        valid_rows.append((1.0 - boundary) * (1.0 - ended_before_target))
    segment_valid = torch.stack(valid_rows[:segment_count], dim=2)
    maximum = torch.zeros(
        (*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device,
    )
    for first in range(0, segment_count, 16):
        start = starts[:, :, first:first + 16, None, None]
        vector = (ends[:, :, first:first + 16] - starts[:, :, first:first + 16])[:, :, :, None, None]
        relative = grid - start
        projection = (relative * vector).sum(dim=-1) / vector.square().sum(dim=-1).clamp_min(1e-8)
        closest = start + projection.clamp(0.0, 1.0).unsqueeze(-1) * vector
        distance = (grid - closest).square().sum(dim=-1)
        occupancy = torch.exp(-distance / (2.0 * sigma * sigma))
        occupancy = occupancy * segment_valid[:, :, first:first + 16, None, None]
        maximum = torch.maximum(maximum, occupancy.amax(dim=2))
    # 한 점짜리 획과 선분 양 끝은 기존 point rasterizer로 보존한다.
    points = soft_rasterize_virtual06(
        coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 32),
        point_weights=1.0 - pen_end,
    )
    return torch.maximum(maximum, points)


def virtual_raster_similarity06(
    coordinates: Tensor, raster: Tensor, *, state_logits: Tensor | None = None,
    size: int = 32, sigma: float = 0.025,
) -> Tensor:
    """필요 변수: 가설 좌표·원본 raster. 작동 원리: 재렌더링 Dice와 양방향 coverage로 라벨 독립 품질을 계산한다."""

    reconstructed = (
        soft_rasterize_virtual_segments06(coordinates, state_logits, size=size, sigma=sigma)
        if state_logits is not None else soft_rasterize_virtual06(
            coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 64),
        )
    )
    target = nn.functional.adaptive_max_pool2d(raster, (size, size))[:, 0]
    target = target[:, None].expand_as(reconstructed)
    intersection = (reconstructed * target).sum(dim=(-1, -2))
    dice = (2.0 * intersection + 1e-5) / (
        reconstructed.sum(dim=(-1, -2)) + target.sum(dim=(-1, -2)) + 1e-5
    )
    # 한쪽만 넓게 칠해 Dice를 속이는 가설을 막기 위해 precision·recall의 기하평균도 함께 본다.
    precision = intersection / reconstructed.sum(dim=(-1, -2)).clamp_min(1e-5)
    recall = intersection / target.sum(dim=(-1, -2)).clamp_min(1e-5)
    coverage = torch.sqrt((precision * recall).clamp_min(0.0))
    return 0.5 * (dice + coverage)


def raster_symmetry_logits06(
    model: "MathInk06Model", output: dict[str, Tensor], *, mode: str = "logsumexp",
) -> tuple[Tensor, Tensor]:
    """필요 변수: 가상 stroke 출력·shared 모델. 작동 원리: 정적 이미지에서 알 수 없는 방향·획순서 네 경우를 동일 encoder로 평가한다."""

    features = virtual_features06(
        output["coordinates"], output["state_logits"], output["stroke_progress"],
        contract=model.virtual_contract,
    )
    batch, hypotheses, steps, channels = features.shape
    if model.use_virtual_adapter:
        features = model.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view(
            batch, hypotheses, steps, channels,
        )
    variants = equivalent_modality_features06(features.view(batch * hypotheses, steps, channels))
    exact, family = model.classify_trajectory(variants.flatten(0, 1))
    exact = exact.view(batch, hypotheses, 4, -1)
    family = family.view(batch, hypotheses, 4, -1)
    if mode == "logsumexp":
        return torch.logsumexp(exact.log_softmax(dim=-1), dim=2), torch.logsumexp(
            family.log_softmax(dim=-1), dim=2,
        )
    if mode == "max":
        return exact.log_softmax(dim=-1).amax(dim=2), family.log_softmax(dim=-1).amax(dim=2)
    raise ValueError(f"지원하지 않는 symmetry mode입니다: {mode}")


class MathInk06Model(nn.Module):
    """필요 변수: 378 exact/family class와 선택 boundary head. 작동 원리: online·virtual stroke를 동일 embedding으로 분류한다."""

    def __init__(
        self, *, exact_classes: int, family_classes: int, hidden_size: int = 128, hypotheses: int = 4,
        raster_architecture: str = "spatial_flat_progress_v2", virtual_contract: str = "legacy_v1",
        use_virtual_adapter: bool = False, use_boundary_head: bool = False,
    ) -> None:
        super().__init__()
        if raster_architecture not in {
            "spatial_flat_v1", "spatial_flat_progress_v2", "cross_attention_v2", "cross_attention_8x8_v3",
            "gated_cross_attention_8x8_v4", "split_auxiliary_v5", "fine_cross_attention_16x16_v6",
            "gated_fine_cross_attention_16x16_v7",
            "ink_pointer_32x32_v8",
        }:
            raise ValueError("지원하지 않는 raster architecture입니다.")
        self.hidden_size = hidden_size
        self.hypotheses = hypotheses
        self.raster_architecture = raster_architecture
        self.virtual_contract = virtual_contract
        self.use_virtual_adapter = use_virtual_adapter
        self.use_boundary_head = use_boundary_head
        self.virtual_adapter_weight = 1.0
        self.trajectory_encoder = SharedTrajectoryEncoder06(hidden_size=hidden_size)
        self.virtual_adapter = VirtualTrajectoryAdapter06()
        self.exact_head = nn.Linear(hidden_size * 3, exact_classes)
        self.family_head = nn.Linear(hidden_size * 3, family_classes)
        self.boundary_head = nn.Linear(hidden_size * 3, 1) if use_boundary_head else None
        self.raster_encoder = DepthwiseRasterEncoder06(hidden_size)
        self.virtual_decoder = VirtualStrokeDecoder06(hidden_size, hypotheses)
        self.auxiliary_virtual_decoder = (
            VirtualStrokeDecoder06(hidden_size, 2) if raster_architecture == "split_auxiliary_v5" else None
        )

    def initialize_auxiliary_from_primary(self) -> None:
        """필요 변수: split auxiliary 모델. 작동 원리: 유효한 primary 0·1번 출력을 auxiliary 초기값으로 복제한다."""

        if self.auxiliary_virtual_decoder is None:
            raise ValueError("split_auxiliary_v5 모델에서만 auxiliary 초기화가 가능합니다.")
        source = self.virtual_decoder.state_dict()
        target = self.auxiliary_virtual_decoder.state_dict()
        for key, target_value in target.items():
            source_value = source[key]
            if source_value.shape == target_value.shape:
                target[key] = source_value.detach().clone()
            elif key == "hypothesis.weight" and source_value.shape[0] >= 2:
                target[key] = source_value[:2].detach().clone()
            else:
                raise ValueError(f"auxiliary 초기화 shape가 일치하지 않습니다: {key}")
        self.auxiliary_virtual_decoder.load_state_dict(target)

    def encode_trajectory(self, sequence: Tensor) -> Tensor:
        """필요 변수: B×128×19 canonical sequence. 작동 원리: 모든 symbol/behavior head가 공유할 trajectory embedding을 한 번 계산한다."""

        return self.trajectory_encoder(sequence)

    def classify_trajectory(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
        """필요 변수: B×128×19. 작동 원리: shared embedding에서 기존 exact/family 출력 계약을 유지한다."""

        embedding = self.encode_trajectory(sequence)
        return self.exact_head(embedding), self.family_head(embedding)

    def classify_trajectory_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """필요 변수: boundary head가 활성화된 sequence. 작동 원리: 한 embedding에서 exact/family/경계 침범 logit을 함께 반환한다."""

        if self.boundary_head is None:
            raise RuntimeError("boundary head가 활성화되지 않았습니다.")
        embedding = self.encode_trajectory(sequence)
        return self.exact_head(embedding), self.family_head(embedding), self.boundary_head(embedding).squeeze(-1)

    def forward_online(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
        """필요 변수: 실제 canonical tap. 작동 원리: raster 우회 없이 shared trajectory 분류를 반환한다."""

        return self.classify_trajectory(sequence)

    def forward_online_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """필요 변수: 실제 canonical tap. 작동 원리: 기존 LiteRT forward를 바꾸지 않고 연구용 boundary logit을 추가 노출한다."""

        return self.classify_trajectory_with_boundary(sequence)

    def decode_raster_trajectories(self, raster: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        """필요 변수: B×1×128×128. 작동 원리: architecture별 top-4 좌표·state·progress·score를 한 경로로 만든다."""

        embedding, spatial_tokens = self.raster_encoder(
            raster, fine_tokens=self.raster_architecture in {
                "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7",
            },
            pointer_tokens=self.raster_architecture == "ink_pointer_32x32_v8",
        )
        if self.raster_architecture == "split_auxiliary_v5":
            if self.auxiliary_virtual_decoder is None:
                raise RuntimeError("split auxiliary decoder가 초기화되지 않았습니다.")
            primary = self.virtual_decoder(embedding)
            auxiliary = self.auxiliary_virtual_decoder(embedding, spatial_tokens, gated_attention=True)
            return tuple(
                torch.cat((primary[index][:, :2], auxiliary[index]), dim=1) for index in range(4)
            )  # type: ignore[return-value]
        memory = spatial_tokens if self.raster_architecture in {
            "cross_attention_v2", "cross_attention_8x8_v3", "gated_cross_attention_8x8_v4",
            "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7",
            "ink_pointer_32x32_v8",
        } else None
        pointer_mode = self.raster_architecture == "ink_pointer_32x32_v8"
        return self.virtual_decoder(
            embedding, memory, gated_attention=self.raster_architecture in {
                "gated_cross_attention_8x8_v4", "gated_fine_cross_attention_16x16_v7",
            },
            pointer_positions=self.raster_encoder.pointer_positions if pointer_mode else None,
            ink_prior=(
                nn.functional.adaptive_max_pool2d(raster, (32, 32)).flatten(2)[:, 0]
                if pointer_mode else None
            ),
        )

    def forward_raster(self, raster: Tensor) -> dict[str, Tensor]:
        """필요 변수: B×1×128×128. 작동 원리: top-4 가상 stroke를 만든 뒤 shared TCN으로만 분류한다."""

        coordinates, states, progress, hypothesis_scores = self.decode_raster_trajectories(raster)
        features = virtual_features06(
            coordinates, states, None if self.raster_architecture == "spatial_flat_v1" else progress,
            contract=self.virtual_contract,
        )
        batch, hypotheses, steps, channels = features.shape
        if self.use_virtual_adapter:
            raw_features = features
            adapted_features = self.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view(
                batch, hypotheses, steps, channels,
            )
            features = raw_features + self.virtual_adapter_weight * (adapted_features - raw_features)
        flat_features = features.view(batch * hypotheses, steps, channels)
        if self.boundary_head is None:
            exact, family = self.classify_trajectory(flat_features)
            boundary = None
        else:
            exact, family, boundary = self.classify_trajectory_with_boundary(flat_features)
        output = {
            "coordinates": coordinates, "state_logits": states, "stroke_progress": progress,
            "hypothesis_scores": hypothesis_scores,
            "exact_logits": exact.view(batch, hypotheses, -1), "family_logits": family.view(batch, hypotheses, -1),
        }
        if boundary is not None:
            output["boundary_logits"] = boundary.view(batch, hypotheses)
        return output

    def forward(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
        """필요 변수: LiteRT용 online tensor. 작동 원리: 기본 forward를 online 경로로 고정한다."""

        return self.forward_online(sequence)


def boundary_auxiliary_loss06(
    boundary_logits: Tensor,
    boundary_targets: Tensor,
    *,
    positive_weight: float = 1.0,
    sample_weight: Tensor | None = None,
) -> Tensor:
    """필요 변수: 후보별 경계 logit·0/1 target·선택 weight. 작동 원리: class imbalance를 보정한 binary auxiliary loss를 계산한다."""

    if boundary_logits.shape != boundary_targets.shape:
        raise ValueError("boundary logit과 target shape가 다릅니다.")
    if positive_weight <= 0.0:
        raise ValueError("boundary positive weight는 0보다 커야 합니다.")
    targets = boundary_targets.to(dtype=boundary_logits.dtype)
    loss = nn.functional.binary_cross_entropy_with_logits(
        boundary_logits,
        targets,
        pos_weight=torch.as_tensor(positive_weight, dtype=boundary_logits.dtype, device=boundary_logits.device),
        reduction="none",
    )
    if sample_weight is not None:
        if sample_weight.shape != loss.shape:
            raise ValueError("boundary sample weight shape가 다릅니다.")
        normalized = sample_weight.to(loss).clamp_min(0.0)
        return (loss * normalized).sum() / normalized.sum().clamp_min(1e-8)
    return loss.mean()


def fuse_raster_logits06(
    output: dict[str, Tensor], *, mode: str = "max", score_weight: float = 1.0,
    family_weight: float = 0.0, geometry_weight: float = 0.0,
    exact_family_index: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
    """필요 변수: top-4 exact/family/quality logit. 작동 원리: 기호별 증거를 합치고 debug 대표 가설을 반환한다."""

    exact = output["exact_logits"].log_softmax(dim=-1)
    score = output["hypothesis_scores"].log_softmax(dim=-1).unsqueeze(-1)
    joint = exact + score_weight * score
    if geometry_weight:
        if "geometry_scores" not in output:
            raise ValueError("geometry_weight를 사용할 때 geometry_scores가 필요합니다.")
        geometry = output["geometry_scores"].clamp_min(1e-6).log().unsqueeze(-1)
        joint = joint + geometry_weight * geometry
    if family_weight:
        if exact_family_index is None:
            raise ValueError("family_weight를 사용할 때 exact_family_index가 필요합니다.")
        family = output["family_logits"].log_softmax(dim=-1)[..., exact_family_index]
        joint = joint + family_weight * family
    if mode == "max":
        fused = joint.amax(dim=1)
    elif mode == "logsumexp":
        fused = torch.logsumexp(joint, dim=1)
    elif mode == "score_pick":
        selected = output["hypothesis_scores"].argmax(dim=1)
        fused = joint[torch.arange(len(joint), device=joint.device), selected]
        return fused, selected
    else:
        raise ValueError(f"지원하지 않는 raster fusion mode입니다: {mode}")
    # 합산 모드의 debug 좌표는 최종 top-1 기호에 가장 크게 기여한 가설로 설명한다.
    predicted = fused.argmax(dim=-1)
    contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1)
    return fused, contribution.argmax(dim=1)


def fuse_hypothesis_class_logits06(
    class_logits: Tensor, hypothesis_scores: Tensor, *, mode: str = "logsumexp", score_weight: float = 1.0,
) -> Tensor:
    """필요 변수: B×H×C 분류 logit·B×H 가설 점수. 작동 원리: exact/family 공통 규칙으로 top-H 증거를 결합한다."""

    if class_logits.ndim != 3 or hypothesis_scores.shape != class_logits.shape[:2]:
        raise ValueError("class logit과 hypothesis score shape가 일치하지 않습니다.")
    joint = class_logits.log_softmax(dim=-1)
    joint = joint + score_weight * hypothesis_scores.log_softmax(dim=-1).unsqueeze(-1)
    if mode == "max":
        return joint.amax(dim=1)
    if mode == "logsumexp":
        return torch.logsumexp(joint, dim=1)
    if mode == "score_pick":
        selected = hypothesis_scores.argmax(dim=1)
        return joint[torch.arange(len(joint), device=joint.device), selected]
    raise ValueError(f"지원하지 않는 hypothesis fusion mode입니다: {mode}")


def hypothesis_quality_features06(output: dict[str, Tensor], exact_family_index: Tensor) -> Tensor:
    """필요 변수: 가설별 logits/state/좌표·family 사상. 작동 원리: raster label 없이 가설 품질 특징을 만든다."""

    exact_log_probability = output["exact_logits"].log_softmax(dim=-1)
    exact_probability = exact_log_probability.exp()
    top_values, top_indices = exact_log_probability.topk(min(2, exact_log_probability.shape[-1]), dim=-1)
    predicted = top_indices[..., 0]
    margin = top_values[..., 0] - top_values[..., -1]
    exact_entropy = -(exact_probability * exact_log_probability).sum(dim=-1) / np.log(max(2, exact_probability.shape[-1]))
    family_log_probability = output["family_logits"].log_softmax(dim=-1)
    predicted_family = exact_family_index[predicted]
    predicted_family_log_probability = family_log_probability.gather(2, predicted_family.unsqueeze(-1)).squeeze(-1)
    state_log_probability = output["state_logits"].log_softmax(dim=-1)
    state_probability = state_log_probability.exp()
    state_entropy = -(state_probability * state_log_probability).sum(dim=-1).mean(dim=-1) / np.log(3.0)
    start_confidence = state_probability[..., 1].amax(dim=-1)
    delta = output["coordinates"][:, :, 1:] - output["coordinates"][:, :, :-1]
    distance = delta.square().sum(dim=-1).sqrt()
    path_length = distance.mean(dim=-1)
    direction = delta / distance.clamp_min(1e-6).unsqueeze(-1)
    turn = (
        direction[:, :, :-1, 0] * direction[:, :, 1:, 1]
        - direction[:, :, :-1, 1] * direction[:, :, 1:, 0]
    ).abs().mean(dim=-1)
    agreement = (predicted[:, :, None] == predicted[:, None, :]).float().mean(dim=-1)
    branch = nn.functional.one_hot(
        torch.arange(predicted.shape[1], device=predicted.device), num_classes=predicted.shape[1],
    ).to(dtype=exact_probability.dtype).unsqueeze(0).expand(len(predicted), -1, -1)
    scalar = torch.stack((
        top_values[..., 0], margin, exact_entropy, family_log_probability.amax(dim=-1),
        predicted_family_log_probability, output["hypothesis_scores"].log_softmax(dim=-1),
        start_confidence, state_entropy, path_length, turn, agreement,
    ), dim=-1)
    return torch.cat((scalar, branch), dim=-1)


class HypothesisSelector06(nn.Module):
    """필요 변수: trajectory-only 품질 특징. 작동 원리: 각 virtual hypothesis의 혼합 logit을 예측한다."""

    def __init__(
        self, input_size: int = 15, hidden_size: int = 24, *, label_classes: int = 0,
        label_embedding_size: int = 0,
    ) -> None:
        super().__init__()
        if (label_classes > 0) != (label_embedding_size > 0):
            raise ValueError("label class와 embedding 크기는 함께 지정해야 합니다.")
        self.label_embedding = (
            nn.Embedding(label_classes, label_embedding_size) if label_classes > 0 else None
        )
        self.network = nn.Sequential(
            nn.LayerNorm(input_size + label_embedding_size),
            nn.Linear(input_size + label_embedding_size, hidden_size), nn.GELU(),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, features: Tensor, predicted_labels: Tensor | None = None) -> Tensor:
        """필요 변수: 품질 특징·선택 top-1 label. 작동 원리: 가설별 scalar quality logit을 반환한다."""

        if self.label_embedding is not None:
            if predicted_labels is None:
                raise ValueError("class-conditional selector에는 predicted_labels가 필요합니다.")
            features = torch.cat((features, self.label_embedding(predicted_labels)), dim=-1)
        return self.network(features).squeeze(-1)


def initialize_from_05(model: MathInk06Model, checkpoint_paths: Sequence[Path]) -> None:
    """필요 변수: 0.6 모델·동일 0.5 seed checkpoint. 작동 원리: 유효한 단일 teacher를 19채널 student 초기값으로 이식한다."""

    if not checkpoint_paths:
        raise ValueError("0.5 checkpoint가 필요합니다.")
    # 서로 다른 seed의 비선형망 weight 평균은 logit ensemble과 동등하지 않고 즉시 정확도를 붕괴시킨다.
    # 첫 seed를 유효 초기값으로 사용하고 3-seed 정보 결합은 별도 distillation loss에서 수행한다.
    checkpoint = torch.load(checkpoint_paths[0], map_location="cpu", weights_only=False)
    state = checkpoint["state_dict"]
    target = model.state_dict()
    mapping = {
        "trajectory_encoder.input_projection": "encoder.input_projection",
        "trajectory_encoder.blocks": "encoder.blocks",
        "trajectory_encoder.attention": "encoder.attention",
        "exact_head": "exact_head", "family_head": "family_head",
    }
    for target_key in list(target):
        source_key = next((target_key.replace(prefix, source) for prefix, source in mapping.items() if target_key.startswith(prefix)), None)
        if source_key is None or source_key not in state:
            continue
        source_value = state[source_key].float()
        if source_value.shape == target[target_key].shape:
            target[target_key] = source_value
        elif target_key.endswith("input_projection.0.weight") and source_value.shape[1] == 15 and target[target_key].shape[1] == 19:
            expanded = torch.zeros_like(target[target_key])
            expanded[:, :15] = source_value
            target[target_key] = expanded
    model.load_state_dict(target)


@dataclass(frozen=True, slots=True)
class SymbolCandidate06:
    """필요 변수: token·확률. 작동 원리: 모바일 공개 결과의 후보 한 개를 표현한다."""

    token: str
    probability: float


class MathInk06Engine:
    """필요 변수: 0.6 checkpoint. 작동 원리: 원본 stroke 또는 raster에서 텍스트 후보만 반환한다."""

    def __init__(
        self, checkpoint: Path, *, adapter_checkpoint: Path | None = None, device: str = "cpu",
    ) -> None:
        """필요 변수: base와 선택 composite adapter. 작동 원리: base→shared state→modality adapter 순서로 런타임을 구성한다."""

        payload = torch.load(checkpoint, map_location=device, weights_only=False)
        self.labels = tuple(str(value) for value in payload["exact_labels"])
        self.family_labels = tuple(str(value) for value in payload["family_labels"])
        self.model_version = str(payload.get("model_version", "aiflow-math-ink-0.6"))
        self.model = MathInk06Model(
            exact_classes=len(self.labels), family_classes=len(payload["family_labels"]),
            hidden_size=int(payload["hidden_size"]), hypotheses=int(payload.get("hypotheses", 4)),
            raster_architecture=str(payload.get("raster_architecture", "spatial_flat_v1")),
            virtual_contract=str(payload.get("virtual_contract", "legacy_v1")),
            use_virtual_adapter=bool(payload.get("use_virtual_adapter", False)),
        ).to(device)
        # candidate1에는 cross-attention 파라미터가 없으므로 구 checkpoint는 flat 경로로 호환 로드한다.
        self.model.load_state_dict(payload["state_dict"], strict=False)
        self.model.virtual_adapter_weight = float(payload.get("virtual_adapter_weight", 1.0))
        self.model.eval()
        self.device = torch.device(device)
        self.raster_fusion = {
            "mode": "max", "score_weight": 1.0, "family_weight": 0.0, "geometry_weight": 0.0,
            "symmetry_weight": 0.0, "symmetry_mode": "logsumexp",
            **dict(payload.get("raster_fusion", {})),
        }
        self.composite_adapter: nn.Module = nn.Identity()
        self.online_adapter: nn.Module = nn.Identity()
        self.raster_adapter: nn.Module = nn.Identity()
        self.online_family_fusion_weight = 0.0
        if adapter_checkpoint is not None:
            adapter_payload = torch.load(
                adapter_checkpoint, map_location=device, weights_only=False,
            )
            shared_state = adapter_payload.get("shared_state_dict") or {}
            if shared_state:
                incompatible = self.model.load_state_dict(shared_state, strict=False)
                if incompatible.unexpected_keys:
                    raise ValueError(
                        f"adapter shared state key 오류: {incompatible.unexpected_keys}"
                    )
            architecture = str(adapter_payload["adapter_architecture"])
            if architecture == "local_v1":
                adapter: nn.Module = VirtualTrajectoryAdapter06()
            else:
                from .skeleton_adapter06 import (
                    DualModalityTrajectoryAdapter06, SkeletonTrajectoryAdapter06,
                )
                if architecture == "tcn_v2":
                    adapter = SkeletonTrajectoryAdapter06()
                elif architecture == "dual_tcn_v3":
                    adapter = DualModalityTrajectoryAdapter06()
                else:
                    raise ValueError(
                        f"지원하지 않는 adapter architecture입니다: {architecture}"
                    )
            adapter.load_state_dict(adapter_payload["state_dict"])
            adapter = adapter.to(self.device).eval()
            self.online_family_fusion_weight = float(
                adapter_payload.get("family_fusion_weight", 0.0)
            )
            if not 0.0 <= self.online_family_fusion_weight <= 1.0:
                raise ValueError("online family fusion weight는 0~1 범위여야 합니다.")
            self.composite_adapter = adapter
            if architecture == "dual_tcn_v3":
                self.online_adapter = adapter.online  # type: ignore[attr-defined]
                self.raster_adapter = adapter.raster  # type: ignore[attr-defined]
            else:
                self.online_adapter = self.raster_adapter = adapter
            self.model_version = (
                f"{self.model_version}+{adapter_payload.get('model_version', architecture)}"
            )
        family_to_index = {label: index for index, label in enumerate(self.family_labels)}
        # checkpoint가 가진 ontology와 동일한 exact→family 사상을 기기 내 상수 tensor로 유지한다.
        from .trajectory_sequence import shape_family
        self.exact_family_index = torch.tensor(
            [family_to_index[shape_family(label)] for label in self.labels], device=self.device,
        )
        selector_payload = payload.get("hypothesis_selector")
        self.hypothesis_selector: HypothesisSelector06 | None = None
        if selector_payload:
            self.hypothesis_selector = HypothesisSelector06(
                input_size=int(selector_payload["input_size"]), hidden_size=int(selector_payload["hidden_size"]),
                label_classes=int(selector_payload.get("label_classes", 0)),
                label_embedding_size=int(selector_payload.get("label_embedding_size", 0)),
            ).to(self.device)
            self.hypothesis_selector.load_state_dict(selector_payload["state_dict"])
            self.hypothesis_selector.eval()

    def fuse_raster_output(self, output: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
        """필요 변수: model raster 출력. 작동 원리: checkpoint에 따라 learned selector 또는 고정 fusion을 적용한다."""

        if self.hypothesis_selector is None:
            return fuse_raster_logits06(
                output, mode=str(self.raster_fusion["mode"]),
                score_weight=float(self.raster_fusion["score_weight"]),
                family_weight=float(self.raster_fusion["family_weight"]),
                geometry_weight=float(self.raster_fusion["geometry_weight"]),
                exact_family_index=self.exact_family_index,
            )
        features = hypothesis_quality_features06(output, self.exact_family_index)
        predicted_labels = output["exact_logits"].argmax(dim=-1)
        selector_log_probability = self.hypothesis_selector(features, predicted_labels).log_softmax(dim=1)
        joint = output["exact_logits"].log_softmax(dim=-1) + selector_log_probability.unsqueeze(-1)
        fused = torch.logsumexp(joint, dim=1)
        predicted = fused.argmax(dim=-1)
        contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1)
        return fused, contribution.argmax(dim=1)

    def _forward_raster_composite06(self, raster: Tensor) -> dict[str, Tensor]:
        """필요 변수: 정규화 raster. 작동 원리: virtual top-4를 외부 raster adapter까지 거쳐 shared head로 분류한다."""

        if isinstance(self.raster_adapter, nn.Identity):
            return self.model.forward_raster(raster)
        coordinates, states, progress, hypothesis_scores = self.model.decode_raster_trajectories(raster)
        features = virtual_features06(
            coordinates, states,
            None if self.model.raster_architecture == "spatial_flat_v1" else progress,
            contract=self.model.virtual_contract,
        )
        batch, hypotheses, steps, channels = features.shape
        if self.model.use_virtual_adapter:
            raw_features = features
            internal = self.model.virtual_adapter(
                features.reshape(batch * hypotheses, steps, channels),
            ).reshape(batch, hypotheses, steps, channels)
            features = (
                raw_features
                + self.model.virtual_adapter_weight * (internal - raw_features)
            )
        flat = self.raster_adapter(
            features.reshape(batch * hypotheses, steps, channels),
        )
        exact, family = self.model.classify_trajectory(flat)
        return {
            "coordinates": coordinates,
            "state_logits": states,
            "stroke_progress": progress,
            "hypothesis_scores": hypothesis_scores,
            "exact_logits": exact.reshape(batch, hypotheses, -1),
            "family_logits": family.reshape(batch, hypotheses, -1),
        }

    def _result(self, logits: Tensor, started: float, top_k: int) -> dict[str, Any]:
        """필요 변수: fused logits·시작시각·k. 작동 원리: stroke/image 없이 모바일 공개 SymbolResult를 만든다."""

        probability = logits.softmax(dim=-1)[0]
        values, indices = probability.topk(min(top_k, len(self.labels)))
        candidates = [SymbolCandidate06(self.labels[int(index)], float(value)) for value, index in zip(values, indices, strict=True)]
        return {
            "candidates": [candidate.__dict__ if hasattr(candidate, "__dict__") else {"token": candidate.token, "probability": candidate.probability} for candidate in candidates],
            "confidence": candidates[0].probability, "modelVersion": self.model_version,
            "latencyMs": (time.perf_counter() - started) * 1000.0,
        }

    def _fuse_online_exact06(self, exact: Tensor, family: Tensor) -> Tensor:
        """필요 변수: exact/family logit. 작동 원리: validation에서 고정한 형태군 prior로 exact 후보만 재정렬한다."""

        if not self.online_family_fusion_weight:
            return exact
        return (
            exact.log_softmax(dim=-1)
            + self.online_family_fusion_weight
            * family.log_softmax(dim=-1)[:, self.exact_family_index]
        )

    def recognize_online(self, strokes: Sequence[dict[str, Any]], *, canvas_width: float, canvas_height: float, top_k: int = 5) -> dict[str, Any]:
        """필요 변수: 원본 stroke·canvas. 작동 원리: 6Hz 재구성 후 기기 밖으로 내보낼 텍스트 후보만 반환한다."""

        started = time.perf_counter()
        ink = canonicalize_ink06(strokes, canvas_width=canvas_width, canvas_height=canvas_height)
        sequence = torch.from_numpy(ink.features).unsqueeze(0).to(self.device)
        with torch.inference_mode():
            exact, family = self.model.forward_online(self.online_adapter(sequence))
            exact = self._fuse_online_exact06(exact, family)
        return self._result(exact, started, top_k)

    def recognize_raster(self, image: Image.Image, *, top_k: int = 5, debug: bool = False) -> dict[str, Any]:
        """필요 변수: PIL image·k·로컬 debug. 작동 원리: 가상 stroke를 거쳐 텍스트만 반환하고 debug 때만 좌표를 붙인다."""

        started = time.perf_counter()
        normalized = image.convert("L").resize((128, 128), Image.Resampling.LANCZOS)
        raster = 1.0 - torch.from_numpy(np.asarray(normalized, dtype=np.float32) / 255.0)
        with torch.inference_mode():
            output = self._forward_raster_composite06(
                raster.view(1, 1, 128, 128).to(self.device),
            )
            symmetry_weight = float(self.raster_fusion["symmetry_weight"])
            if symmetry_weight:
                symmetry_exact, symmetry_family = raster_symmetry_logits06(
                    self.model, output, mode=str(self.raster_fusion["symmetry_mode"]),
                )
                output["exact_logits"] = (
                    (1.0 - symmetry_weight) * output["exact_logits"].log_softmax(dim=-1)
                    + symmetry_weight * symmetry_exact
                )
                output["family_logits"] = (
                    (1.0 - symmetry_weight) * output["family_logits"].log_softmax(dim=-1)
                    + symmetry_weight * symmetry_family
                )
            if float(self.raster_fusion["geometry_weight"]):
                output["geometry_scores"] = virtual_raster_similarity06(
                    output["coordinates"], raster.view(1, 1, 128, 128).to(self.device),
                )
            logits, selected = self.fuse_raster_output(output)
            flat_index = int(selected[0])
        result = self._result(logits, started, top_k)
        if debug:
            result["virtualHypothesis"] = output["coordinates"][0, flat_index].cpu().tolist()
        return result