import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class Scaling1D(nn.Module):
    def __init__(self):
        super(Scaling1D, self).__init__()
    
    def forward(self, x, a):
        """Performs spatial dilation of a periodic function by Fourier 
        transform. Uses the observation that the Fourier transform has the 
        following time scaling relation:

        f(x) <-----> F(k)

        f(ax) <----> 1/|a| F(k/a)

        Args:
            x (torch.tensor): Observations to be scaled. A tensor of shape 
            (n_batch, n_channels, n_grid_points). The scaling will be performed
            across the last dimension.
            a (float): Scaling parameter. x -> lambda * x
        """

        a_ = int(a)

        # Take Fourier transform
        x_dft = torch.fft.fft(x)

        # Scale the frequencies k -> k/a
        # Idea: we want a new DFT vector which has shape x_dft.shape[-1] * a
        # In this vector we want zeros for non-integer values of k/a and the 
        # elements of x_dft for the integer values of k/a. We will do this via
        # 1d convolution with a kernel of zeros except for the first entry

        convo_kernel = torch.zeros(x_dft.shape[0], 1, a_, 
                                    dtype=torch.cfloat)
        convo_kernel[:, 0, 0] = torch.ones_like(convo_kernel[:, 0, 0])
        convo_kernel.requires_grad = False

        print(convo_kernel)

        out_x_dft = torch.zeros((x_dft.shape[0], 
                                    x_dft.shape[1], 
                                    a_ * x_dft.shape[-1]), 
                                dtype=torch.cfloat)
        # print(out_x_dft.shape)
        for i in range(x_dft.shape[-1]):
            lb = a_ * i
            ub = a_ * (i + 1)

            arg_1 = x_dft[:, :, i].view(x_dft.shape[0], x_dft.shape[1], 1)

            # ignoring the batch dimension, convo_kernel is of size (1, a)
            # and arg_1 is of size (n_channels, 1)
            out_x_dft[:, :, lb:ub] = torch.matmul(arg_1, convo_kernel)

        x_dft = out_x_dft
        # print(x_dft.shape)
        

        # Scale the whole thing by 1/|a|
        x = torch.divide(x_dft, np.abs(a))

        x_ifft = torch.fft.ifft(x)

        return x_ifft


