# -*- coding: utf-8 -*-
"""
@author: lywen
"""
import os
import json
import time
import web
import numpy as np
import uuid
from PIL import Image
from model import text_predict, crnn_handle

web.config.debug = True

filelock = "file.lock"
if os.path.exists(filelock):
    os.remove(filelock)

render = web.template.render("templates", base="base")

from config import *
from apphelper.image import union_rbox, adjust_box_to_origin, base64_to_PIL
from application import trainTicket, idcard


# from main import TextOcrModel


billList = ["通用OCR", "火车票", "身份证"]


class OCR:
    """通用OCR识别"""

    def GET(self):
        post = {}
        post["postName"] = "ocr"  ##请求地址
        post["height"] = 1000
        post["H"] = 1000
        post["width"] = 600
        post["W"] = 600
        post["billList"] = billList
        return render.ocr(post)

    def POST(self):
        t = time.time()
        data = web.data()
        uidJob = uuid.uuid1().__str__()

        data = json.loads(data)
        # print(data)
        billModel = data.get("billModel", "")
        # textAngle = data.get('textAngle',False)##文字检测
        textLine = data.get("textLine", False)  ##只进行单行识别

        imgString = data["imgString"].encode().split(b";base64,")[-1]
        img = base64_to_PIL(imgString)
        if img is not None:
            img = np.array(img)

        H, W = img.shape[:2]

        while time.time() - t <= TIMEOUT:
            if os.path.exists(filelock):
                continue
            else:
                with open(filelock, "w") as f:
                    f.write(uidJob)

                if textLine:
                    ##单行识别
                    partImg = Image.fromarray(img)
                    text = crnn_handle.predict(partImg)
                    res = [{"text": text, "name": "0", "box": [0, 0, W, 0, W, H, 0, H]}]
                    os.remove(filelock)
                    break

                else:
                    # detectAngle = textAngle
                    result = text_predict(img)

                    if billModel == "" or billModel == "通用OCR":
                        # result = union_rbox(result,0.2)
                        res = [
                            {
                                "text": x["text"],
                                "name": str(i),
                                "box": {
                                    "cx": x["cx"],
                                    "cy": x["cy"],
                                    "w": x["w"],
                                    "h": x["h"],
                                    "angle": x["degree"],
                                },
                            }
                            for i, x in enumerate(result)
                        ]
                        # res = adjust_box_to_origin(img,angle, res)##修正box

                    elif billModel == "火车票":
                        res = trainTicket.trainTicket(result)
                        res = res.res
                        res = [
                            {"text": res[key], "name": key, "box": {}} for key in res
                        ]

                    elif billModel == "身份证":

                        res = idcard.idcard(result)
                        res = res.res
                        res = [
                            {"text": res[key], "name": key, "box": {}} for key in res
                        ]

                    os.remove(filelock)
                    break

        timeTake = time.time() - t

        return json.dumps(
            {"res": res, "timeTake": round(timeTake, 4)}, ensure_ascii=False
        )


urls = ("/ocr", "OCR")

if __name__ == "__main__":

    app = web.application(urls, globals())
    app.run()
