import json

from flask import Flask, request
from paddlenlp import Taskflow

from utils import process_request
from . import HOSPITAL, DEPARTMENT, PATIENT_NAME, ADMISSION_DATE, DISCHARGE_DATE, DOCTOR, ADMISSION_ID, AGE

app = Flask(__name__)
DISCHARGE_RECORD_SCHEMA = (
        HOSPITAL + DEPARTMENT + PATIENT_NAME + ADMISSION_DATE + DISCHARGE_DATE + DOCTOR + ADMISSION_ID + AGE
)
DISCHARGE = Taskflow('information_extraction', schema=DISCHARGE_RECORD_SCHEMA, model='uie-x-base',
                     task_path='../../model/discharge_record_model', layout_analysis=False, precision='fp16')


@app.route('/nlp/discharge', methods=['POST'])
@process_request
def discharge():
    img_path = request.form.get('img_path')
    layout = request.form.get('layout')
    return DISCHARGE({'doc': img_path, 'layout': json.loads(layout)})


if __name__ == '__main__':
    app.run('0.0.0.0', 5003)