import torch
import dashscope
from http import HTTPStatus
from dashscope import TextEmbedding
from torch import nn
import torch.optim as optim
import pandas as pd
import numpy as np
dashscope.api_key = 'sk-44ccc9ab5e754eddb545cade12b632cf'
cache = {}
answerCache = []


def getem(question):
    global cache
    if question in cache.keys():
        return cache[question]
    resp = TextEmbedding.call(model=TextEmbedding.Models.text_embedding_v1,
                              input=question,
                              text_type='query')
    if resp.status_code == HTTPStatus.OK:
        cache[question] = resp['output']['embeddings'][0]['embedding']
        return resp['output']['embeddings'][0]['embedding']


class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu2 = nn.ReLU()
        self.conv3 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu3 = nn.ReLU()
        self.fc = nn.Linear(1536 * 3, 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        X1 = self.conv1(x)
        X1 = self.relu1(X1)
        X2 = self.conv2(x)
        X2 = self.relu1(X2)
        X3 = self.conv3(x)
        X3 = self.relu1(X3)
        X = torch.cat([X1, X2, X3], dim=2)
        X = X.view(-1, 1536 * 3)
        X = self.fc(X)
        X = self.sigmoid(X)
        return X


# 创建模型实例
model = ConvNet()

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

train = []
test = []
dataall = pd.read_csv('data.csv')
dataall = dataall.iloc[:, 1:4]

dataall = dataall.sample(frac=1)
train = dataall.iloc[0:300]
train = train.reset_index(drop=True)
test = dataall.iloc[300:]
test = test.reset_index(drop=True)
nlossLast = 0
for i in range(5):
    nloss = 0
    for k in range(len(train)):
        va = getem(train.iloc[k]['question'])
        vb = getem(train.iloc[k]['answer'])
        if train.iloc[k]['answer'] not in answerCache:
            answerCache.append(train.iloc[k]['answer'])
        trainTensor = torch.Tensor([va, vb]).reshape([1, 2, len(va)])
        output = model(trainTensor)
        # 计算损失
        if train.iloc[k]['label'] == 1:
            loss = criterion(output,
                             torch.tensor([1, 0]).float().reshape([1, 2]))
        else:
            loss = criterion(output,
                             torch.tensor([0, 1]).float().reshape([1, 2]))
        # 反向传播并更新权重
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        nloss += loss
        if k % 50 == 0:
            print(i, k, 'done')
    print('one loop done', nloss/len(train))


p = 0
for i in range(len(test)):
    va = getem(test.iloc[i]['question'])
    Scores = np.zeros(len(answerCache))
    for j in range(len(answerCache)):
        vb = getem(answerCache[j])
        testTensor = torch.Tensor([va, vb]).reshape([1, 2, len(va)])
        output = model(testTensor)
        Scores[j] = output[0][0]
    for k in range(2):
        if test.iloc[i]['label'] == 1:
            vc = test.iloc[i]['answer']
        else:
            vc = ''
        tt = Scores.argmax()
        if Scores[tt] > 0.5:
            vb = answerCache[tt]
            Scores[tt] = -1
        else:
            vb = ''
        if vb == vc:
            p += 1
            break
print(p/len(test))
