##
# @file   discrete_spectral_transform.py
# @author Yibo Lin
# @date   Jun 2018
#

import os
import sys
import numpy as np
import torch
import torch.nn.functional as F
import pdb
import dreamplace.ops.dct.torch_fft_api as torch_fft_api

""" Discrete spectral transformation leveraging fast fourier transform engine.
The math here mainly uses Prosthaphaeresis properties.
The trigonometric identities exploited by prosthaphaeresis relate products of trigonometric functions to sums.
sin(a) sin(b) = 1/2 * (cos(a-b) - cos(a+b))
cos(a) cos(b) = 1/2 * (cos(a-b) + cos(a+b))
sin(a) cos(b) = 1/2 * (sin(a+b) + sin(a-b))
cos(a) sin(b) = 1/2 * (sin(a-b) - sin(a+b))

A 2D FFT performs
y_{u, v} = \sum_i \sum_j x_{i, j} exp(-j*2*pi*u*i/M) exp(-j*2*pi*v*j/N)
         = \sum_i \sum_j x_{i, j} exp(-j*2*pi*(u*i/M + v*j/N))
         = \sum_i \sum_j x_{i, j} (cos(-2*pi*(u*i/M + v*j/N)) + j sin(-2*pi*(u*i/M + v*j/N))).

By mapping the original image from (i, j) to (i, N-j), we can have (u*i/M - v*j/N) inside exp.
This will enable us to derive various cos/sin transformation by computing FFT twice.
"""

def get_expk(N, dtype, device):
    """ Compute 2*exp(-1j*pi*u/(2N)), but not exactly the same.
    The actual return is 2*cos(pi*u/(2N)), 2*sin(pi*u/(2N)).
    This will make later multiplication easier.
    """
    pik_by_2N = torch.arange(N, dtype=dtype, device=device)
    pik_by_2N.mul_(np.pi/(2*N))
    # cos, sin
    # I use sin because the real part requires subtraction
    # this will be easier for multiplication
    expk = torch.stack([pik_by_2N.cos(), pik_by_2N.sin()], dim=-1)
    expk.mul_(2)

    return expk.contiguous()


def get_expkp1(N, dtype, device):
    """ Compute 2*exp(-1j*pi*(u+1)/(2N)), but not exactly the same.
    The actual return is 2*cos(pi*(u+1)/(2N)), 2*sin(pi*(u+1)/(2N))
    """
    neg_pik_by_2N = torch.arange(1, N+1, dtype=dtype, device=device)
    neg_pik_by_2N.mul_(np.pi/(2*N))
    # sin, -cos
    # I swap -cos and sin because we need the imag part
    # this will be easier for multiplication
    expk = torch.stack([neg_pik_by_2N.cos(), neg_pik_by_2N.sin()], dim=-1)
    expk.mul_(2)

    return expk.contiguous()


def get_exact_expk(N, dtype, device):
    # Compute exp(-j*pi*u/(2N)) = cos(pi*u/(2N)) - j * sin(pi*u/(2N))
    pik_by_2N = torch.arange(N, dtype=dtype, device=device)
    pik_by_2N.mul_(np.pi/(2*N))
    # cos, -sin
    expk = torch.stack([pik_by_2N.cos(), -pik_by_2N.sin()], dim=-1)
    return expk.contiguous()


def get_perm(N, dtype, device):
    """ Compute permutation to generate following array
    0, 2, 4, ..., 2*(N//2)-2, 2*(N//2)-1, 2*(N//2)-3, ..., 3, 1
    """
    perm = torch.zeros(N, dtype=dtype, device=device)
    perm[0:(N-1)//2+1] = torch.arange(0, N, 2, dtype=dtype, device=device)
    perm[(N-1)//2+1:] = torch.arange(2*(N//2)-1, 0, -2, dtype=dtype, device=device)

    return perm

def dct_2N(x, expk=None):
    """ Batch Discrete Cosine Transformation without normalization to coefficients.
    Compute y_u = \sum_i  x_i cos(pi*(2i+1)*u/(2N)),
    Impelements the 2N padding trick to solve DCT with FFT in the following link,
    https://dsp.stackexchange.com/questions/2807/fast-cosine-transform-via-fft

    1. Pad x by zeros
    2. Perform FFT
    3. Multiply by 2*exp(-1j*pi*u/(2N))
    4. Extract the real part
    """
    # last dimension
    N = x.size(-1)
    # pad last dimension
    x_pad = F.pad(x, (0, N), 'constant', 0)

    # the last dimension here becomes -2 because complex numbers introduce a new dimension
    y = torch_fft_api.rfft(x_pad, signal_ndim=1, normalized=False, onesided=True)[..., 0:N, :]
    y.mul_(1.0/N)