import json
import logging.config

from flask import Flask, request
from paddlenlp import Taskflow

from __init__ import IE_KEY
from log import LOGGING_CONFIG
from utils import process_request

app = Flask(__name__)
SETTLEMENT_LIST_SCHEMA = tuple(IE_KEY[key] for key in [
    'name', 'admission_date', 'discharge_date', 'medical_expenses', 'personal_cash_payment',
    'personal_account_payment', 'personal_funded_amount', 'medical_insurance_type', 'admission_id', 'settlement_id',
    'uppercase_medical_expenses'
])
SETTLEMENT_IE = Taskflow('information_extraction', schema=SETTLEMENT_LIST_SCHEMA, model='uie-x-base',
                         task_path='model/settlement_list_model', layout_analysis=False, precision='fp16')


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


@app.route('/text', methods=['POST'])
@process_request
def text():
    t = request.form.get('text')
    return SETTLEMENT_IE(t)


if __name__ == '__main__':
    logging.config.dictConfig(LOGGING_CONFIG)
    app.run('0.0.0.0', 5002)