# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['load_pdf', 'pdf_img_to_np', 'cid_to_char', 'combine_lines', 'get_avg_gap', 'get_max_gap', 'combine_splits',
           'split_words', 'process_text', 'only_bnw', 'color_fill_pdf_text']

# %% ../nbs/00_core.ipynb 3
from .imports import *

# %% ../nbs/00_core.ipynb 4
def load_pdf(pdf):
    if path_or_str(pdf):
        pdf = pdfplumber.open(pdf)
    elif not isinstance(pdf, pdfplumber.pdf.PDF):
        raise Exception(f'pdf must be a path or pdfplumber.pdf.PDF object.')
    return pdf

def pdf_img_to_np(img):
    return np.array(img.annotated)

def cid_to_char(cidx):
    try:
        return chr(int(re.findall(r'\(cid\:(\d+)\)',cidx)[0]) + 29)
    except:
        return cidx

def combine_lines(txt):
    avg_len = np.mean([len(w.split()) for w in txt])
    txt2 = []
    for i,w in enumerate(txt):
        w = process_text(w)
        if i==0:
            txt2.append(w)
        else:
            if len(txt2[-1].split()) < avg_len:
                txt2[-1]+=' '+w
            else:
                txt2.append(w)
    return txt2

def get_avg_gap(words, key0='top', key1='bottom'):
    if key1 is None: key1 = key0
    return np.mean([w[key0]-words[i-1][key1] for i,w in enumerate(words) if i>0])

def get_max_gap(words, key0='top', key1='bottom'):
    if key1 is None: key1 = key0
    return np.max([w[key0]-words[i-1][key1] for i,w in enumerate(words) if i>0])

def combine_splits(splits):
    avg_len = np.mean([len(s) for s in splits])
    splits2 = []
    for i,s in enumerate(splits):
        if i==0:
            splits2.append(s)
        else:
            if len(splits2[-1]) < avg_len:
                splits2[-1]+=s
            else:
                splits2.append(s)
    return splits2

def split_words(words, key0='top', key1='bottom', avg_gap=None, fill_empty=False):
    if key1 is None: key1 = key0
    if avg_gap is None:
        avg_gap = np.mean([w[key0]-words[i-1][key1] for i,w in enumerate(words) if i>0])
    splits = []
    for i,w in enumerate(words):
        if i==0:
            splits.append([w])
        else:
            if w[key0]-words[i-1][key1] > avg_gap:
                if fill_empty:
                    splits.append(['*']*len(splits[-1]) + [w])
                else:
                    splits.append([w])
            else:
                splits[-1].append(w)
    return splits

def process_text(text):
    text = cid_to_char(text)
    text = re.sub(r"\uf0b7", " ", text)
    text = re.sub(r"\(cid:\d{0,3}\)", " ", text)
    text = re.sub(r'• ', " ", text)
    text = re.sub(r'● ', " ", text)
    return text

def only_bnw(img, thresh=50):
    black = np.all(img < thresh, axis=-1)
    img[~black] = 255
    return img

def color_fill_pdf_text(pdf, height=512, color='black'):
    """
    Gets a pdf and returns a list of images with the text filled in with the specified color.
    
    Parameters
    ----------
    pdf: str or Path or pdfplumber.pdf.PDF
        The path to the pdf or a pdfplumber.pdf.PDF object.
    height: int
        The height of the image to be returned.
    color: str or tuple
        The color to fill the text with.
        
    Returns
    -------
    list
        A list of np.ndarray images.
    """
    
    pdf = load_pdf(pdf)
    color = tuple(color_to_rgb(color))
    return [only_bnw(pdf_img_to_np(page.to_image(height=height).draw_rects(page.extract_words(x_tolerance=5),
                                                                  fill=color, stroke=color))) for page in pdf.pages]
