'''
This PINN solver is based on the 2D PINN-Solver at

https://github.com/AshmitB05/2D-PINN-Solver

'''



import numpy as np
import sys
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from matplotlib import rcParams # for changing default values


plt.close('all')
plt.interactive(True)
# set all fontsizes to 16
rcParams["font.size"] = 16

'''
 T = temperature, k = conductivity coefficient

 the 2D heat conduction equation reads  k*T_yy + 2 = 0
 b.c. T(0)=1; T(1) = 0
 integrate the eq.  T_y = (-2y +C1)/k
 integrate once more T = -y**2/k +C1*y/k + C2/k
 b.c. at y=0 gives C2=0
 b.c. at y=1 gives C1=1

 T =  (-y**2 +y)/k

'''

# solve differential equation for heat conductivity coefficient, k

# the grid
nj = 10
# cell faces
yc  = np.linspace(0,1,nj+1)
# cell centers (nodes)
yp = 0.5*(yc[0:-1] + yc[1:])
k = 1/100
T = (-yp**2 +yp)/k # k = const


dTdy_grad = np.gradient(T,yp,edge_order = 2)
d2Tdy2_grad = np.gradient(dTdy_grad,yp,edge_order = 2)
dTdy= (-2*yp +1)/k
d2Tdy2= -2/k

# source 
b = 2*np.ones(nj)

T = torch.tensor(T, requires_grad=False, dtype=torch.float32).view((-1, 1))
d2Tdy2 = torch.tensor(d2Tdy2, requires_grad=False, dtype=torch.float32).view((-1, 1))
dtdy = torch.tensor(dTdy, requires_grad=False, dtype=torch.float32).view((-1, 1))
y = torch.tensor(yp, requires_grad=True, dtype=torch.float32).view((-1, 1))
b = torch.tensor(b, requires_grad=False, dtype=torch.float32).view((-1, 1))

# interior nodes
y_int = y.clone().detach().requires_grad_(True).view(-1,1)

# the y coordinate at the north boundary (y=1)
y_north=1.
north=torch.tensor(y_north).view(-1,1)

# the y coordinate at the south boundary (y=0)
y_south=0.
south=torch.tensor(y_south).view(-1,1)

k_south_bc = k
k_north_bc = k



class MyNet(nn.Module):
    def __init__(self,input,output,h1):
        super().__init__()
        self.fc1=nn.Linear(input,h1)
        self.fc2=nn.Linear(h1,h1)
        self.fc4=nn.Linear(h1,h1)
        self.fc5=nn.Linear(h1,h1)
        self.fc6=nn.Linear(h1,h1)
        self.fc3=nn.Linear(h1,output)
    def forward(self,x):
        x=torch.tanh(self.fc1(x))
        x=torch.tanh(self.fc2(x))
        x=torch.tanh(self.fc4(x))
        x=torch.tanh(self.fc5(x))
        x=torch.tanh(self.fc6(x))
        x=self.fc3(x)
        return x

    
# iniiiate the model 
# first argument: 1 input paramaters (y)
# second argument: 1 input paramaters (k)
# third argument:  40 neurons 
model = MyNet(1,1,40)

# choose Adam optimized witjh learning rate  = 2e-4
optimizer=torch.optim.Adam(model.parameters(),lr=2*(1e-4))

max_epochs=5000
#max_epochs=3
loss_history=[]
loss_min=1e10

for epoch in range(max_epochs):
    optimizer.zero_grad()
# compute k at the interior nodes
    k_pred=model(y_int)

# compute derivative
    dk_dy = torch.autograd.grad(k_pred,y_int, torch.ones_like(k_pred), create_graph=True)[0]

# Differential equation loss
    imbalance = k_pred * d2Tdy2 + b
    pde_loss = torch.mean(imbalance**2)

# compute k at the boundary nodes
    k_north=model(north)
    k_south=model(south)

# set Dirichlet b,c, and compute the loss == errors
    south_loss=1*torch.mean((k_south-k_south_bc)**2)
    north_loss=1*torch.mean((k_north-k_north_bc)**2)

# total loss
    loss=pde_loss+north_loss+south_loss
    loss.backward()
    optimizer.step()
    loss_history.append(loss.detach().numpy())

# print loss and t-max
    loss_np = loss.detach().numpy()
    loss_np = loss_np.item()
    loss_min = np.minimum(loss_np,loss_min)
    k_min = torch.min(k_pred)
    k_max = torch.max(k_pred)
    print(f"{'Epoch: '}{epoch},{'Loss:'}{loss_np:.2e}, {'Loss_min:'} {loss_min:.2e}, {'k_min:'} {k_min:.2e}, {'k_max:'} {k_max:.2e}")


fig, ax = plt.subplots(nrows=1, ncols=1) # Create a figure with one subplot
ax.semilogy(np.arange(len(loss_history)), loss_history,color='blue',label="diff eq error")
plt.xlabel(r'epochs')
ax.set_title(r'Errors')
ax.grid(visible=True)
ax.legend(loc='best') 

k_pred=model(y_int)
k_pred_np =  k_pred.detach().numpy()[:,0]

######################## plot k
fig, ax = plt.subplots(nrows=1, ncols=1) # Create a figure with one subplot
plt.subplots_adjust(left=0.30,bottom=0.20)
k_exact = np.ones(len(y))*k
plt.plot(yp, k_exact,'r-',linewidth=2, label='exact')
plt.plot(yp, k_pred_np,color='k',linestyle='-',linewidth=2, label=r"$\nu_t{\mathrm{pred}}$")
plt.legend(loc='best') 
plt.xlabel('$y$')
plt.ylabel('heat conductivity, $k$')   
plt.savefig('heat-coeff-k.png',bbox_inches='tight')