from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
import base64
import io
import os
from PIL import Image, ImageOps
import numpy as np
import cv2
from scipy.ndimage import interpolation as inter
import pytesseract



def correct_skew(image, delta=1, limit=5):
    def determine_score(arr, angle):
        data = inter.rotate(arr, angle, reshape=False, order=0)
        histogram = np.sum(data, axis=1, dtype=float)
        score = np.sum((histogram[1:] - histogram[:-1]) ** 2, dtype=float)
        return histogram, score

    if len(image.shape)==3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] 

    scores = []
    angles = np.arange(-limit, limit + delta, delta)
    for angle in angles:
        histogram, score = determine_score(thresh, angle)
        scores.append(score)

    best_angle = angles[scores.index(max(scores))]

    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, best_angle, 1.0)
    corrected = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, \
            borderMode=cv2.BORDER_REPLICATE)

    return best_angle, corrected

def getInputData(request_data):
    if not request.json or 'fileType' not in request.json:
         return 'Error: No fileType provided', 400
    file_type = request_data.get('fileType')

    if not request.json or 'file' not in request.json:
         return 'Error: No file provided', 400
    file_data = request_data.get('file')

    if not request.json or 'ocrEngine' not in request.json:
         return 'Error: No ocrEngine provided', 400
    ocr_engine = request_data.get('ocrEngine')

    if not request.json or 'language' not in request.json:
         return 'Error: No language provided', 400
    language = request_data.get('language')

    if not request.json or 'scale' not in request.json:
         return 'Error: No scale provided', 400
    scale = request_data.get('scale')

    if not request.json or 'detectOrientation' not in request.json:
         return 'Error: No detectOrientation provided', 400
    detect_orientation = request_data.get('detectOrientation')

    if not request.json or 'region' not in request.json:
         return 'Error: No region provided', 400
    region = request_data.get('region')

    extracted_data = {
        'fileType': file_type,
        'file_data': file_data,
        'ocrEngine': ocr_engine,
        'language': language,
        'scale': scale,
        'detectOrientation': detect_orientation,
        'region': region
    }

    return extracted_data


app = Flask(__name__)


UPLOAD_FOLDER = 'files/upload/'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp', 'webp', 'pdf', 'gif', 'tif'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER


def allowed_file(ext):
    if ext in ALLOWED_EXTENSIONS:
        return True
    else:
        return False


def makeTransform(pdfFile, image, extracted_data):
    #preprocess to improve quality of results
    #Add border
    img_with_border = ImageOps.expand(image, border=50, fill='red')
    image_np = np.array(img_with_border)

    clahe = cv2.createCLAHE(clipLimit=1)
    if(extracted_data['detectOrientation']):
        angle, image_np = correct_skew(image_np)

    #Scanned documents are ok, CLAHE is not needed
    if(pdfFile==False):
        if len(image_np.shape)==3:
            image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
            image_np = clahe.apply(image_np) + 1
        else:
            image_np = clahe.apply(image_np) + 1
    else:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
         

    custom_config = r'-l ces --oem 1 --psm 6' 
    
    text_opraveno = pytesseract.image_to_string(image_np,config=custom_config)

    return text_opraveno


# Route to receive data and return text
@app.route('/process_data', methods=['POST'])
def process_data():
    # Check if request contains JSON data
    request_data = request.json

    extracted_data = getInputData(request_data)

    encoded_image = extracted_data['file_data']

    if allowed_file(extracted_data['fileType']):
        filename = 'img.'+extracted_data['fileType']
        decoded_image = base64.b64decode(encoded_image)

        imgdata = base64.b64decode(encoded_image)
        img = Image.open(io.BytesIO(imgdata))

        #img.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))

        # checking if it is a file
        pdfFile = (extracted_data['fileType']=='pdf')


        image = img    

        
        if(extracted_data['region']['makeCut']):
            size = extracted_data['region']['coordinates'][0]
            image = image.crop((size['x'], size['y'], size['x'] + size['width'], size['y'] + size['height']))


        text = makeTransform(pdfFile, image, extracted_data)        
    else:
        return 'Error: No supported file format', 400





    parsed_text = {
        "html": "<p>This is converted text from the input image.</p>",
        "plainText": text
    }

    # Return the parsed text as JSON
    return jsonify({"parsedText": parsed_text})


if __name__ == '__main__':
    app.run(debug=True)