import fitz  # pyMuPDF里面的fitz包，不要与pip install fitz混淆
import numpy as np

ocr = get_ocr()
doc = fitz.open(filepath)
resp = ""

b_unit = tqdm.tqdm(
    total=doc.page_count, desc="RapidOCRPDFLoader context page index: 0"
)
for i, page in enumerate(doc):
    b_unit.set_description(
        "RapidOCRPDFLoader context page index: {}".format(i)
    )
    b_unit.refresh()
    text = page.get_text("")
    resp += text + "\n"

    img_list = page.get_image_info(xrefs=True)
    for img in img_list:
        if xref := img.get("xref"):
                bbox = img["bbox"]
                # 检查图片尺寸是否超过设定的阈值
                if (bbox[2] - bbox[0]) / (page.rect.width) < PDF_OCR_THRESHOLD[
                    0
                ] or (bbox[3] - bbox[1]) / (
                    page.rect.height
                ) < PDF_OCR_THRESHOLD[1]:
                    continue
                    pix = fitz.Pixmap(doc, xref)
                    samples = pix.samples
                    if int(page.rotation) != 0:  # 如果Page有旋转角度，则旋转图片
                        img_array = np.frombuffer(
                            pix.samples, dtype=np.uint8
                        ).reshape(pix.height, pix.width, -1)
                        tmp_img = Image.fromarray(img_array)
                        ori_img = cv2.cvtColor(np.array(tmp_img), cv2.COLOR_RGB2BGR)
                        rot_img = rotate_img(img=ori_img, angle=360 - page.rotation)
                        img_array = cv2.cvtColor(rot_img, cv2.COLOR_RGB2BGR)
                        else:
                            img_array = np.frombuffer(
                                pix.samples, dtype=np.uint8
                            ).reshape(pix.height, pix.width, -1)

                            result, _ = ocr(img_array)
                            if result:
                                ocr_result = [line[1] for line in result]
                                resp += "\n".join(ocr_result)

import pypdfium2
    with blob.as_bytes_io() as file_path:
        pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)
        try:
            for page_number, page in enumerate(pdf_reader):
                text_page = page.get_textpage()
                content = text_page.get_text_range()
                text_page.close()
                page.close()
                metadata = {"source": blob.source, "page": page_number}
                yield Document(page_content=content, metadata=metadata)
                finally:
                    pdf_reader.close()

from unstructured.partition.pdf import partition_pdf
# Get elements
raw_pdf_elements = partition_pdf(
    filename=test_file,
    # Using pdf format to find embedded image blocks
    extract_images_in_pdf=True,
    # Use layout model (YOLOX) to get bounding boxes (for tables) and find titles
    # Titles are any sub-section of the document
    infer_table_structure=True,
    # Post processing to aggregate text once we have the title
    chunking_strategy="by_title",
    # Chunking params to aggregate text blocks
    # Attempt to create a new chunk 3800 chars
    # Attempt to keep chunks > 2000 chars
    # Hard max on chunks
    max_characters=4000,
    new_after_n_chars=3800,
    combine_text_under_n_chars=2000,
)

tables = [el for el in raw_pdf_elements if el.category == 'Table']

print(tables[0].text)

print(tables[0].metadata.text_as_html)

class Transformer(nn.Module):
    '''
    define the whole architecture of Transformer in:
        Vaswani et al. Attention is All You Need. NIPS 2017.
    '''
    def __init__(self, num_heads=8, d_model=512, d_ff=2048, num_enc_layers=6, num_dec_layers=6,
                 drop_rate=0.1, warmup_steps=400, pos_encoding_type='sinusoid',
                 ls_epsilon=0.1, use_label_smoothing=True,
                 model_name='transformer', tf_sess_config=None, **kwargs):
        super().__init__()
        self.h = num_heads
        self.d_model = d_model
        self.d_ff = d_ff

        self.num_enc_layers = num_enc_layers
        self.num_dec_layers = num_dec_layers

        # Dropout regularization: added in every sublayer before layer_norm(...) and
        # applied to embedding + positional encoding.
        self.drop_rate = drop_rate

        # Label smoothing epsilon
        self.ls_epsilon = ls_epsilon
        self.use_label_smoothing = use_label_smoothing
        self.pos_encoding_type = pos_encoding_type

        # For computing the learning rate
        self.warmup_steps = warmup_steps

        self.config = dict(
            num_heads=self.h,
            d_model=self.d_model,
            d_ff=self.d_ff,
            num_enc_layers=self.num_enc_layers,
            num_dec_layers=self.num_dec_layers,
            drop_rate=self.drop_rate,
            warmup_steps=self.warmup_steps,
            ls_epsilon=self.ls_epsilon,
            use_label_smoothing=self.use_label_smoothing,
            pos_encoding_type=self.pos_encoding_type,
            model_name=self.model_name,
            tf_sess_config=self.tf_sess_config,
        )
        def forward(self, src, tgt, src_mask, tgt_mask): # 这里进行拼接，将encoder和decoder两大模块拼接在一起
            enc_out = self.encoder(src, src_mask)
            dec_out = self.decoder(enc_out, src_mask, tgt, tgt_mask)
            return dec_out
        

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class TransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_enc_layers) -> None:
        self.num_enc_layers = num_enc_layers
        self.encoder_layers = clones(encoder_layer,num_enc_layers) # 将encoder_layer复制6次

    def forward(self, src, src_mask):
        out = src
        for layer in self.encoder_layers:
            out = layer(out, src_mask)
        return out

out = inp  # now, (batch, seq_len, embed_size)
with tf.variable_scope(scope):
    for i in range(self.num_enc_layers):
        out = self.encoder_layer(out, input_mask, f'enc_{i}')
        return out

class TransformerEncoderLayer(nn.Module):
    """
    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
    About:

    """
    # One multi-head attention + one feed-forward
    def __init__(self, d_model, n_head, dim_feedforward, dropout = 0.1) -> None:
        super().__init__()
        self.self_attn = MultiheadAttention(d_model, n_head)
        self.norm_1 = nn.LayerNorm(d_model)
        # Implementation of Feedforward model(Two linear transformation together with one dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward, )
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm_2 = nn.LayerNorm(d_model)

    def __ff_block(self, x):
        # feed forward layer contains two linear
        out = F.relu(self.linear1(x))
        out = self.dropout(out)
        out = self.linear2(out)
        
    def forward(self, src, src_mask):
        out = src
        out = self.norm_1(out + self.self_attn(out, src_mask))# 这里在pytorch的官方实现中在self_attention后还加了一个dropout
        out = self.norm_2(out + self.__ff_block(out))
    
        return out

def multihead_attention(self, query, memory=None, mask=None, scope='attn'):
        """
        Args:
            query (tf.tensor): of shape (batch, q_size, d_model)
            memory (tf.tensor): of shape (batch, m_size, d_model)
            mask (tf.tensor): shape (batch, q_size, k_size)

        Returns:h
            a tensor of shape (bs, q_size, d_model)
        """
        if memory is None:
            memory = query

        with tf.variable_scope(scope):
            # Linear project to d_model dimension: [batch, q_size/k_size, d_model]
            Q = tf.layers.dense(query, self.d_model, activation=tf.nn.relu)
            K = tf.layers.dense(memory, self.d_model, activation=tf.nn.relu)
            V = tf.layers.dense(memory, self.d_model, activation=tf.nn.relu)

            # Split the matrix to multiple heads and then concatenate to have a larger
            # batch size: [h*batch, q_size/k_size, d_model/num_heads]
            Q_split = tf.concat(tf.split(Q, self.h, axis=2), axis=0)
            K_split = tf.concat(tf.split(K, self.h, axis=2), axis=0)
            V_split = tf.concat(tf.split(V, self.h, axis=2), axis=0)
            mask_split = tf.tile(mask, [self.h, 1, 1])

            # Apply scaled dot product attention
            out = self.scaled_dot_product_attention(Q_split, K_split, V_split, mask=mask_split)

            # Merge the multi-head back to the original shape
            out = tf.concat(tf.split(out, self.h, axis=0), axis=2)  # [bs, q_size, d_model]

            # The final linear layer and dropout.
            # out = tf.layers.dense(out, self.d_model)
            # out = tf.layers.dropout(out, rate=self.drop_rate, training=self._is_training)

        return out