# AUTOGENERATED! DO NOT EDIT! File to edit: Fourier.ipynb (unless otherwise specified).

__all__ = ['FourierModel', 'Seasonal']

# Cell
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

torch.Tensor.ndim = property(lambda x: x.dim())

# Cell
_TWOPI = 2*np.pi

class FourierModel(nn.Module):
    """
    Block that outputs sin's and cos' as basis functions.
    parameters:
    - p: period
    - scale: time axis scaling factor
    - n: number of fourier components
    """
    def __init__(self, p:float=365.25, scale:float=1, n:int=7):
        super().__init__()
        self.np = [(i+1, p/scale) for i in range(n)]
        self.n, self.p = n, p
        if n > 0:
            self.linear = nn.Linear(n * 2, 1, bias=False)

    def forward(self, t:torch.Tensor):
        if len(self.np) > 0:
            cos = [torch.cos(_TWOPI * n * t / p) for n,p in self.np]
            sin = [torch.sin(_TWOPI * n * t / p) for n,p in self.np]

            return self.linear(torch.cat(cos + sin, dim=1))

        else:
            return 0

    def plot(self, ax=None, title=None):
        if self.n > 0:
            t = torch.linspace(0, self.p, steps=100)
            y = self.forward(t[:,None])
            if ax is None:
                fig, ax = plt.subplots(figsize=(12,5))
            ax.plot(t, y)
            ax.set_xlabel('Days')


class Seasonal(nn.Module):
    """
    Combination of Fourier Blocks that gives weekly, monthly and
    yearly seasonality.
    parameters:
    - y_n: Number of yearly components
    - m_n: Number of monthly components
    - w_n: Number of weekly components
    - scale: time axis scaling factor
    """
    def __init__(self, y_n=10, m_n=5, w_n=3, scale=1):
        super().__init__()
        self.y_n, self.m_n, self.w_n = y_n, m_n, w_n
        y_p=365.25; m_p=30.5; w_p=7
        self.yearly = FourierModel(y_p, scale, y_n) # , w[:,:idxs[0]]
        self.monthly = FourierModel(m_p, scale, m_n) # w[:,idxs[0]:idxs[1]]
        self.weekly = FourierModel(w_p, scale, w_n) # w[:,idxs[1]:idxs[2]]

    def forward(self, t):
        return self.yearly(t) + self.monthly(t) + self.weekly(t)

    def plot_components(self):
        plot = [True if c>0 else False for c in
                [self.y_n, self.m_n, self.w_n]]
        period = [self.yearly, self.monthly, self.weekly]
        period = [p for (p, pl) in zip(period, plot) if pl]
        fig, axs = plt.subplots(sum(plot), 1, figsize=(12, 12))

        for p, ax in zip(period, axs):
            p.plot(ax)
