#%%
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))

# %%
[(36, 'torch.device', 'torch.device', (['"""cuda"""'], {}), False, 'import torch\n'), (38, 'numpy.random.seed', 'np.random.seed', (['(904)'], {}), True, 'import numpy as np\n'), (140, 'pandas.read_csv', 'pd.read_csv', (['datapath'], {}), True, 'import pandas as pd\n'), (142, 'read_data.create_data', 'create_data', (['df', 'value'], {}), False, 'from read_data import create_data\n'), (145, 'sklearn.metrics.accuracy_score', 'accuracy_score', (['y_pre', 'y_true'], {}), False, 'from sklearn.metrics import accuracy_score\n'), (153, 'torch.utils.data.DataLoader', 'DataLoader', (['trainset'], {'batch_size': 'batch_size', 'shuffle': '(True)'}), False, 'from torch.utils.data import DataLoader, Dataset\n'), (41, 'sklearn.ensemble.RandomForestClassifier', 'RandomForestClassifier', ([], {'n_estimators': '(100)'}), False, 'from sklearn.ensemble import RandomForestClassifier\n'), (41, 'sklearn.naive_bayes.GaussianNB', 'GaussianNB', ([], {}), False, 'from sklearn.naive_bayes import GaussianNB\n'), (41, 'sklearn.svm.SVC', 'SVC', ([], {'kernel': '"""linear"""'}), False, 'from sklearn.svm import SVC, LinearSVC\n'), (41, 'sklearn.svm.LinearSVC', 'LinearSVC', ([], {}), False, 'from sklearn.svm import SVC, LinearSVC\n'), (59, 'numpy.row_stack', 'np.row_stack', (['label_lis'], {}), True, 'import numpy as np\n'), (73, 'torch.nn.Linear', 'nn.Linear', (['(6)', 'dim[0]'], {}), True, 'import torch.nn as nn\n'), (74, 'torch.nn.Linear', 'nn.Linear', (['dim[0]', 'dim[1]'], {}), True, 'import torch.nn as nn\n'), (75, 'torch.nn.Linear', 'nn.Linear', (['(2 * dim[1])', '(6)'], {}), True, 'import torch.nn as nn\n'), (77, 'torch.nn.Linear', 'nn.Linear', (['(4)', 'dim[0]'], {}), True, 'import torch.nn as nn\n'), (78, 'torch.nn.Linear', 'nn.Linear', (['dim[0]', 'dim[1]'], {}), True, 'import torch.nn as nn\n'), (79, 'torch.nn.Linear', 'nn.Linear', (['(2 * dim[1])', '(4)'], {}), True, 'import torch.nn as nn\n'), (81, 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['dim[0]'], {}), True, 'import torch.nn as nn\n'), (82, 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['dim[1]'], {}), True, 'import torch.nn as nn\n'), (117, 'torch.cat', 'torch.cat', (['[z1, z2]'], {'dim': '(1)'}), False, 'import torch\n'), (118, 'torch.cat', 'torch.cat', (['[z2, z1]'], {'dim': '(1)'}), False, 'import torch\n'), (50, 'numpy.zeros', 'np.zeros', (['testdata.shape[0]'], {}), True, 'import numpy as np\n'), (52, 'numpy.row_stack', 'np.row_stack', (['res_list'], {}), True, 'import numpy as np\n'), (99, 'torch.nn.functional.mse_loss', 'F.mse_loss', (['a', 'a_'], {}), True, 'import torch.nn.functional as F\n'), (100, 'torch.nn.functional.mse_loss', 'F.mse_loss', (['x', 'x_'], {}), True, 'import torch.nn.functional as F\n'), (184, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (194, 'sklearn.metrics.accuracy_score', 'accuracy_score', (['label_lis', 'testlabel'], {}), False, 'from sklearn.metrics import accuracy_score\n'), (58, 'numpy.unique', 'np.unique', (['testlabel'], {}), True, 'import numpy as np\n'), (57, 'numpy.square', 'np.square', (['(attribute_matrix - pre_res)'], {}), True, 'import numpy as np\n'), (147, 'torch.from_numpy', 'torch.from_numpy', (['traindata'], {}), False, 'import torch\n'), (150, 'torch.from_numpy', 'torch.from_numpy', (['testdata'], {}), False, 'import torch\n'), (152, 'torch.from_numpy', 'torch.from_numpy', (['train_attributelabel'], {}), False, 'import torch\n'), (92, 'torch.diagonal', 'torch.diagonal', (['c'], {}), False, 'import torch\n'), (185, 'torch.from_numpy', 'torch.from_numpy', (['train_attributelabel'], {}), False, 'import torch\n'), (189, 'torch.from_numpy', 'torch.from_numpy', (['test_attributelabel'], {}), False, 'import torch\n'), (187, 'torch.cat', 'torch.cat', (['[f_train, traindata]'], {'dim': '(1)'}), False, 'import torch\n'), (191, 'torch.cat', 'torch.cat', (['[f_test, testdata]'], {'dim': '(1)'}), False, 'import torch\n')]