
import os
import sys
sys.path.append(os.getcwd())

import numpy as np
import torch
import flow
from utils import cdfDiscreteLogitstic, cdfMixDiscreteLogistic
from utils import logDiscreteLogistic, logMixDiscreteLogistic

nbins = 4096
_bins = torch.arange(-nbins // 2, nbins // 2).reshape(-1, 1, 1, 1, 1)
decimal = flow.ScalingNshifting(256, -128)


def test_disLogisticCDF():
    logscale = torch.tensor(
        [[[[-3.6826, -3.0157, -3.6032],
          [-3.7063, -3.0269, -3.5338],
          [-3.5311, -2.9907, -3.3516],
          [-3.9300, -3.3121, -3.8110]],

         [[-3.1022, -3.0692, -3.2039],
          [-2.9466, -3.0006, -3.2969],
          [-2.7636, -2.5691, -2.9628],
          [-3.3657, -3.2948, -3.5318]],

         [[-3.9748, -3.0670, -3.2399],
          [-3.9312, -3.0055, -3.1729],
          [-3.8588, -2.9139, -3.1794],
          [-4.1534, -3.2404, -3.5665]]]]
    )

    mean = torch.tensor(
        [[[[ 0.0191,  0.0459,  0.0131],
          [-0.0059,  0.0254, -0.0100],
          [ 0.0359,  0.0406,  0.0242],
          [ 0.0331,  0.0438,  0.0255]],

         [[ 0.0214,  0.0502,  0.0622],
          [ 0.0371,  0.0368,  0.0517],
          [ 0.0217,  0.0855,  0.0874],
          [ 0.0144,  0.0475,  0.0470]],

         [[-0.0602, -0.0791, -0.0784],
          [-0.0443, -0.0765, -0.0701],
          [-0.0654, -0.0709, -0.0788],
          [-0.0608, -0.0721, -0.0688]]]]
    )

    bins = _bins - 1 + torch.round(decimal.forward_(mean))

    cdf = cdfDiscreteLogitstic(bins, mean, logscale, decimal=decimal).detach().numpy()

    pList = []
    for i in range(bins.shape[0]):
        logp = logDiscreteLogistic(bins[i: i + 1], mean, logscale, decimal=decimal).detach().numpy()
        pList.append(np.exp(logp).reshape(mean.shape))
    pList = np.array(pList)

    _cdf = np.cumsum(pList, 0)

    assert np.allclose(cdf, _cdf)


def test_mixDixLogisticCDF():
    mean = torch.tensor(
        [[[[-0.2414,  0.2089, -0.0209, -0.1279]],
         [[ 0.7791,  0.1031,  0.0940,  0.1678]],
         [[ 0.0095,  0.0391, -0.0318, -0.2183]]],
        [[[-0.1466,  0.2090, -0.0594, -0.0837]],
         [[ 0.8711,  0.0540,  0.0940,  0.0859]],
         [[-0.0683, -0.0204, -0.0340, -0.0587]]],
        [[[-0.1994, -0.0442, -0.0307, -0.0823]],
         [[ 1.0158,  0.0636,  0.0832,  0.0717]],
         [[-0.1863, -0.0177, -0.0293, -0.0708]]],
        [[[-0.3517,  0.1062, -0.0362, -0.1661]],
         [[ 0.6567,  0.1452,  0.0294,  0.0864]],
         [[-0.1384, -0.0171, -0.0195, -0.0710]]],
        [[[-0.3158,  0.2068,  0.1114, -0.1251]],
         [[ 0.5600,  0.1987,  0.1891,  0.1754]],
         [[-0.2758, -0.1032, -0.0435, -0.1156]]]])
    logscale = torch.tensor(
        [[[[-3.1292, -4.0168, -3.2886, -2.5948]],
         [[-2.8226, -2.3489, -2.8613, -2.3892]],
         [[-3.3502, -3.4929, -2.9572, -2.7060]]],
        [[[-3.4556, -4.0166, -2.7471, -3.1203]],
         [[-2.6906, -3.6062, -2.8620, -3.0673]],
         [[-3.2775, -3.3661, -3.2897, -4.0553]]],
        [[[-3.4652, -3.3828, -3.3053, -3.6945]],
         [[-2.7657, -2.9172, -3.4067, -3.7734]],
         [[-3.4817, -3.0397, -2.8021, -3.1398]]],
        [[[-2.7246, -3.7798, -4.1237, -2.8605]],
         [[-3.0524, -2.6628, -2.4833, -3.0913]],
         [[-4.0249, -3.8364, -3.7608, -2.7111]]],
        [[[-3.5460, -4.0208, -2.9837, -3.1288]],
         [[-3.2062, -2.1702, -2.2238, -2.6122]],
         [[-3.1754, -3.0892, -2.3359, -2.4321]]]])
    mixing = torch.tensor(
        [[[[ 1.3161,  0.8664,  1.7648, -0.7598, -0.8658],
          [-3.7472, -3.6553,  5.2783,  0.2242, -3.6304],
          [-0.7378,  0.2730,  1.8044,  0.7450, -1.6218],
          [-0.8105,  1.8833,  1.8243, -0.7879, -1.1211]]],
        [[[ 1.3952, -0.8232, -1.0135,  1.8041,  0.9846],
          [-0.4372,  1.1296,  1.5473, -0.0661, -0.5995],
          [-0.5167,  1.5559,  1.2607, -0.3227, -0.8687],
          [-0.6226,  1.5024,  1.4221,  1.4741, -0.4409]]],
        [[[ 1.3045,  1.8551,  0.1755, -0.6253, -1.2045],
          [-0.9858,  1.5529, -0.6332,  1.4569, -1.1089],
          [-0.5954,  1.2305,  1.4068,  0.7919, -0.3811],
          [-0.2997,  0.6804,  2.0660,  1.1353, -0.9155]]]])

    bins = _bins - 1 + torch.round(decimal.forward_(mean.permute([1, 2, 3, 0])) * mixing).sum(-1).reshape(1, *mean.shape[1:])
    cdf = cdfMixDiscreteLogistic(bins, mean, logscale, mixing, decimal=decimal)

    pList = []
    for i in range(bins.shape[0]):
        logp = logMixDiscreteLogistic(bins[i: i + 1], mean, logscale, mixing, decimal=decimal).detach().numpy()
        pList.append(np.exp(logp).reshape(logp.shape[1:]))
    pList = np.array(pList)

    _cdf = np.cumsum(pList, 0)

    assert np.allclose(cdf, _cdf)


if __name__ == "__main__":

    test_disLogisticCDF()
    test_mixDixLogisticCDF()
torch.arange torch.tensor