import torch
from pina.problem import SpatialProblem, TimeDependentProblem
from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux
from pina.geometry import CartesianDomain
from pina.model import FeedForward
from pina.operators import grad, laplacian
from pina import Condition, Trainer
from pina.solvers import PINN
from pytorch_lightning.callbacks import Callback, Timer 
from pytorch_lightning import seed_everything
from pina.plotter import Plotter 

$$ \begin{cases} \frac{\partial u}{\partial t} - \alpha \frac{\partial ^2 u}{\partial x^2} = f \\ u(a,t) = u(b,t) = 0, \partial \Omega\\ u(.,t) = u_0(x) = \exp(-x^2), \Omega\\ \end{cases} $$

class Heat1DEquation(SpatialProblem, TimeDependentProblem):
    output_variables = ['u']
    spatial_domain = CartesianDomain({'x': [0, 1]})
    temporal_domain = CartesianDomain({'t': [0, 1]})
    
    def pde (input_, output_):
        alpha =0.08
        u_t = grad(output_, input_, components=['u'], d=['t'])
        u_x = grad(output_, input_, components=['u'], d=['x'])
        u_xx = grad(u_x, input_, components=['dudx'], d=['x'])
        return u_t - alpha*u_xx
    
    def initial_condition(input_, output_):
        u_0 = torch.exp(-input_.extract(['x'])**2)
        return output_.extract(['u']) - u_0
    
    conditions = {
        'Dirichlet1': Condition(location=CartesianDomain({'x': 0., 't':[0,1]}), equation=FixedValue(0.)),
        'Dirichlet2': Condition(location=CartesianDomain({'x': 1., 't':[0,1]}), equation=FixedValue(0.)),
        't0': Condition(location=CartesianDomain({'x': [0, 1], 't':0.}), equation=Equation(initial_condition)),
        'Domain': Condition(location=CartesianDomain({'x': [0, 1], 't':[0,1]}), equation=Equation(pde))
    }

heat = Heat1DEquation()
heat.discretise_domain(10000, "random", locations=['Dirichlet1', 'Dirichlet2', "t0", "Domain"])

model = FeedForward(
    layers=[50 for _ in range(20)],
    func=torch.nn.Tanh,
    output_dimensions=len(heat.output_variables),
    input_dimensions=len(heat.input_variables)
)

pinn = PINN(heat, model)

trainer = Trainer(
    solver=pinn,
    accelerator='cpu',
    enable_model_summary=False,
    max_epochs = 1000
)

trainer.train()

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs

Training: |          | 0/? [00:00<?, ?it/s]

`Trainer.fit` stopped: `max_epochs=1000` reached.

plotter = Plotter()
plotter.plot(solver=pinn)