import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

import matplotlib.pyplot as plt
from utilities4 import *

import operator
from functools import reduce
from functools import partial

from timeit import default_timer
import scipy.io

torch.manual_seed(0)
np.random.seed(0)


def compl_mul2d(a, b):
    # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t)
    return torch.einsum("bixy,ioxy->boxy", a, b)

    # return torch.stack([
    #     op(a[..., 0], b[..., 0]) - op(a[..., 1], b[..., 1]),
    #     op(a[..., 1], b[..., 0]) + op(a[..., 0], b[..., 1])
    # ], dim=-1)

class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, modes1, modes2):
        super(SpectralConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2

        self.scale = (1 / (in_channels * out_channels))
        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))

    def forward(self, x, size=None):
        if size==None:
            size = x.size(-1)

        batchsize = x.shape[0]
        #Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfftn(x, dim=[2,3], norm="ortho")

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(batchsize, self.out_channels, size, size//2 + 1, device=x.device, dtype=torch.cfloat)
        out_ft[:, :, :self.modes1, :self.modes2] = \
            compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)


        #Return to physical space
        x = torch.fft.irfftn(out_ft, s=(size, size), dim=[2,3], norm="ortho")
        return x

class SimpleBlock2d(nn.Module):
    def __init__(self, in_dim, out_dim, modes1, modes2, width):
        super(SimpleBlock2d, self).__init__()

        self.modes1 = modes1
        self.modes2 = modes2

        self.width_list = [width*2//4, width*3//4, width*4//4, width*4//4, width*5//4]
        self.size_list = [64,] * 5
        self.grid_dim = 2

        self.fc0 = nn.Linear(in_dim+self.grid_dim, self.width_list[0])

        self.conv0 = SpectralConv2d(self.width_list[0], self.width_list[1], self.modes1*4//4, self.modes2*4//4)
        self.conv1 = SpectralConv2d(self.width_list[1], self.width_list[2], self.modes1*3//4, self.modes2*3//4)
        self.conv2 = SpectralConv2d(self.width_list[2], self.width_list[3], self.modes1*2//4, self.modes2*2//4)
        self.conv3 = SpectralConv2d(self.width_list[3], self.width_list[4], self.modes1*2//4, self.modes2*2//4)
        self.w0 = nn.Conv1d(self.width_list[0], self.width_list[1], 1)
        self.w1 = nn.Conv1d(self.width_list[1], self.width_list[2], 1)
        self.w2 = nn.Conv1d(self.width_list[2], self.width_list[3], 1)
        self.w3 = nn.Conv1d(self.width_list[3], self.width_list[4], 1)

        self.fc1 = nn.Linear(self.width_list[4], self.width_list[4]*2)
        self.fc2 = nn.Linear(self.width_list[4]*2, self.width_list[4]*2)
        self.fc3 = nn.Linear(self.width_list[4]*2, out_dim)

    def forward(self, x):

        batchsize = x.shape[0]
        size_x, size_y= x.shape[1], x.shape[2]
        grid = self.get_grid(size_x, batchsize, x.device)
        size_list = self.size_list

        x = torch.cat((x, grid.permute(0, 2, 3, 1)), dim=-1)

        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)

        x1 = self.conv0(x, size_list[1])
        x2 = self.w0(x.view(batchsize, self.width_list[0], size_list[0]**2)).view(batchsize, self.width_list[1], size_list[0], size_list[0])
        # x2 = F.interpolate(x2, size=size_list[1], mode='trilinear')
        x = x1 + x2
        x = F.selu(x)

        x1 = self.conv1(x, size_list[2])
        x2 = self.w1(x.view(batchsize, self.width_list[1], size_list[1]**2)).view(batchsize, self.width_list[2], size_list[1], size_list[1])
        # x2 = F.interpolate(x2, size=size_list[2], mode='trilinear')
        x = x1 + x2
        x = F.selu(x)

        x1 = self.conv2(x, size_list[3])
        x2 = self.w2(x.view(batchsize, self.width_list[2], size_list[2]**2)).view(batchsize, self.width_list[3], size_list[2], size_list[2])
        # x2 = F.interpolate(x2, size=size_list[3], mode='trilinear')
        x = x1 + x2
        x = F.selu(x)

        x1 = self.conv3(x, size_list[4])
        x2 = self.w3(x.view(batchsize, self.width_list[3], size_list[3]**2)).view(batchsize, self.width_list[4], size_list[3], size_list[3])
        # x2 = F.interpolate(x2, size=size_list[4], mode='trilinear')
        x = x1 + x2

        x = x.permute(0, 2, 3, 1)
        x = self.fc1(x)
        x = F.selu(x)
        x = self.fc2(x)
        x = F.selu(x)
        x = self.fc3(x)
        return x

    def get_grid(self, S, batchsize, device):
        gridx = torch.tensor(np.linspace(0, 1, S), dtype=torch.float)
        gridx = gridx.reshape(1, 1, S, 1).repeat([batchsize, 1, 1, S])
        gridy = torch.tensor(np.linspace(0, 1, S), dtype=torch.float)
        gridy = gridy.reshape(1, 1, 1, S).repeat([batchsize, 1, S, 1])
        return torch.cat((gridx, gridy), dim=1).to(device)

class Net2d(nn.Module):
    def __init__(self, in_dim, out_dim, modes, width):
        super(Net2d, self).__init__()
        self.conv1 = SimpleBlock2d(in_dim, out_dim, modes, modes, width)

    def forward(self, x):
        x = self.conv1(x)
        return x

    def count_params(self):
        c = 0
        for p in self.parameters():
            c += reduce(operator.mul, list(p.size()))

        return c


ntrain = 100
ntest = 20

modes = 30
width = 180

batch_size = 20

epochs = 50
learning_rate = 0.0005
scheduler_step = 10
scheduler_gamma = 0.5

loss_k = 2
loss_group = True

print(epochs, learning_rate, scheduler_step, scheduler_gamma)

path = 'KF_fourier_big_N'+str(ntrain)+'_k' + str(loss_k)+'_g' + str(loss_group)+'_ep' + str(epochs) + '_m' + str(modes) + '_w' + str(width)
path_model = 'model/'+path
path_train_err = 'results/'+path+'train.txt'
path_test_err = 'results/'+path+'test.txt'
path_image = 'image/'+path



in_dim = 1
out_dim = 1

sub = 1
S = 64 // sub

T_in = 100
T = 400
T_out = T_in+T
step = 1


t1 = default_timer()
# data = np.load('data/KFvorticity_Re500_N175_T500_dt0.25.npy')
data = np.load('data/KFvorticity_Re500_N1000_T500.npy')
# data = np.load('data/KFvorticity_Re40_N200_T500.npy')
# data = np.load('data/KFvelocity_Re40_N25_part1.npy')
# data = np.load('data/KFvorticity_Re40_N25_part1.npy')
# data = np.load('data/KFstream_Re40_N25_part1.npy')
data = torch.tensor(data, dtype=torch.float)[:, :, ::sub, ::sub]
print(data.shape )

train_a = data[:ntrain,T_in-1:T_out-1].reshape(ntrain*T, S, S, -1)
train_u = data[:ntrain,T_in:T_out].reshape(ntrain*T, S, S, -1)

test_a = data[-ntest:,T_in-1:T_out-1].reshape(ntest*T, S, S, -1)
test_u = data[-ntest:,T_in:T_out].reshape(ntest*T, S, S, -1)

# train_a = w_to_f(train_a)
# train_u = w_to_f(train_u)
# test_a = w_to_f(test_a)
# test_u = w_to_f(test_u)

print(train_a.shape)
print(train_u.shape)
assert (S == train_u.shape[2])



train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u), batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=False)

t2 = default_timer()

print('preprocessing finished, time used:', t2-t1)
device = torch.device('cuda')

model = Net2d(in_dim, out_dim, modes, width).cuda()
# model = torch.load('model/KF_w_fourier500_N900_k2_gTrue_ep50_m20_w64')

print(model.count_params())
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma)


lploss = LpLoss(size_average=False)
h1loss = HsLoss(k=1, group=False, size_average=False)
h2loss = HsLoss(k=2, group=False, size_average=False)
myloss = HsLoss(k=loss_k, group=loss_group, size_average=False)

for ep in range(epochs):
    model.train()
    t1 = default_timer()
    train_l2 = 0
    for x, y in train_loader:
        x = x.to(device).view(batch_size, S, S, in_dim)
        y = y.to(device).view(batch_size, S, S, out_dim)

        out = model(x).reshape(batch_size, S, S, out_dim)
        loss = myloss(out, y)
        train_l2 += loss.item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    test_l2 = 0
    test_u = 0
    test_h1 = 0
    test_h2 = 0
    with torch.no_grad():
        for x, y in test_loader:
            x = x.to(device).view(batch_size, S, S, in_dim)
            y = y.to(device).view(batch_size, S, S, out_dim)

            out = model(x).reshape(batch_size, S, S, out_dim)
            test_l2 += lploss(out, y).item()
            test_u += lploss(w_to_u(out), w_to_u(y)).item()
            test_h1 += h1loss(out, y).item()
            test_h2 += h2loss(out, y).item()


    t2 = default_timer()
    scheduler.step()
    print(ep, t2 - t1, train_l2/(ntrain*T), test_l2/(ntest*T), test_u/(ntest*T), test_h1/(ntest*T), test_h2/(ntest*T) )

torch.save(model, path_model)

# model = torch.load('model/KF_fourier_N180_k0_gTrue_ep200_m20_w64')
# model = torch.load('model/KF_w_fourier500_N900_k0_gTrue_ep50_m20_w64')
#
# model.eval()
# # test_a = np.load('data/KFvorticity_N1_Re40_T1000.npy')[100]
# test_a = np.load('data/KFvorticity_Re500_N25_part1.npy')[-1,100,::4,::4]
# print(test_a.shape)
# test_a = torch.tensor(test_a, dtype=torch.float).reshape(-1,64,64,1).cuda()
#
# input = test_a[0,:,:].reshape(1,S,S,1)
# # test_a = w_to_u(test_a)
# T = 10000
# pred = torch.zeros(S,S,T,in_dim)
# out = input.reshape(1,S,S,in_dim)
# with torch.no_grad():
#     for i in range(T):
#         out = model(out.reshape(1,S,S,in_dim))
#         pred[:,:,i] = out.view(S,S,in_dim)
#
#         # print(i, lploss(out, test_a[i+1]))
#         if i% 100 == 0:
#             print(i, torch.mean(out))
#
# print("complete")
# print(torch.mean(pred))
# # scipy.io.savemat('pred/'+path+'.mat', mdict={'pred': pred.cpu().numpy()})
#
# scipy.io.savemat('pred/KF_w_fourier500_T10000_k0_gTrue_ep50_m20_w64.mat', mdict={'pred': pred.cpu().numpy()})
# # scipy.io.savemat('pred/KF_w_fourier40_T10000_k0_gTrue_ep200_m20_w64.mat', mdict={'pred': pred.cpu().numpy()})
#
# # scipy.io.savemat('pred/KF_w_fourier40_T10000_k2_gTrue_ep200_m20_w64.mat', mdict={'pred': pred.cpu().numpy(), 'truth': test_a[1:].cpu().numpy()})


