import os
import subprocess
import numpy as np
import cv2 as cv
from lxml import etree
from PIL import Image
from src.deskew import deskew_image_pca, get_cropping_corner_points, crop_bb_from_img
from src.alto import create_original_alto, modify_cutout, copy_blocks_to_original, recompute_ids


def process_sample(sample, artifacts_dir='/home/jovyan/work/shared/esperanto-excerpts/artifacts'):
    # Create canonical ALTOXML template
    with Image.open(sample.filepath) as im:
        width, height = im.size
        root = create_original_alto(sample.filepath, width, height)
    
    directory = os.path.basename(os.path.dirname(sample.filepath))
    filename = os.path.splitext(os.path.basename(
            sample.filepath))[0]  # without ext
    
    for idx, detection in enumerate(sample['prediction']['detections']):
        if detection['label'] != 'Text':
            continue

        if os.path.isfile(f'{artifacts_dir}/predictions/{directory}/{filename}/{idx}-deskewed.jpg'):
            continue

        mask = detection['mask'].astype(bool)
        box = detection['bounding_box']
        img_path = sample['filepath']
        img = Image.open(img_path)

        #print('---Box percent---')
        #print(f'topleftx: {box[0]}, toplefty: {box[1]}, width: {box[2]}, height: {box[3]}')

        box_width = int(
            (box[0] + box[2]) * img.width) - int(box[0] * img.width)
        box_height = int(
            (box[1] + box[3]) * img.height) - int(box[1] * img.height)

        if mask.shape[1] > box_width:
            mask = mask[:, :-1]
        if mask.shape[0] > box_height:
            mask = mask[:-1, :]

        if mask.shape[1] < box_width:
            box_width -= 1
        if mask.shape[0] < box_height:
            box_height -= 1

        #print(f'Box width: {box_width}, Box height: {box_height}')

        coords = [
            int(box[0] * img.width),
            int(box[1] * img.height), box_width, box_height
        ]

        #print(f'Coords: {coords}')

        pad_left = coords[0]
        pad_top = coords[1]
        pad_right = img.width - pad_left - coords[2]
        pad_bottom = img.height - pad_top - coords[3]

        mask = np.pad(mask, [(pad_top, pad_bottom), (pad_left, pad_right)])

        mask3d = np.dstack([mask] * 3)
        cut_out_mask = np.invert(mask3d)
        res = np.where(cut_out_mask, 0, img)
        image = Image.fromarray(res)
        os.makedirs(f'{artifacts_dir}/predictions/{directory}/{filename}', exist_ok=True)

        img_path = f'{artifacts_dir}/predictions/{directory}/{filename}/{idx}.jpg'
        image.save(img_path)

        single_image = cv.imread(img_path)
        bbox = get_cropping_corner_points(single_image,
                                          threshold_border=0.001,
                                          debug=False)
        cropped_img = crop_bb_from_img(single_image, bbox, margin=10)
        deskewed_img, angle = deskew_image_pca(cropped_img, debug=False)
        
        detection.angle_deg = angle
        sample.save()
        
        #r1, g1, b1 = 10, 10, 10 # Original value
        #r2, g2, b2 = 255, 255, 255 # Value that we want to replace it with

        #red, green, blue = deskewed_img[:,:,0], deskewed_img[:,:,1], deskewed_img[:,:,2]
        #mask = (red <= r1) & (green <= g1) & (blue <= b1)
        #deskewed_img[:,:,:3][mask] = [r2, g2, b2]
        
        deskewed_path = f'{artifacts_dir}/predictions/{directory}/{filename}/{idx}-deskewed.jpg'

        # Add white 10 px border to counter Error in boxClipToRectangle: box outside rectangle
        img_with_border = cv.copyMakeBorder(deskewed_img,
                                            10,
                                            10,
                                            10,
                                            10,
                                            cv.BORDER_CONSTANT,
                                            value=[255, 255, 255])
        cv.imwrite(deskewed_path,
                   img_with_border,
                   params=(cv.IMWRITE_JPEG_QUALITY, 95,
                           cv.IMWRITE_JPEG_PROGRESSIVE, 1,
                           cv.IMWRITE_JPEG_OPTIMIZE, 1))

        ocr_path = f'{artifacts_dir}/predictions/{directory}/{filename}/{idx}'
        
        langs = 'fra+deu+Fraktur+frk+eng+spa+ita+epo+rus+nld+por+pol+ces+cat+swe+dan+hun+lat+bul+fin+isl+ron'

        command = f'tesseract -l {langs} --oem 1 --psm 1 {deskewed_path} {ocr_path}-deskewed.alto alto txt'
        result = subprocess.check_output(command,
                                         stderr=subprocess.STDOUT,
                                         shell=True,
                                         universal_newlines=True)
        if 'Too few characters' in result:
            # print('Too few characters detected. Retrying with psm = 7')
            # Retry with single-line mode of tesseract (psm = 7)
            command = f'tesseract -l {langs} --psm 7 {deskewed_path} {ocr_path}-deskewed.alto alto'
            result = subprocess.check_output(command,
                                             stderr=subprocess.STDOUT,
                                             shell=True,
                                             universal_newlines=True)
        
        cutout = modify_cutout(f'{ocr_path}-deskewed.alto.xml', coords, detection.angle_deg)
        root = copy_blocks_to_original(root, width, height, cutout)
    
    final_alto_path = f'{artifacts_dir}/predictions/{directory}/{filename}/{filename}.alto.xml'
    filename_path = f'{artifacts_dir}/predictions/{directory}/{filename}'
    if not os.path.exists(filename_path):
        os.makedirs(filename_path)
    if not os.path.exists(final_alto_path):
        # Recompute IDs once all cutouts are in the new root
        root = recompute_ids(root)
    
        # Save final ALTOXML
        xml_string = etree.tostring(root,
                                    pretty_print=True,
                                    xml_declaration=True,
                                    encoding='utf-8')
    
        with open(final_alto_path, 'wb') as f:
            f.write(xml_string)
        
        # Delete leftover individual alto files and .jpg files
        leftover = [f for f in os.listdir(filename_path) if 'deskewed' in f or f.endswith('.jpg')]
        for f in leftover:
            os.remove(os.path.join(filename_path, f))