#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun 10 15:59:12 2019

@author: adamreidsmith
"""

'''
Standard neural network with dropout.  Input is the peaks and corresponding 
frequencies of the Fourier transform of a solution to the Van der Pol equation
x'' - a*(1-x^2)*x' + b*x = f(t) there f is periodic or multiperiodic with a 
random phase that is encorporated into training.  The network predicts values 
of the parameters a and b given training data.

Datafiles for this network can be generated by running 'vdp_data.py'.
'''

#Path of the datafile created by 'vdp_data.py'.
file_path = './datafiles/vdp_data_800pts_[soln,FT,hist,phase(True),param].npy'

from os import path
assert path.exists(file_path), 'Datafile not found. Please run \'vdp_data.py\' \
                               to generate a datafile.'
                               
import torch
import numpy as np
from torch import nn
from scipy.signal import find_peaks
from torch.utils.data import Dataset, DataLoader, random_split
import matplotlib.pyplot as plt

###############################################################################
'''
Inputs:
    n_epochs:           Number of epochs to train for.
    
    batch_size:         Batch size.
    
    lr:                 Learning Rate.
    
    weight_decay:       Weight decay factor.
    
    lr_factor:          Learning rate decay factor.  Learning rate is multiplied 
                            by this factor every epoch.
                            
    loss_function:      Loss function. Can be:
                            'mean square':   loss = sum((x_i - y_i)^2)
                            'log cosh':      loss = sum(log(cosh(x_i - y_i)))
    
    include_phase:      Include or exclude the phase in the training of the 
                            neural network.
    
    num_peaks:          Number of peaks in Fourier transforms to include in training.
    
    inputs:             Values to use as inputs to the network. This is a list 
                            of strings each indicating a set of inputs to the 
                            network.  Valid entries are:
                                'refft':         Real part of the FFT.
                                'imfft':         Imaginary part of the FFT.
                                'absfft':        Absolute value of the FFT.
                                'refft_peaks':   Real part of peaks in FFT. Peaks 
                                                    are computed from absolute value 
                                                    of the FFT.
                                'imfft_peaks':   Imaginary part of peaks in FFT. 
                                                    Peaks are computed from absolute 
                                                    value of the FFT.
                                'absfft_peaks':  Peaks in the absolute value of the FFT.
                                'phase':         Phase(s) in the forcing function.
'''
###############################################################################

def main(n_epochs=150,
         batch_size=4,
         lr=0.001,
         lr_factor=0.98,
         weight_decay=1e-8,
         loss_function='log cosh',
         include_phase=True,
         num_peaks=6,
         inputs=['absfft_peaks', 'phase']):
    
    class Data(Dataset):
        
        def __init__(self):
            print('\nLoading data...')
            self.data = np.load(file_path)
            
            dlen = len(self.data)
            
            #Parameters used in solution of Van der Pol oscillator
            self.parameters = torch.Tensor([self.data[i] for i in range(dlen) if (i+1) % 5 == 0])
            
            #Phase phi included in the forcing function cos(wt+phi)
            self.phase = torch.Tensor([[self.data[i]] for i in range(dlen) if (i+2) % 5 == 0])
            
            #Time points at which the Van dar Pol equation was evaluated.
            self.time = [self.data[i] for i in range(dlen) if i % 5 == 0]
            self.time = [self.time[0][i][0] for i in range(len(self.time[0]))]
            
            #Tensor of complex values of Fourier transform
            self.fft = [self.data[i] for i in range(dlen) if (i+4) % 5 == 0]
            self.fft = [self.fft[i][:,1][:len(self.time)//2+1] for i in range(len(self.fft))]
            
            #Tensor of absolute value of Fourier transform values
            self.absfft = torch.Tensor(np.abs(self.fft))
            
            #Tensor of real part of Fourier transform values
            if 'refft' in inputs:
                self.refft = torch.Tensor(np.real(self.fft))
            
            #Tensor of imaginary part of Fourier transform values
            if 'imfft' in inputs:
                self.imfft = torch.Tensor(np.imag(self.fft))
            
            #Indices of peaks in absolute value of Fourier transform
            self.peaks_indices = [find_peaks(self.absfft[i])[0] for i in range(len(self.absfft))]
            
            #Indices of the top n peaks in each Fourier transform
            self.num_max_peaks = num_peaks  #Number of heighest peaks to consider
            self.max_n_peaks_indices = []
            peaks_indices = self.peaks_indices.copy()
            for i in range(len(self.peaks_indices)):
                for j in range(self.num_max_peaks):
                    key_func = lambda index: self.absfft[i][index]
                    try:
                        max_peak_index = max(peaks_indices[i], key=key_func)
                    except:
                        max_peak_index = None
                    if j == 0:
                        #appendee = [max_peak_index] if max_peak_index is not None else self.max_n_peaks_indices[-1][0]
                        self.max_n_peaks_indices.append([max_peak_index])
                    else:
                        appendee = max_peak_index if max_peak_index is not None else self.max_n_peaks_indices[-1][0]
                        self.max_n_peaks_indices[-1].append(appendee)                        
                    
                    index = np.argwhere(peaks_indices[i] == max_peak_index)
                    peaks_indices[i] = np.delete(peaks_indices[i], index)
            
            #Values and frequencies of top n peaks in each Fourier transform
            self.max_n_peaks = [[self.fft[i][j] for j in self.max_n_peaks_indices[i]] for i in range(len(self.fft))]
            self.max_n_peaks_time = [[self.time[j].item() for j in self.max_n_peaks_indices[i]] for i in range(len(self.fft))]
            
            self.len = self.parameters.shape[0]
            
        def __getitem__(self, index):            
            
            added_time = False
            items_to_cat = []
            for string in inputs:
                if string != 'phase':
                    if string == 'refft':
                        items_to_cat.append(self.refft[index])
                        
                    elif string == 'imfft':
                        items_to_cat.append(self.imfft[index])
                    
                    elif string == 'absfft':
                        items_to_cat.append(self.absfft[index])

                    elif string == 'refft_peaks':
                        items_to_cat.append(torch.Tensor(np.real(self.max_n_peaks[index])))
                        if not added_time:
                            items_to_cat.append(torch.Tensor(self.max_n_peaks_time[index]))
                            added_time = True
                            
                    elif string == 'imfft_peaks':
                        items_to_cat.append(torch.Tensor(np.imag(self.max_n_peaks[index])))
                        if not added_time:
                            items_to_cat.append(torch.Tensor(self.max_n_peaks_time[index]))
                            added_time = True
                    
                    elif string == 'absfft_peaks':
                        items_to_cat.append(torch.Tensor(np.abs(self.max_n_peaks[index])))
                        if not added_time:
                            items_to_cat.append(torch.Tensor(self.max_n_peaks_time[index]))
                            added_time = True
                    else:
                        raise RuntimeError('Undefined input string: ' + str(string))
            
            if len(items_to_cat) > 1:
                item = torch.cat(items_to_cat)
            else:
                item = items_to_cat[0]
                                                
            return [item, self.phase[index]], self.parameters[index]
            
        __len__ = lambda self: self.len
                       
    
    dataset = Data()
    
    # Lengths of the training and validation datasets
    train_len = int(0.75*dataset.len)
    valid_len = dataset.len - train_len
        
    #Randomly split the data into training and validation datasets
    train_data, valid_data = random_split(dataset, (train_len, valid_len))
        
    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size, shuffle=True)
   
    
    #Model of the neural network
    class Model(nn.Module):
        
        def __init__(self):
            super(Model, self).__init__()
            
            n_inputs = self.compute_ninputs(inputs)
            
            #Fully connected linear layers
            #self.dropout1 = nn.Dropout(p=0.1)
            self.fc1 = nn.Linear(in_features=n_inputs, out_features=500)
            self.dropout2 = nn.Dropout(p=0.1)
            self.fc2 = nn.Linear(in_features=500, out_features=50)
            self.dropout3 = nn.Dropout(p=0.1)
            self.fc3 = nn.Linear(in_features=50, out_features=2)
                       
        def forward(self, x, phi):
                        
            #Append phi to the input vector
            if 'phase' in inputs:
                x = torch.cat((x,phi.flatten(1)),1)
                        
            #Linear layers wih dropout
            #x = self.dropout1(x)
            x = torch.sigmoid(self.fc1(x))
            x = self.dropout2(x)
            x = torch.sigmoid(self.fc2(x))
            x = self.dropout3(x) 
            return self.fc3(x)
        
        def compute_ninputs(self,inputs):
            added_time = False
            n_inputs = 0
            for string in inputs:
                if string == 'phase':
                    n_inputs += len(dataset.phase[0][0])
                    
                elif string in ['refft', 'imfft', 'absfft']:
                    n_inputs += len(dataset.absfft[0])
                    
                elif string in ['refft_peaks', 'imfft_peaks', 'absfft_peaks']:
                    n_inputs += dataset.num_max_peaks
                    if not added_time:
                        n_inputs += dataset.num_max_peaks
                        added_time = True
                        
            return n_inputs

        
    model = Model()
    
    if loss_function == 'mean square':
        loss_func = nn.MSELoss()
    elif loss_function == 'log cosh':
        loss_func = lambda x, y: torch.log(torch.cosh(2*(x - y))).sum()
    else:
        raise RuntimeError('loss_function not recognized. \
                           Set loss_function to \'mean square\' or \'log cosh\'')
    
    #Optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=lr_factor)
    
    def evaluate():
        #Evaluation mode
        model.eval()
        for data in valid_loader:
    
            #Split batch into inputs and outputs
            x, phi, y = data[0][0], data[0][1], data[1].squeeze()
                        
            #Forward propagation
            out = model(x, phi)
            
            #Loss computation
            loss = loss_func(out, y)
            
            #Save training loss in this batch
            valid_loss.append(loss.item())
            
            #Compute the average percent error over a validation batch
            percent_error = 100*torch.div(abs(out - y), y)
            all_percent_error.extend(percent_error.flatten().squeeze(0).tolist())
        
        return valid_loss
        
    def train():
        #Training mode
        model.train()
        for data in train_loader:
            
            #Split batch into inputs and outputs
            x, phi, y = data[0][0], data[0][1], data[1].squeeze()
            
            def closure():
                #Reset gradients to zero
                optimizer.zero_grad()
                
                #Forward propagation
                out = model(x, phi)
                        
                #Loss computation
                loss = loss_func(out, y)
                
                #Backpropagation
                loss.backward()
                
                return loss
            
            #Weight optimiation
            optimizer.step(closure)
            
            #Save training loss in this batch
            train_loss.append(closure().item())
            
        return train_loss
    
    def plot_hist():
        #Plot histograms of the error (Predicted - True) in the predicted data
        error = []    
        model.eval()
        for data in valid_loader:
            #Split batch into inputs and outputs
            x, phi, y = data[0][0], data[0][1], data[1].squeeze()
            
            out = model(x, phi)
            error.append((out - y).detach().numpy())
        
        error = np.array(error)
        a_error = np.array([error[i][j][0] for i in range(len(error)) for j in range(batch_size)])
        b_error = np.array([error[i][j][1] for i in range(len(error)) for j in range(batch_size)])
        
        plt.figure(figsize=(8,6))
        plt.hist(a_error, bins=30, color='b')
        plt.title('Prediction error in parameter \'a\' in validation data')
        plt.xlabel('Predicted - True')
        plt.figure(figsize=(8,6))
        plt.hist(b_error, bins=30, color='k')
        plt.title('Prediction error in parameter \'b\' in validation data')
        plt.xlabel('Predicted - True')
        
        plt.figure(figsize=(8,6))
        p_err_less_100 = [i for i in all_percent_error if i <= 100]
        n_more_100 = len(all_percent_error) - len(p_err_less_100)  
        plt.hist(p_err_less_100, bins=30)
        plt.text(x=plt.xlim()[1]-35, y=plt.ylim()[1]-20, s='More than 100% error:\n'+str(n_more_100))
        plt.xlabel('Percent Error')
        plt.title('Histogram of percent errors in predictions of validation data')
        
        plt.show()
    
    #Print statistics about the current run
    print('\nModel Information:\n', model, sep='')
    print('\nRun Start', 
          '\n  Batch size:', batch_size, 
          '\n  Epochs:',  n_epochs,
          '\n  Number of peaks:', dataset.num_max_peaks,
          '\n  Training data size:', len(train_loader)*batch_size,
          '\n  Validation data size:', len(valid_loader)*batch_size,
          '\n  Learning rate:', lr,
          '\n  LR decay factor:', lr_factor,
          '\n  Weight decay:', weight_decay,
          '\n  Loss function:', loss_function,
          '\n  Optimizer:', repr(optimizer).partition('(')[0],
          '\n  LR scheduler:', repr(scheduler)[repr(scheduler).find('er.')+3:repr(scheduler).find(' obj')],
          '\n')
          
    #Training and evaluation loop
    for epoch in range(n_epochs): #An epoch is a run of the entire training dataset
    
        train_loss, valid_loss, all_percent_error = [], [], []
        
        #Train the network
        train_loss = train()

        #Evaluate the network
        valid_loss = evaluate()
        
        if (epoch+1) % 5 == 0:
            print('Epoch:', epoch+1,
                  '\n  Learning rate:             ', scheduler.get_lr()[0],
                  '\n  Mean epoch training loss:  ', np.mean(train_loss),
                  '\n  Mean epoch validation loss:', np.mean(valid_loss),
                  '\n  Overfitting factor:        ', np.mean(valid_loss)/np.mean(train_loss),
                  '\n  Median percent error:      ', np.median(np.array(all_percent_error)), '%')
        
        #Update the learing rate
        scheduler.step()
            
    plot_hist()

main()





