"""Offline unit checks for the Lecture 3 examples."""
import csv
import itertools
from pathlib import Path
import sys
from types import SimpleNamespace
import unittest

import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "lecture_3"))
import esm2_flow_guidance as flow
import esm2_diffusion_guidance as diffusion


class GuidanceTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(2)

    def test_example_annotations(self):
        with (ROOT / "lecture_3" / "esm2_example.csv").open(newline="") as handle:
            rows = list(csv.DictReader(handle))
        self.assertEqual(len(rows), 64)
        self.assertEqual({len(row["sequence"]) for row in rows}, {24})
        labels = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
        for module in (flow, diffusion):
            computed = module.composition_proxies([row["sequence"] for row in rows])
            torch.testing.assert_close(computed, labels, atol=1e-7, rtol=0)
        for row in rows:
            self.assertEqual(int(row["c"]), int(float(row["r1"]) > 0))

    def test_scalarization_weights(self):
        for module in (flow, diffusion):
            with self.subTest(module=module.__name__):
                weights = module.normalize_weights((7, 3)).cpu()
                torch.testing.assert_close(weights, torch.tensor([0.7, 0.3]))
                for invalid in ((0, 0), (-1, 2), (1,), (float("nan"), 1)):
                    with self.assertRaises(ValueError):
                        module.normalize_weights(invalid)

    def test_reward_gradient_without_weight_updates(self):
        for module in (flow, diffusion):
            with self.subTest(module=module.__name__):
                torch.manual_seed(3)
                reward = module.RewardModel(4, 5).to(module.DEVICE)
                reward.eval().requires_grad_(False)
                state = torch.randn(2, 4, 5, device=module.DEVICE)
                time = torch.tensor([0.3, 0.7], device=module.DEVICE)
                weights = module.normalize_weights((0.7, 0.3))
                gradient = module.reward_gradient(reward, state, time, weights)
                delta = torch.zeros_like(state)
                delta[0, 1, 2] = 1e-3
                difference = reward(state + delta, time) - reward(state - delta, time)
                estimate = (difference * weights).sum() / (2e-3)
                torch.testing.assert_close(gradient[0, 1, 2], estimate, atol=3e-5, rtol=3e-3)
                self.assertFalse(gradient.requires_grad)
                self.assertTrue(torch.isfinite(gradient).all())
                self.assertTrue(all(p.grad is None for p in reward.parameters()))

    def test_ddpm_schedule(self):
        self.assertEqual(diffusion.alpha_bars[0].item(), 1.0)
        self.assertLess(diffusion.alpha_bars[-1].item(), 1e-4)
        self.assertEqual(diffusion.posterior_variances[1].item(), 0.0)
        self.assertTrue((diffusion.posterior_variances >= 0).all())
        self.assertTrue(torch.isfinite(diffusion.posterior_variances).all())

    def test_constrained_decoder_against_exhaustive_search(self):
        torch.manual_seed(11)
        scores = torch.randn(3, 4, 20)
        tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda tokens: list(range(20)))
        esm = SimpleNamespace(lm_head=torch.nn.Identity())
        stats = {"z_mean": torch.zeros(1, 1, 20), "z_std": torch.ones(1, 1, 20)}
        for module in (flow, diffusion):
            polar = [i for i, a in enumerate(module.AMINO_ACIDS) if a in module.POLAR_RESIDUES]
            nonpolar = [i for i in range(20) if i not in polar]
            for minimum in range(5):
                sequences = module.decode(scores, esm, tokenizer, stats, min_polar=minimum)
                for batch, sequence in enumerate(sequences):
                    count = sum(a in module.POLAR_RESIDUES for a in sequence)
                    self.assertGreaterEqual(count, minimum)
                    actual = sum(float(scores[batch, j, module.AMINO_ACIDS.index(a)])
                                 for j, a in enumerate(sequence))
                    optimal = max(
                        sum(float(scores[batch, j, polar if flag else nonpolar].max())
                            for j, flag in enumerate(assignment))
                        for assignment in itertools.product((0, 1), repeat=4)
                        if sum(assignment) >= minimum
                    )
                    self.assertAlmostEqual(actual, optimal, places=5)
            with self.assertRaises(ValueError):
                module.decode(scores, esm, tokenizer, stats, min_polar=5)


if __name__ == "__main__":
    unittest.main()