from unittest import TestCase

import torch

from chroma.layers.norm import MaskedBatchNorm1d


class TestBatchNorm(TestCase):
    def test_norm(self):
        device = (
            torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        )
        B, C, L = (3, 5, 7)
        x1 = torch.randn(B, C, L).to(device)
        mean1 = x1.sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) / (B * L)
        var1 = ((x1 - mean1) ** 2).sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) / (
            B * L
        )
        x2 = torch.randn(B, C, L).to(device)
        mean2 = x2.sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) / (B * L)
        var2 = ((x2 - mean2) ** 2).sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) / (
            B * L
        )

        mbn = MaskedBatchNorm1d(C)
        mbn = mbn.to(device)

        # Test without mask in train
        mbn.train()
        out = mbn(x1)
        self.assertTrue(mean1.allclose(mbn.running_mean))
        self.assertTrue(var1.allclose(mbn.running_var))
        normed = (x1 - mean1) / torch.sqrt(var1 + mbn.eps) * mbn.weight + mbn.bias
        self.assertTrue(normed.allclose(out))
        out = mbn(x2)
        normed = (x2 - mean2) / torch.sqrt(var2 + mbn.eps) * mbn.weight + mbn.bias
        self.assertTrue(normed.allclose(out))
        self.assertTrue(
            mbn.running_mean.allclose((1 - mbn.momentum) * mean1 + mbn.momentum * mean2)
        )
        self.assertTrue(
            mbn.running_var.allclose((1 - mbn.momentum) * var1 + mbn.momentum * var2)
        )

        # Without mask in eval
        mbn.eval()
        out = mbn(x1)
        self.assertTrue(
            mbn.running_mean.allclose((1 - mbn.momentum) * mean1 + mbn.momentum * mean2)
        )
        self.assertTrue(
            mbn.running_var.allclose((1 - mbn.momentum) * var1 + mbn.momentum * var2)
        )
        normed = (x1 - mbn.running_mean) / torch.sqrt(
            mbn.running_var + mbn.eps
        ) * mbn.weight + mbn.bias
        self.assertTrue(normed.allclose(out))

        # Check that masking with all ones doesn't change values
        mask = x1.new_ones((B, 1, L))
        outm = mbn(x1, input_mask=mask)
        self.assertTrue(outm.allclose(out))
        mbn.eval()
        out = mbn(x2)
        outm = mbn(x2, input_mask=mask)
        self.assertTrue(outm.allclose(out))

        # With mask in train
        mask = torch.randn(B, 1, L)
        mask = mask > 0.0
        mask = mask.to(device)
        n = mask.sum()
        mean1 = (x1 * mask).sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) / n
        var1 = (((x1 * mask) - mean1) ** 2).sum(dim=0, keepdim=True).sum(
            dim=2, keepdim=True
        ) / n
        mbn = MaskedBatchNorm1d(C)
        mbn = mbn.to(device)
        mbn.train()
        out = mbn(x1, input_mask=mask)
        self.assertTrue(mean1.allclose(mbn.running_mean))
        self.assertTrue(var1.allclose(mbn.running_var))
        normed = (x1 * mask - mean1) / torch.sqrt(
            var1 + mbn.eps
        ) * mbn.weight + mbn.bias
        self.assertTrue(normed.allclose(out))
        # With mask in eval
        mbn.eval()
        out = mbn(x1, input_mask=mask)
        normed = (x1 * mask - mbn.running_mean) / torch.sqrt(
            mbn.running_var + mbn.eps
        ) * mbn.weight + mbn.bias
        self.assertTrue(normed.allclose(out))