"""
Conv1D forecaster used by daltoolboxdp via reticulate.

R entry points (see R/ts_conv1d.R):
  - ts_conv1d_create(in_channels, input_dim)
  - ts_conv1d_fit(model, df_train, n_epochs, lr)
  - ts_conv1d_predict(model, df_test)

Data expectations:
  - df_train/df_test contain lagged features and a 't0' target column.
  - Shapes are adjusted to (batch, channels, sequence_len) for Conv1D.
"""

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import TensorDataset
import torch.nn.functional as F
import sys
import functools
import operator

class TsConv1DNet(nn.Module):
    def __init__(self, in_channels, input_dim):
        super(TsConv1DNet,self).__init__()
        ksize = 2
        if (input_dim == 1): 
          ksize = 1
        
        self.feature_extractor = nn.Sequential(
            nn.Conv1d(in_channels = in_channels, out_channels = 64, kernel_size = ksize),
            nn.ReLU(inplace=True)
        )
        
        # https://datascience.stackexchange.com/questions/40906/determining-size-of-fc-layer-after-conv-layer-in-pytorch
        num_features_before_fcnn = functools.reduce(operator.mul, list(self.feature_extractor(torch.rand(1, input_dim)).shape))
        
        self.regressor = nn.Sequential(
            nn.Linear(in_features=num_features_before_fcnn, out_features=50),
            nn.ReLU(inplace=True),
            nn.Linear(50, 1)
        )

    def forward(self,x):
        out = self.feature_extractor(x)
        out = nn.Flatten(1, -1)(out)
        out = self.regressor(out)
        return out
      

def ts_conv1d_create(in_channels, input_dim):
  """Create Conv1D model given number of channels and input length (called from R)."""
  in_channels = int(in_channels)
  input_dim = int(input_dim)
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  model = TsConv1DNet(in_channels, input_dim).to(device)
  return(model)  


def ts_conv1d_train(epochs, lr, model, train_loader, opt_func=torch.optim.SGD):
  # Internal training routine; returns (model, avg_train_losses)
  
  train_losses = []
  # to track the average training loss per epoch as the model trains
  avg_train_losses = []
  
  criterion = nn.MSELoss()
  last_error = sys.float_info.max
  last_epoch = 0
  
  convergency = epochs/10
  if (convergency < 100):
    convergency = 100
  
  optimizer = opt_func(model.parameters(), lr)
  for epoch in range(epochs):
    #####
    # train the model #
    #####
    model.train() # prep model for training
    for data, target in train_loader:
      # clear the gradients of all optimized variables
      model.zero_grad()
      # forward pass: compute predicted outputs by passing inputs to the model
      device = next(model.parameters()).device
      output = model(data.float().to(device))
      target = target.to(device)
            
      # calculate the loss
      loss = criterion(output, target.float())
      
      # backward pass: compute gradient of the loss with respect to model parameters
      loss.backward()
      # perform a single optimization step (parameter update)
      optimizer.step()
      # record training loss
      train_losses.append(loss.item())
    
    # validate the model #
    model.eval() # prep model for evaluation
    
    # calculate average loss over an epoch
    train_loss = np.average(train_losses)
    avg_train_losses.append(train_loss)
    
    if ((last_error - train_loss) > 0.001):
      last_error = train_loss
      last_epoch = epoch

    # clear lists to track next epoch
    train_losses = []
    
    if (train_loss == 0):
      break
        
    if (epoch - last_epoch > convergency):
      break

  return model, avg_train_losses


def ts_conv1d_fit(model, df_train, n_epochs = 3000, lr = 0.001):
  """Entry from R to fit the Conv1D model using df_train with column 't0'."""
  n_epochs = int(n_epochs)
  
  X_train = df_train.drop('t0', axis=1).to_numpy()
  y_train = df_train.t0.to_numpy()
  
  X_train = X_train[:, :, np.newaxis]
  y_train = y_train[:, np.newaxis]
  
  train_x = torch.from_numpy(X_train)
  train_y = torch.from_numpy(y_train)
  
  train_x = torch.permute(train_x, (0, 2, 1))

  train_ds = TensorDataset(train_x, train_y)
  
  BATCH_SIZE = 8
  train_loader = torch.utils.data.DataLoader(train_ds, batch_size = BATCH_SIZE, shuffle = False)
  
  model = model.float()
  model, train_loss = ts_conv1d_train(n_epochs, lr, model, train_loader, opt_func=torch.optim.Adam)

  return model


def ts_conv1d_predict(model, df_test):
  """Entry from R to predict using the trained Conv1D; returns a 1D numpy array."""
  X_test = df_test.drop('t0', axis=1).to_numpy()
  y_test = df_test.t0.to_numpy()
  
  X_test = X_test[:, :, np.newaxis]
  y_test = y_test[:, np.newaxis]
  
  test_x = torch.from_numpy(X_test)
  test_y = torch.from_numpy(y_test)
  
  test_x = torch.permute(test_x, (0, 2, 1))	

  test_ds = TensorDataset(test_x, test_y)
  
  BATCH_SIZE = 8
  test_loader = torch.utils.data.DataLoader(test_ds, batch_size = BATCH_SIZE, shuffle = False)
  
  outputs = []
  with torch.no_grad():
    for xb, yb in test_loader:
      device = next(model.parameters()).device
      output = model(xb.float().to(device))
      outputs.append(output)
  
  test_predictions = torch.vstack(outputs).squeeze(1)  
  test_predictions = test_predictions.flatten().cpu().numpy()
  
  return test_predictions
