import requests
import os
import bilbystats as bs
import numpy as np

def china_importance(texts, local = False):
    if local:
        model_path = bs.check_dir("marco_importance/orig_title_model/checkpoint-664")
        model_name = "hfl/chinese-roberta-wwm-ext" 
        model_dict = bs.load_model(model_path, model_name)
        predictions = bs.predict(texts, model_dict)
        labels = predictions['pred_labels']
        scores = predictions['logits']
        confidence = np.max(scores, axis = 1)
    else:
        response = requests.post(
            "https://sdavenport--china-importance-classifier-model-classify-batch.modal.run",
            json={"texts": texts},
            headers={
                "Modal-Key": os.getenv("MODAL_KEY"),
                "Modal-Secret": os.getenv("MODAL_SECRET")
                },
            )
        result = response.json()
        labels = [result['results'][i]['label'] for i in range(len(result['results']))]
        dict2use = {'None': 0, 'High': 1, 'Medium': 2}
        labels = [dict2use[label] for label in labels]
        
        confidence = [result['results'][i]['confidence'] for i in range(len(result['results']))]
        scores = [result['results'][i]['scores'] for i in range(len(result['results']))]
        logits = [result['results'][i]['logits'] for i in range(len(result['results']))]

    return labels, confidence, scores, logits

def scores2matrix(data, columns=None):
    if columns is None:
        columns = list(data[0].keys())
    return np.array([[d[c] for c in columns] for d in data])

