class SemanticIndex(nn.Module):
    def __init__(self, ptm):
        self.ptm = ptm
        self.reducer_linear = nn.Linear(768, 256) # 为了减少计算复杂度，降维
        self.dp = nn.Dropout(0.1)

    def get_pooled_embedding(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None):
        _, cls_embedding = self.ptm(input_ids, token_type_ids, position_ids, attention_mask)

        if self.output_emb_size > 0:
            cls_embedding = self.emb_reduce_linear(cls_embedding) # 768维太大了，使用256进行降维
        cls_embedding = self.dropout(cls_embedding)
        cls_embedding = F.normalize(cls_embedding, p=2, axis=-1)

        return cls_embedding

    def forward(
        self,
        query_input_ids,
        title_input_ids,
        query_token_type_ids=None,
        query_position_ids=None,
        query_attention_mask=None,
        title_token_type_ids=None,
        title_position_ids=None,
        title_attention_mask=None,
    ):

        query_cls_embedding = self.get_pooled_embedding(
            query_input_ids, query_token_type_ids, query_position_ids, query_attention_mask
        )

        title_cls_embedding = self.get_pooled_embedding(
            title_input_ids, title_token_type_ids, title_position_ids, title_attention_mask
        )

        cosine_sim = paddle.matmul(query_cls_embedding, title_cls_embedding, transpose_y=True)

        # Substract margin from all positive samples cosine_sim()
        margin_diag = paddle.full(
            shape=[query_cls_embedding.shape[0]], fill_value=self.margin, dtype=paddle.get_default_dtype()
        )

        cosine_sim = cosine_sim - paddle.diag(margin_diag)

        # Scale cosine to ease training converge
        cosine_sim *= self.sacle

        labels = paddle.arange(0, query_cls_embedding.shape[0], dtype="int64")
        labels = paddle.reshape(labels, shape=[-1, 1])

        loss = F.cross_entropy(input=cosine_sim, label=labels)

        return loss

    # def cosine_sim(
    #     self,
    #     query_input_ids,
    #     title_input_ids,
    #     query_token_type_ids=None,
    #     query_position_ids=None,
    #     query_attention_mask=None,
    #     title_token_type_ids=None,
    #     title_position_ids=None,
    #     title_attention_mask=None,
    # ):

    #     query_cls_embedding = self.get_pooled_embedding(
    #         query_input_ids, query_token_type_ids, query_position_ids, query_attention_mask
    #     )

    #     title_cls_embedding = self.get_pooled_embedding(
    #         title_input_ids, title_token_type_ids, title_position_ids, title_attention_mask
    #     )

    #     cosine_sim = paddle.sum(query_cls_embedding * title_cls_embedding, axis=-1)
    #     return cosine_sim

# -*- coding: utf8 -*-
#


import torch


class WordEncoder(torch.nn.Module):
    def __init__(self, n_in, p=0.1):
        super(WordEncoder, self).__init__()
        self.n_in = n_in
        self.p = p
        self.linear = torch.nn.Linear(in_features=n_in, out_features=1)
        self.dp = torch.nn.Dropout(p)

    def forward(self, x, indices):
        """
        如果indices中为padding，那么就(0, bert_out.size(1))
        :param x:
        :param indices:
        :return:
        """
        # x: BS, SL, HS
        # indices: BS, L
        start_end_indices_to_tensor = torch.tensor(indices).to(x.device)
        start_indices = start_end_indices_to_tensor[:, :, 0]
        end_indices = start_end_indices_to_tensor[:, :, 1]
        mask = torch.arange(x.shape[1]).expand(x.shape[0], len(indices[0]), x.shape[1]).to(x.device)
        attn_mask = (mask >= start_indices.unsqueeze(2)) * (mask < end_indices.unsqueeze(2))
        attn_mask = torch.log(attn_mask.to(torch.float))

        attn_scores = self.linear(x).permute(0, 2, 1)
        attn_scores = attn_scores.expand(x.shape[0], len(indices[0]), x.shape[1])
        attn_scores = attn_scores + attn_mask
        output = torch.softmax(attn_scores, dim=2).bmm(x)
        return self.dp(output)


if __name__ == '__main__':
    # 我 是 中国人
    # 我 中国人 null
    word_feature = torch.arange(60, dtype=torch.float).view(2, 5, 6)
    start_end_indices = [
        [(0, 1), (1, 2), (2, 5)],
        [(0, 1), (1, 4), (0, 0)]
    ]
    w = WordEncoder(6)
    w(word_feature, start_end_indices)

    # single
    # word_feature = torch.arange(50, dtype=torch.float).view(1, 5, 10)
    # # 那么index是：
    # start_end_indices = [[(0, 1), (1, 2), (2, 5)]]
    #
    # we = WordEncoder(n_in=10, p=0.1)
    # output = we.forward(word_feature, start_end_indices)
    # print(output)

instance={
	'id': '087d4f27c9e100ab27eca2d01f426edc',
	'text': '无畏的约翰，弗兰德斯公爵菲利普二世的儿子，勃艮第菲利普三世的老爹因为领导勃艮第派与阿玛尼克-奥尔良派争夺法王可爱查理的摄政权而大打出手，葬送了查理五世亲手建立的大好局势无畏的约翰先是刺杀了奥尔良公爵路易并夺取摄政权，后又里通外国和英王亨利五世眉来眼去',
	'tokens': ['无', '畏', '的', '约', '翰', '，', '弗', '兰', '德', '斯', '公', '爵', '菲', '利', '普', '二', '世', '的', '儿', '子', '，', '勃', '艮', '第', '菲', '利', '普', '三', '世', '的', '老', '爹', '因', '为', '领', '导', '勃', '艮', '第', '派', '与', '阿', '玛', '尼', '克', '-', '奥', '尔', '良', '派', '争', '夺', '法', '王', '可', '爱', '查', '理', '的', '摄', '政', '权', '而', '大', '打', '出', '手', '，', '葬', '送', '了', '查', '理', '五', '世', '亲', '手', '建', '立', '的', '大', '好', '局', '势', '无', '畏', '的', '约', '翰', '先', '是', '刺', '杀', '了', '奥', '尔', '良', '公', '爵', '路', '易', '并', '夺', '取', '摄', '政', '权', '，', '后', '又', '里', '通', '外', '国', '和', '英', '王', '亨', '利', '五', '世', '眉', '来', '眼', '去'],
	'entity': [{
		'type': '人物',
		'offset': [12, 13, 14, 15, 16],
		'text': '菲利普二世'
	}, {
		'type': '人物',
		'offset': [0, 1, 2, 3, 4],
		'text': '无畏的约翰'
	}],
	'relation': [{
		'type': '父亲',
		'args': [{
			'type': '人物',
			'offset': [0, 1, 2, 3, 4],
			'text': '无畏的约翰'
		}, {
			'type': '人物',
			'offset': [12, 13, 14, 15, 16],
			'text': '菲利普二世'
		}]
	}],
	'event': [],
	'spot_asoc': [{
		'span': '无畏的约翰',
		'label': '人物',
		'asoc': [
			['父亲', '菲利普二世']
		]
	}, {
		'span': '菲利普二世',
		'label': '人物',
		'asoc': []
	}],
	'spot': ['人物'],
	'asoc': ['父亲']
}

' '.join(self.tokenizer.convert_ids_to_tokens(prefix_id))
Out[3]: '<spot> 人 物 <spot> 地 理 位 置 <spot> 组 织 机 构 <asoc> 国 籍 <asoc> 妻 子 <asoc> 父 亲 <asoc> 祖 籍 <asoc> 丈 夫 <asoc> 毕 业 院 校 <asoc> 母 亲'

'<extra_id_0> <extra_id_0> 组织机构 <extra_id_5> 哈尔滨市第五中学 <extra_id_1> <extra_id_0> 组织机构 <extra_id_5> 哈五中 <extra_id_1> <extra_id_1>'

(
  (<extra_id_7>
    <extra_id_5>
    天城公主
    (祖籍 <extra_id_5> 通之)
    (丈夫 <extra_id_5> 陈国峻)
    (配偶 <extra_id_5> 陈国峻)
    (国籍 <extra_id_5> 中国))
  (<extra_id_7>
    <extra_id_5>
    陈国峻
    (妻子 <extra_id_5> 天城公主)
    (国籍 <extra_id_5> 中国)
    (毕业院校 <extra_id_5> 北京大学)
    (地理位置 <extra_id_5> 通之)
    (祖籍 <extra_id_5> 通之)
    (配偶 <extra_id_5> 天城公主))
  (<extra_id_7> <extra_id_5> 天城 (人物 <extra_id_5> 陈国峻))
  (<extra_id_7> <extra_id_5> 陈国 (国籍 <extra_id_5> 中国)))

# 实现代码
import os.path
import json
import cv2
PATH = "/layoutlm/data/funsd-dataset/dataset/training_data"

with open(os.path.join(PATH, 'annotations', '00040534.json')) as f:
    data = json.loads(f.read())
    # 选择一个box
# 找到对应的image
img_path = os.path.join(PATH, 'images', '00040534.png')
image = cv2.imread(img_path)
draw_1 = cv2.rectangle(image, (145, 98), (507, 116), (0, 0, 255), 2)
cv2.imshow('draw_0', draw_1)
cv2.waitKey(0)
cv2.destroyWindow("draw_0")

# -*- coding: utf8 -*-
#
import os

import fitz

path = 'aaa.pdf'

doc = fitz.open(path)
pages = [doc[i] for i in range(doc.page_count)]
page = pages[0]

blocks = page.get_text("dict")['blocks']
text_blocks = [block for block in blocks if block['type'] == 0]
print('对第一张操作哦')


def recalculate_bbox(span):
    a = span["ascender"] / 1000 * span["size"]
    d = span["descender"] / 1000 * span["size"]
    o = fitz.Point(span["origin"])
    r = fitz.Rect(span["bbox"])
    r.y1 = o.y - span["size"] * d / (a - d)
    r.y0 = r.y1 - span["size"]
    return r


#
for block in text_blocks:
    for lines in text_blocks:
        for line in lines['lines']:
            for span in line['spans']:
                page.draw_rect(recalculate_bbox(span), width=0.3, color=(0, 1, 1))

page.get_pixmap().save('1.png')

# -*- coding: utf8 -*-
#
import math
from typing import List


def trace_path(vec, pathes):
    if not vec:
        return pathes
    valid_pathes = []
    for item in vec[0]:
        for path in pathes:
            if item > path[-1]:
                valid_pathes.append([*path, item])

    return trace_path(vec=vec[1:], pathes=valid_pathes)


def get_all_valid_path(vec: List[List[int]]) -> List[List[int]]:
    if not vec: return []
    for v in vec:
        if not v:
            return []
    return trace_path(vec=vec[1:], pathes=[[i] for i in vec[0]])


def _sub_sum(path: List[int]) -> int:
    _sum = 0
    for index, item in enumerate(path):
        try:
            next_item = path[index + 1]
            _sum += (next_item - item)
        except IndexError:
            pass
    return _sum


def get_shortest_path(vec: List[List[int]]) -> List[int]:
    """
    输出最短的那个路径
    :param vec:
    :return:
    """
    pathes = get_all_valid_path(vec=vec)
    short_path, short_sum = [], math.inf

    for path in pathes:
        _sum = _sub_sum(path=path)
        if short_sum > _sum:
            short_path = path
            short_sum = _sum
    return short_path


if __name__ == '__main__':
    # 结果
    # 1, 2, 6
    # 1, 2, 7
    # 1, 4, 6
    # 1, 4, 7
    # 3, 4, 6
    # 3, 4, 7
    # 我这里的vec代表元素的position～
    vec = [
        [1, 3, 5],
        [2, 4],
        [6, 7]
    ]
    print(get_all_valid_path(vec))

{
        "negative": ["bad"],
        "positive": ["good", "wonderful", "great"],
}

def aggregate(label_words_logits: torch.Tensor, label_words_mask: torch.Tensor) -> torch.Tensor:
    r"""Use weight to aggregate the logits of label words.

    Args:
        label_words_logits(:obj:`torch.Tensor`): The logits of the label words.

    Returns:
        :obj:`torch.Tensor`: The aggregated logits from the label words.
    """
    label_words_logits = (label_words_logits * label_words_mask).sum(-1) / label_words_mask.sum(-1)
    return label_words_logits

output = torch.rand(2, 28996)
# [tokenizer.vocab.get(i) for i in ['bad', 'good', 'wonderful', 'great']]
# [2213, 1363, 7310, 1632]
verbalizer = torch.tensor([[2213,    0,    0],
                    [1363, 7310, 1632]])
mask = torch.tensor([[1, 0, 0],
                    [1, 1, 1]])
logits = output[:, verbalizer]
# tensor([[[0.0508, 0.0000, 0.0000],
#          [0.1770, 0.9680, 0.0091]],
#         [[0.2045, 0.0000, 0.0000],
#          [0.9504, 0.2127, 0.1417]]])
print(aggreate(logits, mask))
#  negative， positive
# tensor([[0.0508, 0.3847],
#         [0.2045, 0.4349]])

print(aggreate(logits, mask))
# tensor([[0.0508, 0.3847],
#         [0.2045, 0.4349]])

# -*- coding: utf8 -*-
#
# step1:
import torch

from openprompt.data_utils import InputExample

classes = [  # There are two classes in Sentiment Analysis, one for negative and one for positive
    "negative",
    "positive"
]
dataset = [  # For simplicity, there's only two examples
    # text_a is the input text of the data, some other datasets may have multiple input sentences in one example.
    InputExample(
        guid=0,
        text_a="Albert Einstein was one of the greatest intellects of his time.",
    ),
    InputExample(
        guid=1,
        text_a="The film was badly made.",
    ),
]
# step 2
from openprompt.plms import load_plm

plm, tokenizer, model_config, WrapperClass = load_plm("bert", "bert-base-cased")
# step3
from openprompt.prompts import ManualTemplate

promptTemplate = ManualTemplate(
    text='{"placeholder":"text_a"} It was {"mask"}',
    tokenizer=tokenizer,
)
# step4
from openprompt.prompts import ManualVerbalizer

promptVerbalizer = ManualVerbalizer(
    classes=classes,
    label_words={
        "negative": ["bad"],
        "positive": ["good", "wonderful", "great"],
    },
    tokenizer=tokenizer,
)
# step 5
from openprompt import PromptForClassification

promptModel = PromptForClassification(
    template=promptTemplate,
    plm=plm,
    verbalizer=promptVerbalizer,
)
# step 6
from openprompt import PromptDataLoader

data_loader = PromptDataLoader(
    dataset=dataset,
    tokenizer=tokenizer,
    template=promptTemplate,
    tokenizer_wrapper_class=WrapperClass,
)
# step 7
# making zero-shot inference using pretrained MLM with prompt
promptModel.eval()
with torch.no_grad():
    for batch in data_loader:
        logits = promptModel(batch)
        preds = torch.argmax(logits, dim=-1)
        print(classes[preds])
# predictions would be 1, 0 for classes 'positive', 'negative'

from paddlenlp import Taskflow


# 跨任务抽取
schema = ['法院', {'原告': '委托代理人'}, {'被告': '委托代理人'}]
ie = Taskflow('information_extraction', schema=schema)
ie.set_schema(schema)
print(ie("北京市海淀区人民法院\n民事判决书\n(199x)建初字第xxx号\n原告：张三。\n委托代理人李四，北京市 A律师事务所律师。\n被告：B公司，法定代表人王五，开发公司总经理。\n委托代理人赵六，北京市 C律师事务所律师。")) # Better print results using pprint