import unittest
from types import SimpleNamespace

import torch

from diffusion_cli.sampling import (
    betaScheduler,
    cfgDenoise,
    halfLogSnrToSigma,
    latentShape,
    offsetFirstSigmaForSnr,
    sampleSeeds3,
    sigmaToHalfLogSnr,
)


class FakeConditioning:
    pass


class FakeModel:
    def __init__(self):
        self.calls = []

    def forward(self, latent, sigma, conditioning):
        self.calls.append((sigma.detach().clone(), conditioning))
        if conditioning is POSITIVE:
            return latent * 0.25
        return latent * -0.25


POSITIVE = FakeConditioning()
NEGATIVE = FakeConditioning()
CONDITIONING = SimpleNamespace(positive=POSITIVE, negative=NEGATIVE)


class SamplingTest(unittest.TestCase):
    def testLatentShapeUsesFluxDownscale(self):
        self.assertEqual(latentShape(1, 1248, 832), (1, 16, 156, 104))

    def testBetaSchedulerAppendsZero(self):
        sigmas = betaScheduler(10, device="cpu", dtype=torch.float32)

        self.assertEqual(sigmas.shape, (11,))
        self.assertEqual(sigmas[-1].item(), 0.0)
        self.assertTrue(torch.all(torch.isfinite(sigmas)))
        self.assertGreater(sigmas[0].item(), 0.99)
        self.assertLess(sigmas[-2].item(), 0.2)

    def testCfgDenoiseUsesNegativePathAtCfgOne(self):
        model = FakeModel()
        latent = torch.ones(1, 1, 1, 1)
        sigma = torch.ones(1)

        output = cfgDenoise(model, latent, sigma, CONDITIONING, 1.0)

        self.assertTrue(torch.equal(output, latent * 0.25))
        self.assertEqual(len(model.calls), 2)
        self.assertIs(model.calls[0][1], POSITIVE)
        self.assertIs(model.calls[1][1], NEGATIVE)

    def testFlowLogSnrConversionsRoundTrip(self):
        sigma = torch.tensor([0.2, 0.5, 0.8])

        half_log_snr = sigmaToHalfLogSnr(sigma)
        output = halfLogSnrToSigma(half_log_snr)

        self.assertTrue(torch.allclose(output, sigma))

    def testOffsetFirstSigmaAvoidsInfiniteLogSnr(self):
        sigmas = torch.tensor([1.0, 0.5, 0.0])

        output = offsetFirstSigmaForSnr(sigmas)

        self.assertLess(output[0].item(), 1.0)
        self.assertEqual(output[1].item(), 0.5)
        self.assertEqual(output[2].item(), 0.0)

    def testSampleSeeds3RunsAndRemainsFinite(self):
        model = FakeModel()
        latent = torch.zeros(1, 1, 2, 2)
        sigmas = torch.tensor([1.0, 0.5, 0.0])

        output = sampleSeeds3(
            model,
            latent,
            sigmas,
            CONDITIONING,
            1.0,
            eta=0.0,
            s_noise=0.0,
        )

        self.assertEqual(output.shape, latent.shape)
        self.assertTrue(torch.all(torch.isfinite(output)))
        self.assertEqual(len(model.calls), 8)


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