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


def seasonality_model(thetas, t, device):
    p = thetas.size()[-1]
    assert p < 10, 'thetas_dim is too big.'
    p1, p2 = (p // 2, p // 2) if p % 2 == 0 else (p // 2, p // 2 + 1)
    s1 = torch.tensor([np.cos(2 * np.pi * i * t) for i in range(p1)]).float()  # H/2-1
    s2 = torch.tensor([np.sin(2 * np.pi * i * t) for i in range(p2)]).float()
    S = torch.cat([s1, s2])
    return thetas.mm(S.to(device))


def trend_model(thetas, t, device):
    p = thetas.size()[-1]
    assert p <= 4, 'thetas_dim is too big.'
    T = torch.tensor([t ** i for i in range(p)]).float()
    return thetas.mm(T.to(device))


def linspace(backcast_length, forecast_length):
    lin_space = np.linspace(-backcast_length, forecast_length, backcast_length + forecast_length)
    b_ls = lin_space[:backcast_length]
    f_ls = lin_space[backcast_length:]
    return b_ls, f_ls


class Block(nn.Module):

    def __init__(self, units, thetas_dim, device, backcast_length=10, forecast_length=5, share_thetas=False):
        super(Block, self).__init__()
        self.units = units
        self.thetas_dim = thetas_dim
        self.backcast_length = backcast_length
        self.forecast_length = forecast_length
        self.share_thetas = share_thetas
        self.fc1 = nn.Linear(backcast_length, units)
        self.fc2 = nn.Linear(units, units)
        self.fc3 = nn.Linear(units, units)
        self.fc4 = nn.Linear(units, units)
        self.device = device
        self.backcast_linspace, self.forecast_linspace = linspace(backcast_length, forecast_length)
        if share_thetas:
            self.theta_f_fc = self.theta_b_fc = nn.Linear(units, thetas_dim)
        else:
            self.theta_b_fc = nn.Linear(units, thetas_dim)
            self.theta_f_fc = nn.Linear(units, thetas_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x.to(self.device)))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        return x

    def __str__(self):
        block_type = type(self).__name__
        return f'{block_type}(units={self.units}, thetas_dim={self.thetas_dim}, ' \
            f'backcast_length={self.backcast_length}, forecast_length={self.forecast_length}, ' \
            f'share_thetas={self.share_thetas}) at @{id(self)}'


class SeasonalityBlock(Block):

    def __init__(self, units, thetas_dim, device, backcast_length=10, forecast_length=5):
        super(SeasonalityBlock, self).__init__(units, thetas_dim, device, backcast_length,
                                               forecast_length, share_thetas=True)

    def forward(self, x):
        x = super(SeasonalityBlock, self).forward(x)
        backcast = seasonality_model(self.theta_b_fc(x), self.backcast_linspace, self.device)
        forecast = seasonality_model(self.theta_f_fc(x), self.forecast_linspace, self.device)
        return backcast, forecast


class TrendBlock(Block):

    def __init__(self, units, thetas_dim, device, backcast_length=10, forecast_length=5):
        super(TrendBlock, self).__init__(units, thetas_dim, device, backcast_length,
                                         forecast_length, share_thetas=True)

    def forward(self, x):
        x = super(TrendBlock, self).forward(x)
        backcast = trend_model(self.theta_b_fc(x), self.backcast_linspace, self.device)
        forecast = trend_model(self.theta_f_fc(x), self.forecast_linspace, self.device)
        return backcast, forecast


class GenericBlock(Block):

    def __init__(self, units, thetas_dim, device, backcast_length=10, forecast_length=5):
        super(GenericBlock, self).__init__(units, thetas_dim, device, backcast_length, forecast_length)

        self.backcast_fc = nn.Linear(thetas_dim, backcast_length)
        self.forecast_fc = nn.Linear(thetas_dim, forecast_length)

    def forward(self, x):
        # no constraint for generic arch.
        x = super(GenericBlock, self).forward(x)

        theta_b = F.relu(self.theta_b_fc(x))
        theta_f = F.relu(self.theta_f_fc(x))

        backcast = self.backcast_fc(theta_b)  # generic. 3.3.
        forecast = self.forecast_fc(theta_f)  # generic. 3.3.

        return backcast, forecast


class NBeatsNet(nn.Module):
    SEASONALITY_BLOCK = 'seasonality'
    TREND_BLOCK = 'trend'
    GENERIC_BLOCK = 'generic'

    def __init__(self,
                 device,
                 stack_types=[TREND_BLOCK, SEASONALITY_BLOCK],
                 nb_blocks_per_stack=3,
                 forecast_length=5,
                 backcast_length=10,
                 thetas_dims=[4, 8],
                 share_weights_in_stack=False,
                 hidden_layer_units=256, ):
        super(NBeatsNet, self).__init__()
        self.forecast_length = forecast_length
        self.backcast_length = backcast_length
        self.hidden_layer_units = hidden_layer_units
        self.nb_blocks_per_stack = nb_blocks_per_stack
        self.share_weights_in_stack = share_weights_in_stack
        self.stack_types = stack_types
        self.stacks = []
        self.thetas_dim = thetas_dims
        self.parameters = []
        self.device = device
        print(f'| N-Beats')
        for stack_id in range(len(self.stack_types)):
            self.stacks.append(self.create_stack(stack_id))
        self.parameters = nn.ParameterList(self.parameters)
        self.to(self.device)

    def create_stack(self, stack_id):
        stack_type = self.stack_types[stack_id]
        print(f'| --  Stack {stack_type.title()} (#{stack_id}) (share_weights_in_stack={self.share_weights_in_stack})')
        blocks = []
        for block_id in range(self.nb_blocks_per_stack):
            block_init = NBeatsNet.select_block(stack_type)
            if self.share_weights_in_stack and block_id != 0:
                block = blocks[-1]  # pick up the last one to make the
            else:
                block = block_init(self.hidden_layer_units, self.thetas_dim[stack_id],
                                   self.device, self.backcast_length, self.forecast_length)
                self.parameters.extend(block.parameters())
            print(f'     | -- {block}')
            blocks.append(block)
        return blocks

    @staticmethod
    def select_block(block_type):
        if block_type == NBeatsNet.SEASONALITY_BLOCK:
            return SeasonalityBlock
        elif block_type == NBeatsNet.TREND_BLOCK:
            return TrendBlock
        else:
            return GenericBlock

    def forward(self, backcast):
        forecast = torch.zeros(size=(backcast.size()[0], self.forecast_length,))  # maybe batch size here.
        for stack_id in range(len(self.stacks)):
            for block_id in range(len(self.stacks[stack_id])):
                b, f = self.stacks[stack_id][block_id](backcast)
                backcast = backcast.to(self.device) - b
                forecast = forecast.to(self.device) + f
        return backcast, forecast

[((13, 8, 13, 27), 'torch.cat', 'torch.cat', ({(13, 18, 13, 26): '[s1, s2]'}, {}), '([s1, s2])', False, 'import torch\n'), ((25, 16, 25, 97), 'numpy.linspace', 'np.linspace', ({(25, 28, 25, 44): '-backcast_length', (25, 46, 25, 61): 'forecast_length', (25, 63, 25, 96): 'backcast_length + forecast_length'}, {}), '(-backcast_length, forecast_length, backcast_length +\n    forecast_length)', True, 'import numpy as np\n'), ((40, 19, 40, 52), 'torch.nn.Linear', 'nn.Linear', ({(40, 29, 40, 44): 'backcast_length', (40, 46, 40, 51): 'units'}, {}), '(backcast_length, units)', False, 'from torch import nn\n'), ((41, 19, 41, 42), 'torch.nn.Linear', 'nn.Linear', ({(41, 29, 41, 34): 'units', (41, 36, 41, 41): 'units'}, {}), '(units, units)', False, 'from torch import nn\n'), ((42, 19, 42, 42), 'torch.nn.Linear', 'nn.Linear', ({(42, 29, 42, 34): 'units', (42, 36, 42, 41): 'units'}, {}), '(units, units)', False, 'from torch import nn\n'), ((43, 19, 43, 42), 'torch.nn.Linear', 'nn.Linear', ({(43, 29, 43, 34): 'units', (43, 36, 43, 41): 'units'}, {}), '(units, units)', False, 'from torch import nn\n'), ((97, 27, 97, 65), 'torch.nn.Linear', 'nn.Linear', ({(97, 37, 97, 47): 'thetas_dim', (97, 49, 97, 64): 'backcast_length'}, {}), '(thetas_dim, backcast_length)', False, 'from torch import nn\n'), ((98, 27, 98, 65), 'torch.nn.Linear', 'nn.Linear', ({(98, 37, 98, 47): 'thetas_dim', (98, 49, 98, 64): 'forecast_length'}, {}), '(thetas_dim, forecast_length)', False, 'from torch import nn\n'), ((141, 26, 141, 59), 'torch.nn.ParameterList', 'nn.ParameterList', ({(141, 43, 141, 58): 'self.parameters'}, {}), '(self.parameters)', False, 'from torch import nn\n'), ((47, 48, 47, 76), 'torch.nn.Linear', 'nn.Linear', ({(47, 58, 47, 63): 'units', (47, 65, 47, 75): 'thetas_dim'}, {}), '(units, thetas_dim)', False, 'from torch import nn\n'), ((49, 30, 49, 58), 'torch.nn.Linear', 'nn.Linear', ({(49, 40, 49, 45): 'units', (49, 47, 49, 57): 'thetas_dim'}, {}), '(units, thetas_dim)', False, 'from torch import nn\n'), ((50, 30, 50, 58), 'torch.nn.Linear', 'nn.Linear', ({(50, 40, 50, 45): 'units', (50, 47, 50, 57): 'thetas_dim'}, {}), '(units, thetas_dim)', False, 'from torch import nn\n'), ((11, 23, 11, 48), 'numpy.cos', 'np.cos', ({(11, 30, 11, 47): '2 * np.pi * i * t'}, {}), '(2 * np.pi * i * t)', True, 'import numpy as np\n'), ((12, 23, 12, 48), 'numpy.sin', 'np.sin', ({(12, 30, 12, 47): '2 * np.pi * i * t'}, {}), '(2 * np.pi * i * t)', True, 'import numpy as np\n')]