from __future__ import absolute_import
from __future__ import division

from peachpy import *
from peachpy.x86_64 import *
from common import butterfly, sqrt2_over_2


from common import butterfly, sqrt2_over_2, cos_npi_over_8, interleave


def fft8_bitreverse(n):
    return int(format(n, "03b")[::-1], 2)


def load_ymm_variable(variable):
    assert isinstance(variable, (YMMRegister, LocalVariable))
    ymm_variable = variable
    if isinstance(variable, LocalVariable):
        assert variable.size == YMMRegister.size
        ymm_variable = YMMRegister()
        VMOVAPS(ymm_variable, variable)
    return ymm_variable


def store_ymm_result(variable, result):
    assert isinstance(result, YMMRegister)

    if isinstance(variable, YMMRegister):
        SWAP.REGISTERS(variable, result)
    else:
        VMOVAPS(variable, result)


def forward_vfft(reg_t0, reg_t8, reg_t_stride, data_out, reg_row_start=None, reg_row_end=None, ymm_load_mask=None):
    assert isinstance(reg_t0, GeneralPurposeRegister64)
    assert isinstance(reg_t8, GeneralPurposeRegister64)
    assert isinstance(reg_t_stride, GeneralPurposeRegister64)
    assert isinstance(data_out, list) and len(data_out) == 16
    assert ymm_load_mask is None or isinstance(ymm_load_mask, YMMRegister)

    out_real, out_imag = data_out[0::2], data_out[1::2]

    real, imag = [YMMRegister() for _ in range(8)], [YMMRegister() for _ in range(8)]
    imag[0] = LocalVariable(YMMRegister.size)
    imag[4] = LocalVariable(YMMRegister.size)
    data = interleave(real, imag)

    for i, (data_lo, data_hi) in enumerate(zip(data[0:8], data[8:16])):
        row_lo = i
        row_hi = row_lo + 8

        ymm_data_lo, ymm_data_hi = data_lo, data_hi
        if isinstance(data_lo, LocalVariable):
            ymm_data_lo = YMMRegister()
        if isinstance(data_hi, LocalVariable):
            ymm_data_hi = YMMRegister()

        VXORPS(ymm_data_lo, ymm_data_lo, ymm_data_lo)
        skip_data_lo = Label()
        if reg_row_start:
            CMP(reg_row_start, row_lo)
            JA(skip_data_lo)
        if reg_row_end:
            CMP(reg_row_end, row_lo)
            JBE(skip_data_lo)
        if ymm_load_mask is None:
            VMOVUPS(ymm_data_lo, [reg_t0])
        else:
            VMASKMOVPS(ymm_data_lo, ymm_load_mask, [reg_t0])
        if i + 1 != 8:
            ADD(reg_t0, reg_t_stride)
        LABEL(skip_data_lo)

        VMOVAPS(ymm_data_hi, ymm_data_lo)