'''
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
plt.close('all')
plt.interactive(True)

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


'''
 T = temperature, k = conductivity coefficient

 the 2D heat conduction equation reads  k*(T_xx + T_yy) + 2 = 0, k = 0.01
 b.c. t(x=0)=1; t(x=1) = 0; t(y=0)=1; u(y=1) = 0

 load the finite volume solution

'''
folder = 'poisson/'
datax= np.loadtxt(str(folder) + "x2d.dat")
x=datax[0:-1]
ni=int(datax[-1])
datay= np.loadtxt(str(folder) + "y2d.dat")
y=datay[0:-1]
nj=int(datay[-1])
T2d_pycalc = np.load(str(folder) + 'u2d_saved.npy')

x2d=np.zeros((ni+1,nj+1))
y2d=np.zeros((ni+1,nj+1))

x2d=np.reshape(x,(ni+1,nj+1))
y2d=np.reshape(y,(ni+1,nj+1))

# compute cell centers
xp2d=0.25*(x2d[0:-1,0:-1]+x2d[0:-1,1:]+x2d[1:,0:-1]+x2d[1:,1:])
yp2d=0.25*(y2d[0:-1,0:-1]+y2d[0:-1,1:]+y2d[1:,0:-1]+y2d[1:,1:])


xp = xp2d[:,0]
yp = yp2d[0,:]

x= xp2d.flatten()
y= yp2d.flatten()

# set heat conductivity coefficient
k = 1/100
# the source tern
b = torch.ones(len(x))*2

# set b.c.
T_west_bc = 0
T_east_bc = 0
T_south_bc = 0
T_north_bc = 1

# the x, y coordinate for the inner nodes
x_int = torch.tensor(x, requires_grad=True, dtype=torch.float32).view(-1,1)
y_int = torch.tensor(y, requires_grad=True, dtype=torch.float32).view(-1,1)

# merge x and y into a tensor with two collumns
interior=(torch.cat((x_int,y_int),1)).requires_grad_(True)

# transform x and y (at the cell centers) to tensors
xp2d = torch.tensor(xp2d,dtype=torch.float32)
yp2d = torch.tensor(yp2d,dtype=torch.float32)

# the x, y coordinate at the west boundary (x=0)
y_west=yp2d[0,:].requires_grad_(True)
x_west=torch.full_like(y_west,0)
y_west=y_west.view(-1,1)
x_west=x_west.view(-1,1)
west=torch.cat((x_west,y_west),1)

# the x, y coordinate at the east boundary (x=1)
y_east=yp2d[-1,:].requires_grad_(True)
x_east=torch.full_like(y_east,1)
y_east=y_east.view(-1,1)
x_east=x_east.view(-1,1)
east=torch.cat((x_east,y_east),1)

# the x, y coordinate at the north boundary (y=1)
x_north=xp2d[:,-1].requires_grad_(True)
y_north=(torch.full_like(x_north,1)).requires_grad_(True)
y_north=y_north.view(-1,1)
x_north=x_north.view(-1,1)
north=torch.cat((x_north,y_north),1)

# the x, y coordinate at the south boundary (y=0)
x_south=xp2d[:,0].requires_grad_(True)
y_south=(torch.full_like(x_south,0)).requires_grad_(True)
y_south=y_south.view(-1,1)
x_south=x_south.view(-1,1)
south=torch.cat((x_south,y_south),1).requires_grad_(True)

# iniiiate the model 
# first argument: 2 input paramaters (x and y)
# second argument: 1 input paramaters (t)
# third argument:  40 neurons 
model=MyNet(2,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 t at the interior nodes
    T_interior=model(interior)

# compute 2nd derivatives
    dT_dx = torch.autograd.grad(T_interior, x_int, torch.ones_like(T_interior), create_graph=True)[0]
    d2T_dx2 = torch.autograd.grad(dT_dx, x_int, torch.ones_like(dT_dx), create_graph=True)[0]
    dT_dy = torch.autograd.grad(T_interior, y_int, torch.ones_like(T_interior), create_graph=True)[0]
    d2T_dy2 = torch.autograd.grad(dT_dy, y_int, torch.ones_like(dT_dy), create_graph=True)[0]

# compute the loss == PDE
    pde_loss= torch.mean((k*d2T_dx2 + k*d2T_dy2+b)**2)

# compute t at the boundary nodes
    T_west=model(west)
    T_east=model(east)
    T_north=model(north)
    T_south=model(south)

# set Dirichlet b,c, and compute the loss == errors
    west_loss=1*torch.mean((T_west-T_west_bc)**2)
    east_loss=1*torch.mean((T_east-T_east_bc)**2)
    north_loss=1*torch.mean((T_north-T_north_bc)**2)
    south_loss=1*torch.mean((T_south-T_south_bc)**2)

# set Neumann b.c. 
#   dT_north = torch.autograd.grad(T_north, y_north, torch.ones_like(T_north), create_graph=True)[0]
#   north_loss=1*torch.mean((dT_north-neuman_north)**2)
#   dT_south = torch.autograd.grad(T_south, y_south, torch.ones_like(T_south), create_graph=True)[0]
#   south_loss=1*torch.mean((dT_south-0)**2)

# total loss
    loss=pde_loss+west_loss+east_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)
    T_max = torch.max(T_interior)
    print(f"{'Epoch: '}{epoch},{'Loss:'}{loss_np:.2e}, {'Loss_min:'} {loss_min:.2e}, {'T_max:'} {T_max:.2e}")

# plot loss history
plt.plot(range(max_epochs),loss_history)


######################## plot u
fig, ax = plt.subplots(nrows=1, ncols=1) # Create a figure with one subplot
plt.subplots_adjust(left=0.20,bottom=0.20)
T_np =T_interior.detach().numpy()[:,0]
T_2d=np.reshape(T_np,(ni,nj))
ax.plot(xp, T_2d[:,5],'b-',linewidth=2, label='PINN')
ax.plot(xp, T2d_pycalc[:,5],'ro',linewidth=2, label='pyCALC')
ax.legend(loc='best')
plt.xlabel('$y$')
plt.ylabel(r'$u$')
ax.grid(visible=True)
plt.savefig('t-poisson-2D.png',bbox_inches='tight')