import streamlit as st
import pandas as pd
from PyPDF2 import PdfReader
from difflib import ndiff
import os
import tempfile
import matplotlib.pyplot as plt
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import fitz
import numpy as np

# Function to find similar paragraphs using cosine similarity
def find_similar(vector_representation, all_representations, threshold=0.5):
    similarity_matrix = cosine_similarity(vector_representation, all_representations)
    similarities = similarity_matrix[0]

    # Find indices where similarity is above the threshold
    similar_indices = np.where(similarities > threshold)[0]
    sorted_indices = sorted(similar_indices, key=lambda i: similarities[i], reverse=True)

    return similar_indices.tolist()

def process_pdf(doc, diff_line):
    data_geom = []
    data_word = []
    data_page = []

    for page_num, page in enumerate(doc['pages']):
        for block in page["blocks"]:
            for line in block["lines"]:
                temp = []
                data_page.append(page_num)
                data_geom.append(line["geometry"])
                for word in line["words"]:
                    temp.append(word["value"])
                data_word.append(" ".join(temp))

    data = {
        "halaman": data_page,
        "kalimat": data_word,
        "posisi": data_geom
    }

    df = pd.DataFrame(data)

    y1_coor = []
    y2_coor = []
    for pos in df["posisi"]:
        y1_coor.append(pos[0][1])
        y2_coor.append(pos[1][1])

    df["y1_coor"] = y1_coor
    df["y2_coor"] = y2_coor

    df = df.sort_values(by=["halaman", "y2_coor"])

    temp_kalimat = []
    merge_kalimat = []
    merge_y1 = []
    merge_y2 = []
    merge_halaman = []

    for i in range(0, df.shape[0]):
        if i == 0:
            temp_y2 = df["y2_coor"].iloc[i]
            temp_y1 = df["y1_coor"].iloc[i]
            temp_hal = df["halaman"].iloc[i]
            temp_kalimat.append(df["kalimat"].iloc[i])
        else:
            if abs(temp_y2 - df["y1_coor"].iloc[i]) < diff_line:
                temp_kalimat.append(df["kalimat"].iloc[i])
                temp_y2 = df["y2_coor"].iloc[i]
            else:
                merge_kalimat.append(" ".join(temp_kalimat))
                merge_y1.append(temp_y1)
                merge_y2.append(temp_y2)
                merge_halaman.append(temp_hal)

                temp_y2 = df["y2_coor"].iloc[i]
                temp_y1 = df["y1_coor"].iloc[i]
                temp_hal = df["halaman"].iloc[i]

                temp_kalimat = []
                temp_kalimat.append(df["kalimat"].iloc[i])

    data = {
        "halaman": merge_halaman,
        "kalimat": merge_kalimat,
        "y1_coor": merge_y1,
        "y2_coor": merge_y2
    }

    return pd.DataFrame(data)

def highlight_pdf(doc, delete_df):
    first_page = doc[0]
    rect = first_page.rect
    width = rect.width
    height = rect.height
    
    for i in range(delete_df.shape[0]):
        norm_y1 = delete_df["y1_coor"].iloc[i]
        norm_y2 = delete_df["y2_coor"].iloc[i]
        page = int(delete_df["halaman"].iloc[i])

        doc[page].draw_rect(
            fitz.Rect(1, norm_y1, doc[page].rect.width, norm_y2),
            color=(0, 1, 0),
            width=2
        )


def get_download_link(file_path, text):
    """Generate a download link for the given file."""
    absolute_path = os.path.abspath(file_path)
    href = f'<a href="{absolute_path}" download="{absolute_path}">{text}</a>'
    return href

# Main Streamlit app
def main():
    st.title("PDF Paragraph Highlighter")

    # Upload PDF file
    uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])

    # Sidebar for threshold and line_diff values
    threshold = st.sidebar.number_input("Threshold", min_value=0.1, max_value=1.0, value=0.4, step=0.1)
    line_diff = st.sidebar.number_input("Line Diff Value", min_value=1.0, max_value=100.0, value=50.0, step=5.0)
    line_diff = line_diff / 10000

    if uploaded_file is not None:
        st.subheader("Uploaded PDF Preview")
        st.write(f"Filename: {uploaded_file.name}")

        # Processing PDF content
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            temp_file.write(uploaded_file.read())
            temp_file_path = temp_file.name

        doc = DocumentFile.from_pdf(temp_file_path)
        predictor = ocr_predictor(pretrained=True)
        result = predictor(doc)
        json_export = result.export()
        res_df = process_pdf(json_export, line_diff)

        # Search for similar paragraphs
        search_string = "deletion erasure removal purge expunge wipe clear annulment obliteration destruction"
        embeddings_distilbert = model.encode(res_df['kalimat'].values)
        search_vect = model.encode([search_string])
        distilbert_similar_indexes = find_similar(search_vect, embeddings_distilbert, threshold)
        delete_df = res_df.iloc[distilbert_similar_indexes, :]

        # Display processed DataFrame
        st.subheader("Processed DataFrame")
        st.dataframe(delete_df)
        
        # Highlight paragraphs in the PDF
        doc = fitz.open(temp_file)
        highlight_pdf(doc, delete_df)

        # Save the highlighted PDF to a temporary file
        highlighted_temp_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name
        doc.save(highlighted_temp_path)
        
        # Display download link for the highlighted PDF
        st.subheader("Download Highlighted PDF")
        st.markdown(get_download_link(highlighted_temp_path, 'Download Highlighted PDF'), unsafe_allow_html=True)

        # Cleanup: Remove temporary file
        os.remove(temp_file_path)

if __name__ == "__main__":
    # Load the SentenceTransformer model
    model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens')

    # Run the Streamlit app
    main()