import platform
import pickle
import os

# IMPORTANT: import the module that defines the custom wrapper classes BEFORE
# loading the pickle. Without this, unpickling fails in the API container
# because BertweetVectorizer / BertweetClassifier can't be resolved.
import sentiment_deploy  # noqa: F401

from flask import Flask, jsonify, request
from flask_cors import CORS

# ////////////////////////////////////////////////////////////////////////
GROUP_ID = 'modelling-giants'                       # TODO: your groupID
MODEL_FILE = 'route_c_bertweet_large_fp16.model'    # the shrunk file
MODEL_VERSION = 'v1.0-large-fp16'


def batch_predict(model, items):
    # Predict the whole batch in ONE call so the model isn't rebuilt per item.
    texts = [item['text'] for item in items]
    X = model['vectorizer'].transform(texts)        # transform, NOT fit_transform
    labels = model['classifier'].predict(X)
    return [
        {"id": item['id'], "label": int(label)}
        for item, label in zip(items, labels)
    ]


# ////////////////////////////////////////////////////////////////////////
# Do not modify below.
# ////////////////////////////////////////////////////////////////////////

app = Flask(__name__)
CORS(app)

with open(MODEL_FILE, 'rb') as file:
    model = pickle.load(file)

meta_data = {
    "groupID": GROUP_ID,
    "modelFile": MODEL_FILE,
    "modelVersion": MODEL_VERSION,
    "pythonVersion": platform.python_version()
}


@app.route("/", methods=['GET', 'POST'])
def main():
    if request.method == 'POST':
        items = request.json['items']
        return jsonify({"items": batch_predict(model, items)})
    else:
        return jsonify({"meta": meta_data})


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8000))
    app.run(host="0.0.0.0", port=port, debug=True)