#!/usr/bin/env python3
"""CPU checks for high-CFG reconstruction from independently decoded FP64 z."""

from io import BytesIO
from pathlib import Path
import sys
import unittest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch

from transformer_flow import Model
from nf_roundtrip import closed_loop_forward, closed_loop_reverse, error_metrics


class ClosedLoopTests(unittest.TestCase):
    def make_model(self, sos=True, patch_size=1, softplus=True):
        torch.manual_seed(42)
        model = Model(in_channels=4, img_size=4, patch_size=patch_size, channels=64,
                      num_blocks=3, layers_per_block=[1, 0, 2], head_dim=64,
                      rope=True, sos=sos, txt_size=4, txt_dim=16,
                      cond_top_only=True, use_softplus=softplus,
                      use_final_norm=True, soft_clip=4, seq_order='L2R').eval()
        with torch.no_grad():
            for block in model.blocks:
                if hasattr(block, 'proj_out'):
                    block.proj_out.weight.normal_(std=0.08)
                    block.proj_out.bias.normal_(std=0.05)
        return model

    @torch.no_grad()
    def test_saved_z_independent_decode_and_feedback(self):
        for sos, patch_size, softplus in ((True, 1, True), (False, 1, False), (True, 2, True)):
            model = self.make_model(sos, patch_size, softplus)
            x, y = torch.randn(2, 4, 4, 4), torch.randn(2, 4, 16)
            x, y = x.bfloat16().float(), y.bfloat16().float()
            for guidance in (0, 3.6, 7.2):
                with self.subTest(sos=sos, patch_size=patch_size, softplus=softplus, cfg=guidance):
                    condition = torch.cat([y, torch.zeros_like(y)]) if guidance else y
                    z, states, simulated = closed_loop_forward(model, x, condition, guidance)
                    self.assertEqual(z.dtype, torch.float64)
                    self.assertTrue(all(p.dtype == torch.float32 for p in model.parameters()))
                    buffer = BytesIO()
                    torch.save(z, buffer)
                    buffer.seek(0)
                    loaded_z = torch.load(buffer, weights_only=True)
                    sequence = closed_loop_reverse(model, loaded_z, condition, guidance, return_sequence=True)
                    self.assertEqual(len(sequence), len(model.blocks) + 1)
                    torch.testing.assert_close(sequence[1].float(), simulated, atol=0, rtol=0)
                    torch.testing.assert_close(sequence[1], model.unpatchify(states[-2]).double(),
                                               atol=1e-13, rtol=1e-13)
                    torch.testing.assert_close(sequence[-1], x, atol=3e-5, rtol=3e-5)
                    repeated = closed_loop_reverse(model, loaded_z, condition, guidance)
                    torch.testing.assert_close(repeated, sequence[-1], atol=0, rtol=0)
                    perturbed = closed_loop_reverse(model, loaded_z + 0.1, condition, guidance)
                    self.assertGreater((perturbed - repeated).abs().mean().item(), 0.01)

    def test_rejects_truncated_z_and_low_precision_weights(self):
        model = self.make_model()
        x, condition = torch.zeros(1, 4, 4, 4), torch.zeros(2, 4, 16)
        with self.assertRaisesRegex(ValueError, 'original FP64 z'):
            closed_loop_reverse(model, x, condition, 3.6)
        with self.assertRaisesRegex(ValueError, 'FP32 model weights'):
            closed_loop_forward(model.bfloat16(), x, condition, 3.6)

    def test_metrics_preserve_fp64_errors(self):
        metrics = error_metrics(torch.tensor([1 + 1e-12], dtype=torch.float64),
                                torch.ones(1, dtype=torch.float64))
        self.assertGreater(metrics['mae'], 1e-13)
        self.assertLess(metrics['mae'], 2e-12)


if __name__ == '__main__':
    torch.set_num_threads(2)
    unittest.main()