#%%
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC,LinearSVC
from torch import device
from torch.optim import optimizer
from torch.utils.data import DataLoader, Dataset

from read_data import create_data


#%%
class my_dataset(Dataset):
    def __init__(self,data,attribute_label):
        super(my_dataset,self).__init__()
        self.data=data
        self.attribute_label=attribute_label

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, index):
        batch_data=self.data[index]
        batch_label=self.attribute_label[index]
        return batch_data,batch_label
        
#%%
device=torch.device('cuda')

np.random.seed(904)

def pre_model(model, traindata, train_attributelabel, testdata, testlabel, attribute_matrix):
    model_dict = {'rf': RandomForestClassifier(n_estimators=100),'NB': GaussianNB(),'SVC_linear': SVC(kernel='linear'),'LinearSVC':LinearSVC()}

    res_list = []
    for i in range(train_attributelabel.shape[1]):
        clf = model_dict[model]
        if max(train_attributelabel[:, i]) != 0:
            clf.fit(traindata, train_attributelabel[:, i])
            res = clf.predict(testdata)
        else:
            res = np.zeros(testdata.shape[0])
        res_list.append(res.T)
    test_pre_attribute = np.mat(np.row_stack(res_list)).T

    label_lis = []
    for i in range(test_pre_attribute.shape[0]):
        pre_res = test_pre_attribute[i, :]
        loc = (np.sum(np.square(attribute_matrix - pre_res), axis=1)).argmin()
        label_lis.append(np.unique(testlabel)[loc])
    label_lis = np.mat(np.row_stack(label_lis))
    return test_pre_attribute,label_lis, testlabel

#%%
def off_diagonal(x):
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()

#%%
class Embedding_Net(nn.Module):

    def __init__(self,dim,lambda_):
        super(Embedding_Net,self).__init__()
        self.l11=nn.Linear(6,dim[0])
        self.l12=nn.Linear(dim[0],dim[1])
        self.l13=nn.Linear(2*dim[1],6)

        self.l21=nn.Linear(4,dim[0])
        self.l22=nn.Linear(dim[0],dim[1])
        self.l23=nn.Linear(2*dim[1],4)

        self.bn1=nn.BatchNorm1d(dim[0])
        self.bn2=nn.BatchNorm1d(dim[1])
        
        self.lambda_=lambda_


    def compability_loss(self,z1,z2):
        N,D=z1.shape

        c=self.bn2(z1).T @ self.bn2(z2)/N

        on_diag=torch.diagonal(c).add_(-1).pow_(2).sum()
        off_diag=off_diagonal(c).pow_(2).sum()
        loss=on_diag+self.lambda_[3]*off_diag

        return loss

    def compute_loss(self,z1,z2,x,a,x_,a_):
        loss_R1=self.lambda_[0]*F.mse_loss(a,a_)
        loss_R2=self.lambda_[1]*F.mse_loss(x,x_)
        loss_CM=self.compability_loss(z1,z2)
        loss_CM=self.lambda_[2]*loss_CM
        loss=loss_R1+loss_R2+loss_CM
        return loss_R1,loss_R2,loss_CM,loss

    def transform(self,x,a):
        z1=self.l11(x)
        z1=torch.relu(self.bn1(z1))
        z1=self.l12(z1)

        z2=self.l21(a)
        z2=torch.relu(self.bn1(z2))
        z2=self.l22(z2)
        return z1,z2

    def reconstruction(self,z1,z2):
        f1=torch.cat([z1,z2],dim=1)
        f2=torch.cat([z2,z1],dim=1)
        x_=self.l13(f1)
        a_=torch.sigmoid(self.l23(f2))
        return x_,a_


    def forward(self,x,a):
        z1,z2=self.transform(x,a)
        x_,a_=self.reconstruction(z1,z2)

        loss_R1,loss_R2,loss_CM,loss=self.compute_loss(z1,z2,x,a,x_,a_)
        package={'z1':z1,'z2':z2,'x':x,'x_':x_,'r1':loss_R1,
                'r2':loss_R2,'cm':loss_CM,'loss':loss}
        
        return package

#%%
datapath='data/classData.csv'
modes=['NB'] #'rf'
test_classes={'test_class':[2,3]}
for key,value in test_classes.items():
    print('========================================{}:[{}:{}]========================================='.format(modes,key,value))
    df = pd.read_csv(datapath)
    df['fault_type'] = df['G'].astype('str') + df['C'].astype('str') + df['B'].astype('str') + df['A'].astype('str')
    traindata,trainlabel,train_attributelabel, train_attributematrix,testdata,testlabel,test_attributelabel,test_attributematrix,attribute_matrix=create_data(df,value)

    _,y_pre,y_true=pre_model(modes[0], traindata, train_attributelabel, testdata, testlabel, test_attributematrix)
    original_acc=accuracy_score(y_pre,y_true)

    traindata=torch.from_numpy(traindata).float().to(device)
    label=torch.from_numpy(trainlabel.squeeze()).long().to(device)

    testdata=torch.from_numpy(testdata).float().to(device)
    batch_size=400
    trainset=my_dataset(traindata,torch.from_numpy(train_attributelabel).float().to(device))
    train_loader=DataLoader(trainset,batch_size=batch_size,shuffle=True)

    lambda_=[1,1e-5,1,0.25]
    dim=[6,12]
    model=Embedding_Net(dim,lambda_=lambda_)
    model.to(device)

    optimizer=optim.RMSprop(model.parameters(),lr=1e-2)

    L1,L2,L3,L=[],[],[],[]
    model.train()

    accs=[]
    best_acc=0
    for epoch in range(200):
        model.train()
        for batch,(batch_data,batch_label) in enumerate(train_loader):

            optimizer.zero_grad() 
            package=model(batch_data,batch_label)
            loss_R1,loss_R2,loss_CM,loss=package['r1'],package['r2'],package['cm'],package['loss']
            
            loss.backward()
            optimizer.step()

            L1.append(loss_R1.item())
            L2.append(loss_R2.item())
            L3.append(loss_CM.item())
            L.append(loss.item())

        model.eval()
        with torch.no_grad():
            train_package=model(traindata,torch.from_numpy(train_attributelabel).float().to(device))
            f_train=train_package['z1']
            f_train=torch.cat([f_train,traindata],dim=1).detach().cpu().numpy()

            test_package=model(testdata,torch.from_numpy(test_attributelabel).float().to(device))
            f_test=test_package['z1']
            f_test=torch.cat([f_test,testdata],dim=1).detach().cpu().numpy()

            test_preattribute,label_lis, testlabel=pre_model(modes[0], f_train, train_attributelabel, f_test, testlabel, test_attributematrix)
            acc=accuracy_score(label_lis, testlabel)
            accs.append(acc)
            if acc>best_acc:
                best_acc=acc

        print('epoch:{:d}, best_acc:{:.4f}'.format(epoch,best_acc))

print('finished! FDAT:{:.4f}, SCE:{:.4f}'.format(original_acc,best_acc))

# %%
[((947, 967), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (959, 967), False, 'import torch\n'), ((971, 990), 'numpy.random.seed', 'np.random.seed', (['(904)'], {}), '(904)\n', (985, 990), True, 'import numpy as np\n'), ((4261, 4282), 'pandas.read_csv', 'pd.read_csv', (['datapath'], {}), '(datapath)\n', (4272, 4282), True, 'import pandas as pd\n'), ((4548, 4570), 'read_data.create_data', 'create_data', (['df', 'value'], {}), '(df, value)\n', (4559, 4570), False, 'from read_data import create_data\n'), ((4706, 4735), 'sklearn.metrics.accuracy_score', 'accuracy_score', (['y_pre', 'y_true'], {}), '(y_pre, y_true)\n', (4720, 4735), False, 'from sklearn.metrics import accuracy_score\n'), ((5061, 5118), 'torch.utils.data.DataLoader', 'DataLoader', (['trainset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), '(trainset, batch_size=batch_size, shuffle=True)\n', (5071, 5118), False, 'from torch.utils.data import DataLoader, Dataset\n'), ((1113, 1153), 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(100)'}), '(n_estimators=100)\n', (1135, 1153), False, 'from sklearn.ensemble import RandomForestClassifier\n'), ((1160, 1172), 'sklearn.naive_bayes.GaussianNB', 'GaussianNB', ([], {}), '()\n', (1170, 1172), False, 'from sklearn.naive_bayes import GaussianNB\n'), ((1187, 1207), 'sklearn.svm.SVC', 'SVC', ([], {'kernel': '"""linear"""'}), "(kernel='linear')\n", (1190, 1207), False, 'from sklearn.svm import SVC, LinearSVC\n'), ((1220, 1231), 'sklearn.svm.LinearSVC', 'LinearSVC', ([], {}), '()\n', (1229, 1231), False, 'from sklearn.svm import SVC, LinearSVC\n'), ((1916, 1939), 'numpy.row_stack', 'np.row_stack', (['label_lis'], {}), '(label_lis)\n', (1928, 1939), True, 'import numpy as np\n'), ((2269, 2289), 'torch.nn.Linear', 'nn.Linear', (['(6)', 'dim[0]'], {}), '(6, dim[0])\n', (2278, 2289), True, 'import torch.nn as nn\n'), ((2307, 2332), 'torch.nn.Linear', 'nn.Linear', (['dim[0]', 'dim[1]'], {}), '(dim[0], dim[1])\n', (2316, 2332), True, 'import torch.nn as nn\n'), ((2350, 2374), 'torch.nn.Linear', 'nn.Linear', (['(2 * dim[1])', '(6)'], {}), '(2 * dim[1], 6)\n', (2359, 2374), True, 'import torch.nn as nn\n'), ((2392, 2412), 'torch.nn.Linear', 'nn.Linear', (['(4)', 'dim[0]'], {}), '(4, dim[0])\n', (2401, 2412), True, 'import torch.nn as nn\n'), ((2430, 2455), 'torch.nn.Linear', 'nn.Linear', (['dim[0]', 'dim[1]'], {}), '(dim[0], dim[1])\n', (2439, 2455), True, 'import torch.nn as nn\n'), ((2473, 2497), 'torch.nn.Linear', 'nn.Linear', (['(2 * dim[1])', '(4)'], {}), '(2 * dim[1], 4)\n', (2482, 2497), True, 'import torch.nn as nn\n'), ((2515, 2537), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['dim[0]'], {}), '(dim[0])\n', (2529, 2537), True, 'import torch.nn as nn\n'), ((2556, 2578), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['dim[1]'], {}), '(dim[1])\n', (2570, 2578), True, 'import torch.nn as nn\n'), ((3503, 3529), 'torch.cat', 'torch.cat', (['[z1, z2]'], {'dim': '(1)'}), '([z1, z2], dim=1)\n', (3512, 3529), False, 'import torch\n'), ((3540, 3566), 'torch.cat', 'torch.cat', (['[z2, z1]'], {'dim': '(1)'}), '([z2, z1], dim=1)\n', (3549, 3566), False, 'import torch\n'), ((1524, 1551), 'numpy.zeros', 'np.zeros', (['testdata.shape[0]'], {}), '(testdata.shape[0])\n', (1532, 1551), True, 'import numpy as np\n'), ((1617, 1639), 'numpy.row_stack', 'np.row_stack', (['res_list'], {}), '(res_list)\n', (1629, 1639), True, 'import numpy as np\n'), ((2987, 3004), 'torch.nn.functional.mse_loss', 'F.mse_loss', (['a', 'a_'], {}), '(a, a_)\n', (2997, 3004), True, 'import torch.nn.functional as F\n'), ((3037, 3054), 'torch.nn.functional.mse_loss', 'F.mse_loss', (['x', 'x_'], {}), '(x, x_)\n', (3047, 3054), True, 'import torch.nn.functional as F\n'), ((5511, 5532), 'torch.optim.optimizer.zero_grad', 'optimizer.zero_grad', ([], {}), '()\n', (5530, 5532), False, 'from torch.optim import optimizer\n'), ((5741, 5757), 'torch.optim.optimizer.step', 'optimizer.step', ([], {}), '()\n', (5755, 5757), False, 'from torch.optim import optimizer\n'), ((5950, 5965), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (5963, 5965), False, 'import torch\n'), ((6572, 6608), 'sklearn.metrics.accuracy_score', 'accuracy_score', (['label_lis', 'testlabel'], {}), '(label_lis, testlabel)\n', (6586, 6608), False, 'from sklearn.metrics import accuracy_score\n'), ((1865, 1885), 'numpy.unique', 'np.unique', (['testlabel'], {}), '(testlabel)\n', (1874, 1885), True, 'import numpy as np\n'), ((1782, 1819), 'numpy.square', 'np.square', (['(attribute_matrix - pre_res)'], {}), '(attribute_matrix - pre_res)\n', (1791, 1819), True, 'import numpy as np\n'), ((4752, 4779), 'torch.from_numpy', 'torch.from_numpy', (['traindata'], {}), '(traindata)\n', (4768, 4779), False, 'import torch\n'), ((4883, 4909), 'torch.from_numpy', 'torch.from_numpy', (['testdata'], {}), '(testdata)\n', (4899, 4909), False, 'import torch\n'), ((4984, 5022), 'torch.from_numpy', 'torch.from_numpy', (['train_attributelabel'], {}), '(train_attributelabel)\n', (5000, 5022), False, 'import torch\n'), ((2748, 2765), 'torch.diagonal', 'torch.diagonal', (['c'], {}), '(c)\n', (2762, 2765), False, 'import torch\n'), ((6010, 6048), 'torch.from_numpy', 'torch.from_numpy', (['train_attributelabel'], {}), '(train_attributelabel)\n', (6026, 6048), False, 'import torch\n'), ((6234, 6271), 'torch.from_numpy', 'torch.from_numpy', (['test_attributelabel'], {}), '(test_attributelabel)\n', (6250, 6271), False, 'import torch\n'), ((6131, 6169), 'torch.cat', 'torch.cat', (['[f_train, traindata]'], {'dim': '(1)'}), '([f_train, traindata], dim=1)\n', (6140, 6169), False, 'import torch\n'), ((6351, 6387), 'torch.cat', 'torch.cat', (['[f_test, testdata]'], {'dim': '(1)'}), '([f_test, testdata], dim=1)\n', (6360, 6387), False, 'import torch\n')]