import torch
import numpy as np

import torch.nn.functional as F

from torch_geometric.data import Data, DataLoader
from utilities import *
from nn_conv import NNConv_old

from timeit import default_timer


class KernelNN3(torch.nn.Module):
    def __init__(self, width_node, width_kernel, depth, ker_in, in_width=1, out_width=1):
        super(KernelNN3, self).__init__()
        self.depth = depth

        self.fc1 = torch.nn.Linear(in_width, width_node)

        kernel = DenseNet([ker_in, width_kernel // 2, width_kernel, width_node**2], torch.nn.ReLU)
        self.conv1 = NNConv_old(width_node, width_node, kernel, aggr='mean')

        self.fc2 = torch.nn.Linear(width_node, 1)

    def forward(self, data):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
        x = self.fc1(x)
        for k in range(self.depth):
            x = F.relu(self.conv1(x, edge_index, edge_attr))

        x = self.fc2(x)
        return x


TRAIN_PATH = 'data/piececonst_r241_N1024_smooth1.mat'
TEST_PATH = 'data/piececonst_r241_N1024_smooth2.mat'

for m in (100,200,400,800):
    mtest1 = 100
    mtest2 = 200
    mtest3 = 400
    mtest4 = 800

    r = 2
    s = int(((241 - 1)/r) + 1)
    n = s**2
    # m = 200
    k = 5

    radius_train = 0.15
    radius_test = 0.15
    print('resolution', s)


    ntrain = 100
    ntest = 100


    batch_size = 10
    batch_size2 = 10
    width = 64
    ker_width = 1000
    depth = 6
    edge_features = 6
    node_features = 6

    epochs = 200
    learning_rate = 0.0001
    scheduler_step = 50
    scheduler_gamma = 0.5

    if m==800:
        batch_size = 2
        epochs = 100

    path_model = 'model/UAI5_s'+str(s)+'_m'+ str(m)
    path_train_err = 'results/UAI5_s'+str(s)+'_m'+ str(m) + 'train.txt'
    path_test_err1 = 'results/UAI5_s'+str(s)+'_m'+ str(m)+'_mtest'+ str(mtest1) + 'test.txt'
    path_test_err2 = 'results/UAI5_s' + str(s) + '_m' + str(m) + '_mtest' + str(mtest2) + 'test.txt'
    path_test_err3 = 'results/UAI5_s' + str(s) + '_m' + str(m) + '_mtest' + str(mtest3) + 'test.txt'
    path_test_err4 = 'results/UAI5_s' + str(s) + '_m' + str(m) + '_mtest' + str(mtest4) + 'test.txt'


    t1 = default_timer()


    reader = MatReader(TRAIN_PATH)
    train_a = reader.read_field('coeff')[:ntrain,::r,::r].reshape(ntrain,-1)
    train_a_smooth = reader.read_field('Kcoeff')[:ntrain,::r,::r].reshape(ntrain,-1)
    train_a_gradx = reader.read_field('Kcoeff_x')[:ntrain,::r,::r].reshape(ntrain,-1)
    train_a_grady = reader.read_field('Kcoeff_y')[:ntrain,::r,::r].reshape(ntrain,-1)
    train_u = reader.read_field('sol')[:ntrain,::r,::r].reshape(ntrain,-1)

    reader.load_file(TEST_PATH)
    test_a = reader.read_field('coeff')[:ntest,::r,::r].reshape(ntest,-1)
    test_a_smooth = reader.read_field('Kcoeff')[:ntest,::r,::r].reshape(ntest,-1)
    test_a_gradx = reader.read_field('Kcoeff_x')[:ntest,::r,::r].reshape(ntest,-1)
    test_a_grady = reader.read_field('Kcoeff_y')[:ntest,::r,::r].reshape(ntest,-1)
    test_u = reader.read_field('sol')[:ntest,::r,::r].reshape(ntest,-1)


    a_normalizer = GaussianNormalizer(train_a)
    train_a = a_normalizer.encode(train_a)
    test_a = a_normalizer.encode(test_a)
    as_normalizer = GaussianNormalizer(train_a_smooth)
    train_a_smooth = as_normalizer.encode(train_a_smooth)
    test_a_smooth = as_normalizer.encode(test_a_smooth)
    agx_normalizer = GaussianNormalizer(train_a_gradx)
    train_a_gradx = agx_normalizer.encode(train_a_gradx)
    test_a_gradx = agx_normalizer.encode(test_a_gradx)
    agy_normalizer = GaussianNormalizer(train_a_grady)
    train_a_grady = agy_normalizer.encode(train_a_grady)
    test_a_grady = agy_normalizer.encode(test_a_grady)

    u_normalizer = UnitGaussianNormalizer(train_u)
    train_u = u_normalizer.encode(train_u)
    # test_u = y_normalizer.encode(test_u)



    meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=m)
    data_train = []
    for j in range(ntrain):
        for i in range(k):
            idx = meshgenerator.sample()
            grid = meshgenerator.get_grid()
            edge_index = meshgenerator.ball_connectivity(radius_train)
            edge_attr = meshgenerator.attributes(theta=train_a[j,:])
            #data_train.append(Data(x=init_point.clone().view(-1,1), y=train_y[j,:], edge_index=edge_index, edge_attr=edge_attr))
            data_train.append(Data(x=torch.cat([grid, train_a[j, idx].reshape(-1, 1),
                                                train_a_smooth[j, idx].reshape(-1, 1), train_a_gradx[j, idx].reshape(-1, 1),
                                                train_a_grady[j, idx].reshape(-1, 1)
                                                ], dim=1),
                                   y=train_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
                                   ))


    meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=mtest1)
    data_test1 = []
    for j in range(ntest):
        idx = meshgenerator.sample()
        grid = meshgenerator.get_grid()
        edge_index = meshgenerator.ball_connectivity(radius_test)
        edge_attr = meshgenerator.attributes(theta=test_a[j,:])
        data_test1.append(Data(x=torch.cat([grid, test_a[j, idx].reshape(-1, 1),
                                           test_a_smooth[j, idx].reshape(-1, 1), test_a_gradx[j, idx].reshape(-1, 1),
                                           test_a_grady[j, idx].reshape(-1, 1)
                                           ], dim=1),
                              y=test_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
                              ))
    #
    meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=mtest2)
    data_test2 = []
    for j in range(ntest):
        idx = meshgenerator.sample()
        grid = meshgenerator.get_grid()
        edge_index = meshgenerator.ball_connectivity(radius_test)
        edge_attr = meshgenerator.attributes(theta=test_a[j,:])
        data_test2.append(Data(x=torch.cat([grid, test_a[j, idx].reshape(-1, 1),
                                           test_a_smooth[j, idx].reshape(-1, 1), test_a_gradx[j, idx].reshape(-1, 1),
                                           test_a_grady[j, idx].reshape(-1, 1)
                                           ], dim=1),
                              y=test_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
                              ))
    #
    meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=mtest3)
    data_test3 = []
    for j in range(ntest):
        idx = meshgenerator.sample()
        grid = meshgenerator.get_grid()
        edge_index = meshgenerator.ball_connectivity(radius_test)
        edge_attr = meshgenerator.attributes(theta=test_a[j,:])
        data_test3.append(Data(x=torch.cat([grid, test_a[j, idx].reshape(-1, 1),
                                           test_a_smooth[j, idx].reshape(-1, 1), test_a_gradx[j, idx].reshape(-1, 1),
                                           test_a_grady[j, idx].reshape(-1, 1)
                                           ], dim=1),
                              y=test_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
                              ))
    #
    meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=mtest4)
    data_test4 = []
    for j in range(ntest):
        idx = meshgenerator.sample()
        grid = meshgenerator.get_grid()
        edge_index = meshgenerator.ball_connectivity(radius_test)
        edge_attr = meshgenerator.attributes(theta=test_a[j,:])
        data_test4.append(Data(x=torch.cat([grid, test_a[j, idx].reshape(-1, 1),
                                           test_a_smooth[j, idx].reshape(-1, 1), test_a_gradx[j, idx].reshape(-1, 1),
                                           test_a_grady[j, idx].reshape(-1, 1)
                                           ], dim=1),
                              y=test_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
                              ))
    #
    train_loader = DataLoader(data_train, batch_size=batch_size, shuffle=True)
    test_loader1 = DataLoader(data_test1, batch_size=batch_size2, shuffle=False)
    test_loader2 = DataLoader(data_test2, batch_size=batch_size2, shuffle=False)
    test_loader3 = DataLoader(data_test3, batch_size=batch_size2, shuffle=False)
    test_loader4 = DataLoader(data_test4, batch_size=2, shuffle=False)

    t2 = default_timer()

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

    model = KernelNN3(width, ker_width,depth,edge_features,in_width=node_features).cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=5e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma)

    myloss = LpLoss(size_average=False)
    u_normalizer.cuda()
    ttrain = np.zeros((epochs, ))
    ttest1 = np.zeros((epochs,))
    ttest2 = np.zeros((epochs,))
    ttest3 = np.zeros((epochs,))
    ttest4 = np.zeros((epochs,))
    model.train()
    for ep in range(epochs):
        t1 = default_timer()
        train_mse = 0.0
        train_l2 = 0.0
        for batch in train_loader:
            batch = batch.to(device)

            optimizer.zero_grad()
            out = model(batch)
            mse = F.mse_loss(out.view(-1, 1), batch.y.view(-1,1))
            mse.backward()

            l2 = myloss(u_normalizer.decode(out.view(batch_size, -1), sample_idx=batch.sample_idx.view(batch_size, -1)),
                        u_normalizer.decode(batch.y.view(batch_size, -1),
                                            sample_idx=batch.sample_idx.view(batch_size, -1)))
            optimizer.step()
            train_mse += mse.item()
            train_l2 += l2.item()

        scheduler.step()
        t2 = default_timer()

        model.eval()
        test1_l2 = 0.0
        test2_l2 = 0.0
        test3_l2 = 0.0
        test4_l2 = 0.0
        with torch.no_grad():
            for batch in test_loader1:
                batch = batch.to(device)
                out = model(batch)
                out = u_normalizer.decode(out.view(batch_size2,-1), sample_idx=batch.sample_idx.view(batch_size2,-1))
                test1_l2 += myloss(out, batch.y.view(batch_size2, -1)).item()
            for batch in test_loader2:
                batch = batch.to(device)
                out = model(batch)
                out = u_normalizer.decode(out.view(batch_size2,-1), sample_idx=batch.sample_idx.view(batch_size2,-1))
                test2_l2 += myloss(out, batch.y.view(batch_size2, -1)).item()
            for batch in test_loader3:
                batch = batch.to(device)
                out = model(batch)
                out = u_normalizer.decode(out.view(batch_size2,-1), sample_idx=batch.sample_idx.view(batch_size2,-1))
                test3_l2 += myloss(out, batch.y.view(batch_size2, -1)).item()
            for batch in test_loader4:
                batch = batch.to(device)
                out = model(batch)
                out = u_normalizer.decode(out.view(2,-1), sample_idx=batch.sample_idx.view(2,-1))
                test4_l2 += myloss(out, batch.y.view(2, -1)).item()

        ttrain[ep] = train_l2/(ntrain * k)
        ttest1[ep] = test1_l2/ntest
        ttest2[ep] = test2_l2 / ntest
        ttest3[ep] = test3_l2 / ntest
        ttest4[ep] = test4_l2 / ntest

        print(m, t2-t1, train_mse/len(train_loader), train_l2/(ntrain * k))
        print(test1_l2/ntest, test2_l2/ntest, test3_l2/ntest, test4_l2/ntest)

    np.savetxt(path_train_err, ttrain)
    np.savetxt(path_test_err1, ttest1)
    np.savetxt(path_test_err2, ttest2)
    np.savetxt(path_test_err3, ttest3)
    np.savetxt(path_test_err4, ttest4)
    torch.save(model, path_model)
