"""Test C function dealing with Leech lattice vectors mod 3 of type 2

In this script we test functions dealing with vectors of type 2 in the 
Leech lattice modulo 3. These functions are implemented in file 
gen_leech3.c and available in the extension mmgroup.generators.

We use the terminology defined in
the document *The C interface of the mmgroup project*, 
section *Description of the mmgroup.generators extension*.
"""

from __future__ import absolute_import, division, print_function
from __future__ import  unicode_literals



from random import randint, choices #, shuffle, sample
from functools import reduce
from operator import __or__
from numbers import Integral
from collections import defaultdict
from multiprocessing import Pool, TimeoutError

import numpy as np
import scipy
import scipy.special
from scipy.stats import chisquare
import pytest

from mmgroup import MM0
from mmgroup.mat24 import MAT24_ORDER, ploop_theta
from mmgroup.mat24 import bw24 as mat24_bw24
from mmgroup.generators import gen_leech3to2_short
from mmgroup.generators import gen_leech2to3_short
from mmgroup.generators import gen_leech3_neg
from mmgroup.generators import gen_leech2_type
from mmgroup.generators import gen_leech2_subtype
from mmgroup.generators import gen_leech2_type2

from mmgroup.mm import mm_aux_index_extern_to_sparse
from mmgroup.mm import mm_aux_index_sparse_to_leech2
 

from mmgroup.tests.test_gen_xi.test_gen_type4 import str_v3


#####################################################################
# Creating test vectors
#####################################################################




def type2_testvectors():
    """Yield short test vectors in the Leech lattice mod 2

    The function yields pairs ``(v2, vtype)``, where ``v2 ``is a short
    vector (of type 2) in the Leech lattice mod 2 in **Leech lattice
    encoding**, and ``vtype`` is the subtype of ``v``.

    This function uses function from file mm_aux.c for converting
    random indices of the rep 196884x to short vectors.
    """
    TEST_PARS = [
        (300, 2 * 276, 50, 0x20),
        (852, 759 * 64, 500, 0x22),
        (49428, 2048 * 24, 500, 0x21),
    ] 
    for (start, nvectors, ntests, vtype) in TEST_PARS:
         for i in range(ntests):
             v_extern = randint(start, start + nvectors - 1)
             v_sparse = mm_aux_index_extern_to_sparse(v_extern)
             v2 = mm_aux_index_sparse_to_leech2(v_sparse)
             assert v2 != 0, (v_extern, hex(v_sparse), hex(v2))
             #print(v_extern, hex(v_sparse), hex(v2), hex(vtype))
             yield v2, vtype


#####################################################################
# Test mapping of type 2 Leech vectors modulo 3 to vectors modulo 2
#####################################################################


def v3_to_v2(v3):
    """Map type-2 vector from Leech lattice mod 3 to Leech lattice mod 2

    Here parameter ``v3`` is a type-2 vector in the Leech lattice mod 2 
    in Leech lattice mod 3 encoding. 

    The function returns a type-2 vector in the Leech lattice mod 2
    corresponding to ``v3`` in Leech lattice encoding.

    The result is unique. The function returns 0 if ``v3`` is not of
    type 2 in the Leech lattice mod 3
    
    Tis function is a wrapper for the C function ``gen_leech3to2_short`` 
    in file ``gen_leech3.c``. 
    """
    result = gen_leech3to2_short(v3)
    assert result != 0, (str_v3(v3), weight_v3(v3), hex(result))
    return result

def v2_to_v3(v2):
    """Map type-2 vector from Leech lattice mod 2 to Leech lattice mod 3

    This function reverses the peration of function ``v3_to_v2``.
    The result is unique up to sign only.
    
    Tis function is a wrapper for the C function ``gen_leech2to3_short`` 
    in file ``gen_leech3.c``. 
    """
    result = gen_leech2to3_short(v2)
    assert result != 0, (hex(v2), hex(result))
    return result


@pytest.mark.gen_xi
def test_type3(verbose = 0):
    r"""Test conversion of type-2 vectors 

    The function converts random sshort vectors in the Leech
    lattice mod2 2 to vectors in the Leech lattice mod 3, and
    vice versa, and check that these conversions are consistent. 

    It also checks functions in module ``checks  
    """
    weights = defaultdict(int)
    for ntest, (v2, vtype) in enumerate(type2_testvectors()):
        if verbose:
            print("\nTEST %s" % (ntest+1))
            print("v2 = " , hex(v2))
        assert gen_leech2_subtype(v2) == vtype
        assert gen_leech2_type2(v2) == vtype

        v3 = v2_to_v3(v2) 
        assert v3_to_v2(v3) == v2
        v3n = gen_leech3_neg(v3)
        assert v3_to_v2(v3n) == v2
    
    
    
