# -*- coding: utf-8 -*-
"""
Created on Mon May 30 17:51:23 2022

@author: BraemD

Example: NODE fit of a 1cp PK model with multiple subjects
"""

#Packages import
import time
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim

import matplotlib.pyplot as plt

from torchdiffeq import odeint_adjoint as odeint

from scipy.integrate import solve_ivp

#Set seed for reproducibility
np.random.seed(8)
torch.manual_seed(0)

#Set device (cpu or gpu if available)
device = torch.device('cpu')

#Define function to plot  
def visualize_fit_multi(t,true,pred,legend=False):
    for i in range(true.shape[1]):
        plt.scatter(t,true[:,i,0],label="PK data "+str(i))
    plt.plot(t,pred[:,0,0], label="nODE fit")
    if legend:
        plt.legend()
    plt.xlabel("Time")
    plt.ylabel("Concentration")
    plt.show()

# Define ODE to simulate data
def ode_fun(t,z,parms):
    kel = parms
    C = z
    dC = -kel * C
    return [dC]

seq_length = 5 #length of sequence
rs = 0.2 #residual error

t = np.array([0,8,16,24,40]) #Time points to simulate concentrations

parms = np.array([0.1, #kel
         ])

V = 2 #Volume of distribution
dose = 1 #Dose

pop = 20 #Number of subjects

z0 = [dose/V] #Initial condition for simulation

sol = solve_ivp(ode_fun,[0,np.max(t)],z0,t_eval=t,args=([parms]), method = 'LSODA') #Simulation

z = sol.y
og_z = z #Single data without noise

#Create dataset for pop subjects with residual error (pop, seq_length)
z = z + np.random.normal(scale=rs,size=seq_length) * z
for i in range(pop-1):
    z_ = og_z + np.random.normal(scale=rs,size=seq_length) * og_z
    z = np.concatenate((z,z_),axis=0)

#Illustrate simulated data
for i in range(pop):
    plt.plot(t,z[i,:],label=i)
plt.show()

#Convert numpy-dataset to tensor with correct dimensions (seq_length, pop, 1)
true_y = torch.Tensor(z[0,:]).view(seq_length,1,1)
for i in range(pop-1):
    true_y = torch.cat((true_y,torch.Tensor(z[(i+1),:]).view(seq_length,1,1)),dim=1)
og_y = torch.Tensor(og_z[0]).view(seq_length,1,1)


#Define Neural Network to be used in the NODE
class ann_dy(nn.Module):

    def __init__(self):
        super(ann_dy, self).__init__()
        
        self.net = nn.Sequential(
            nn.Linear(1,20,bias=True),
            nn.ReLU(),
            nn.Linear(20,1,bias=True))
        
        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0, std=0.1)
                nn.init.normal_(m.bias, mean=0, std=0.5)

    def forward(self, t, y):
        
        return self.net(y)
    
#Define Neural Network for the Volume of distribution
class V_parm(nn.Module):
    def __init__(self):
        super(V_parm,self).__init__()
        self.V = nn.Parameter(torch.Tensor([2.5]),requires_grad=True)
        
    def forward(self,x):
        out = x/self.V
        return out
   
#Create Neural Netorks
V_fit = V_parm()
func = ann_dy().to(device)

#Time vector to tensor
t_fit = torch.Tensor(t)

#Parameters to optimize
opt_params = (list(func.parameters()) + list(V_fit.parameters()))

#initiate optimizer
optimizer = optim.Adam(opt_params, lr=1e-3)

#Define number of iterations and reset of learning rate
iters = 2000
opt_reset = 200

#loss
loss_over_time = torch.zeros(iters)

#Training
for itr in range(1, iters + 1):
    #Initial prediction
    if itr == 1 :
        start = time.time()
        with torch.no_grad():
            true_y0 = torch.ones_like(true_y[0]) * V_fit(dose)
            pred_y = odeint(func, true_y0, t_fit)
            loss = torch.mean((pred_y - true_y)**2)
            print('Iter {:04d} | Total Loss {:.6f}'.format(itr, loss.item()))
            visualize_fit_multi(t, true_y, pred_y)
            
    #Reset gradient
    optimizer.zero_grad()
    
    #Initial condition for NODE fit
    true_y0 = torch.ones_like(true_y[0]) * V_fit(dose)
    
    #Predict concentrations
    pred_y = odeint(func, true_y0, t_fit,rtol=1e-3,atol=1e-4).to(device)
    
    #Calculate loss (difference between predicted and observed concentrations)
    loss = torch.mean((pred_y - true_y)**2)
    
    #Backpropagation to calculate gradient of weights and biases
    loss.backward(retain_graph=True)
    
    #Optimize weights and biases along gradient
    optimizer.step()
    
    loss_over_time[itr-1] = loss.item()
        
    #Show progress every 100th iteration
    if itr % 100 == 0:
        with torch.no_grad():
            true_y0 = torch.ones_like(true_y[0]) * V_fit(dose)
            pred_y = odeint(func, true_y0, t_fit)
            loss = torch.mean((pred_y - true_y)**2)
            print('Iter {:04d} | Total Loss {:.6f}'.format(itr, loss.item()))
            visualize_fit_multi(t, true_y, pred_y)
            
    #Reset learning rate to accelerate optimization
    if itr % opt_reset ==0:
        optimizer = optim.Adam(opt_params, lr=1e-3)

    end = time.time()
   

dose_2 = 1 #Same dose as in fit
t_dense = np.linspace(0,40,40) #Dense timepoints for prediction

z0_2 = [dose_2/V] #Initial condition for simulation

sol_2 = solve_ivp(ode_fun,[0,np.max(t_dense)],z0_2,t_eval=t_dense,args=([parms]), method = 'LSODA') #Dense simulation

z_2 = sol_2.y

#Conversion to tensors
true_y_2 = torch.Tensor(z_2[0]).view(40,1,1)
true_y0_2 = true_y_2[0]
t_fit_dense = torch.Tensor(t_dense)


with torch.no_grad():
    #Dense prediction
    pred_y_2 = odeint(func, true_y0_2, t_fit_dense)
    #Derivatives given by Neural Network within NODE
    dx = func.net(true_y_2[:,:,0])
    
#Show final fit with dense prediction
with torch.no_grad():
    plt.figure(figsize=(7,4.5))
    plt.scatter(t,true_y[:,0,0],label="PK data")
    for i in range(1,pop):
        plt.scatter(t,true_y[:,i,0])
    plt.plot(t_dense,pred_y_2[:,0,0],label="NODE fit")
    plt.legend(fontsize=18)
    plt.ylabel("Concentration")
    plt.xlabel("Time")
    plt.rc('font', size = 18)
    plt.tight_layout()
    plt.show()

#Show Derivative vs Concentration of NODE
with torch.no_grad():
    dx = func.net(true_y_2)
    plt.figure(figsize=(7,4.5))
    plt.plot(z_2[0,:],dx[:,0,0],label="NODE fit")
    plt.plot(z_2[0,:],-0.1 * z_2[0,:],label="PK model",color="red",linestyle=(0, (5, 5)))
    plt.legend(fontsize=18)
    plt.xlabel("Concentration")
    plt.ylabel("Derivative")
    plt.rc('font', size = 18)
    plt.tight_layout()
    plt.show()