import torch
import pytest

from pina.problem import AbstractProblem
from pina import Condition, LabelTensor
from pina.solvers import ReducedOrderModelSolver
from pina.trainer import Trainer
from pina.model import FeedForward
from pina.loss import LpLoss


# class NeuralOperatorProblem(AbstractProblem):
#     input_variables = ['u_0', 'u_1']
#     output_variables = [f'u_{i}' for i in range(100)]
#     conditions = {'data' : Condition(input_points=
#                                      LabelTensor(torch.rand(10, 2),
#                                                  input_variables), 
#                                      output_points=
#                                      LabelTensor(torch.rand(10, 100),
#                                                  output_variables))}


# # make the problem + extra feats
# class AE(torch.nn.Module):
#     def __init__(self, input_dimensions, rank):
#         super().__init__()
#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])
# class AE_missing_encode(torch.nn.Module):
#     def __init__(self, input_dimensions, rank):
#         super().__init__()
#         self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4])
# class AE_missing_decode(torch.nn.Module):
#     def __init__(self, input_dimensions, rank):
#         super().__init__()
#         self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4])

# rank = 10
# problem = NeuralOperatorProblem()
# interpolation_net = FeedForward(len(problem.input_variables),
#                                 rank)
# reduction_net = AE(len(problem.output_variables), rank)

# def test_constructor():
#     ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net,
#     interpolation_network=interpolation_net)
#     with pytest.raises(SyntaxError):
#         ReducedOrderModelSolver(problem=problem,
#                      reduction_network=AE_missing_encode(
#                          len(problem.output_variables), rank),
#                      interpolation_network=interpolation_net)
#         ReducedOrderModelSolver(problem=problem,
#                      reduction_network=AE_missing_decode(
#                          len(problem.output_variables), rank),
#                      interpolation_network=interpolation_net)


# def test_train_cpu():
#     solver = ReducedOrderModelSolver(problem = problem,reduction_network=reduction_net,
#     interpolation_network=interpolation_net, loss=LpLoss())
#     trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20)
#     trainer.train()


# def test_train_restore():
#     tmpdir = "tests/tmp_restore"
#     solver = ReducedOrderModelSolver(problem=problem,
#                reduction_network=reduction_net,
#                interpolation_network=interpolation_net,
#                loss=LpLoss())
#     trainer = Trainer(solver=solver,
#                       max_epochs=5,
#                       accelerator='cpu',
#                       default_root_dir=tmpdir)
#     trainer.train()
#     ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu')
#     t = ntrainer.train(
#         ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
#     import shutil
#     shutil.rmtree(tmpdir)


# def test_train_load():
#     tmpdir = "tests/tmp_load"
#     solver = ReducedOrderModelSolver(problem=problem,
#                reduction_network=reduction_net,
#                interpolation_network=interpolation_net,
#                loss=LpLoss())
#     trainer = Trainer(solver=solver,
#                       max_epochs=15,
#                       accelerator='cpu',
#                       default_root_dir=tmpdir)
#     trainer.train()
#     new_solver = ReducedOrderModelSolver.load_from_checkpoint(
#         f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
#         problem = problem,reduction_network=reduction_net,
#         interpolation_network=interpolation_net)
#     test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables)
#     assert new_solver.forward(test_pts).shape == (20, 100)
#     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
#     torch.testing.assert_close(
#         new_solver.forward(test_pts),
#         solver.forward(test_pts))
#     import shutil
#     shutil.rmtree(tmpdir)