"""
CO-PO Mapping Model - Usage Example
"""

import pickle
import numpy as np
from huggingface_hub import hf_hub_download

def load_model():
    """Download and load the CO-PO mapping model"""
    model_path = hf_hub_download(
        repo_id="Jrine/co-po",
        filename="co_po_model_complete.pkl"
    )
    
    with open(model_path, 'rb') as f:
        package = pickle.load(f)
    
    return package['model'], package['vectorizer']

def predict_co_po(statement, model, vectorizer):
    """Predict PO correlations for a CO statement"""
    # Vectorize
    vec = vectorizer.transform([statement])
    
    # Predict
    pred = model.predict(vec)[0]
    pred_rounded = np.clip(np.round(pred), 0, 3).astype(int)
    
    return pred_rounded

# Example usage
if __name__ == "__main__":
    # Load model
    print("Loading model...")
    model, vectorizer = load_model()
    print("Model loaded!\n")
    
    # Example CO
    co = "Design and implement machine learning solutions for real-world problems"
    
    # Predict
    predictions = predict_co_po(co, model, vectorizer)
    
    # Display
    po_names = ['Engineering Knowledge', 'Problem Analysis', 'Design/Development',
                'Investigation', 'Modern Tool Usage', 'Engineer and Society',
                'Environment', 'Ethics', 'Team Work', 'Communication', 
                'Project Management']
    
    levels = ['None', 'Low', 'Medium', 'High']
    
    print(f"Course Outcome: {co}\n")
    print("PO Mapping:")
    print("-" * 70)
    
    for i, (name, score) in enumerate(zip(po_names, predictions), 1):
        print(f"PO{i:2d} ({name:30s}): {score} ({levels[score]})")