# Import data
f = h5py.File('../NavierStokes_64x64_IN.h5', 'r')
x = []
y = []
for key in f.keys():
    x.append(f[key]['input'])
    y.append(f[key]['output'])
X = np.array(x)
y = np.array(y)

# Permute axis
X = X[:, np.newaxis, ...]
y = y[:, np.newaxis, ...]

# Split data up into train, val, test
X_train, y_train = X[:768], y[:768]
X_valid, y_valid = X[768:768 + 128], y[768:768 + 128]
X_test, y_test = X[768 + 128:768 + 128*2], y[768 + 128:768 + 128*2]

# Transform data (it seems you this this normalization based off your code)
min_data, max_data = np.min(X_train), np.max(X_train)
#min_model, max_model = np.min(y_train), np.max(y_train) (this is not used - I only apply the transformation on X)

class NormalizeMinMax(torch.nn.Module):
    def __init__(self, img_min, img_max):
        self.img_min = img_min
        self.img_max = img_max
        
    def __call__(self, img):
        new_img = (img - self.img_min) / (self.img_max - self.img_min)
        return new_img
    
transform = transforms.Compose(
    [
        NormalizeMinMax(min_data, max_data),
    ]
)

# Convert to tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)

X_valid = torch.tensor(X_valid, dtype=torch.float32)
y_valid = torch.tensor(y_valid, dtype=torch.float32)

X_train = transform(X_train)
X_valid = transform(X_valid)

loss_function = nn.L1Loss() # CNO and UNet
loss_function = nn.SmoothL1Loss() # FNO

batch_size = 32 # CNO and FNO
batch_size = 10 # UNet

optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3, weight_decay=1e-6) # FNO
optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3, weight_decay=1e-10) # CNO
optimizer = torch.optim.AdamW(self.parameters(), lr=5e-4, weight_decay=1e-6) # UNet

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.98) #unet, cno, and fno

# CNO
class CNO2d(nn.Module):
    def __init__(self,
                in_dim = 1,                    # Number of input channels.
                out_dim = 1,                   # Number of input channels.
                size = 64,                      # Input and Output spatial size (required )
                N_layers = 3,                  # Number of (D) or (U) blocks in the network
                N_res = 1,                 # Number of (R) blocks per level (except the neck)
                N_res_neck = 8,            # Number of (R) blocks in the neck
                channel_multiplier = 32,   # How the number of channels evolve?
                use_bn = False,             # Add BN? We do not add BN in lifting/projection layer
                ):
# FNO
class FNO2d(nn.Module):
    def __init__(self, in_channels = 1, out_channels = 1, device=None):
        super(FNO2d, self).__init__()

        """
        The overall network. It contains 4 layers of the Fourier layer.
        1. Lift the input to the desire channel dimension by self.fc0 .
        2. 4 layers of the integral operators u' = (W + K)(u).
            W defined by self.w; K defined by self.conv .
        3. Project from the channel space to the output space by self.fc1 and self.fc2 .

        input: the solution of the coefficient function and locations (a(x, y), x, y)
        input shape: (batchsize, x=s, y=s, c=3)
        output: the solution
        output shape: (batchsize, x=s, y=s, c=1)
        """
        self.modes1 = 16 #16
        self.modes2 = 16 #16
        self.width = 128 #64
        self.n_layers = 5 #5
        self.retrain_fno = 4 #4
        self.padding = 0 #0
        self.include_grid = 1 #1
        self.input_dim = in_channels
        self.act  = nn.LeakyReLU() 
        self.device = device
 # UNet
 class UNet(nn.Module):
    def __init__(self, n_channels = 1, n_classes = 1, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = (DoubleConv(n_channels, 64))
        self.down1 = (Down(64, 128))
        self.down2 = (Down(128, 256))
        self.down3 = (Down(256, 512))
        factor = 2 if bilinear else 1
        self.down4 = (Down(512, 1024 // factor))
        self.up1 = (Up(1024, 512 // factor, bilinear))
        self.up2 = (Up(512, 256 // factor, bilinear))
        self.up3 = (Up(256, 128 // factor, bilinear))
        self.up4 = (Up(128, 64, bilinear))
        self.outc = (OutConv(64, n_classes))

X_test = torch.tensor(X_test, dtype=torch.float32)
X_test = transform(X_test)
test_loader = DataLoader(list(X_test), shuffle=False, batch_size=1)
y_predicted = trainer.predict(network, test_loader)

def relative_l1_error(y_pred, y_true):
    return torch.mean(torch.abs(y_pred - y_true)) / np.mean(np.abs(y_true))