from phenotype_inference import PhenotypeClassifier, print_prediction
from phenotype_next_tests import PhenotypeNextTestRecommender, print_next_test_report

MODEL_PATH = "phenotype_tinytransformer_v1_temperature_scaled.pt"
REFERENCE_PATH = "phenotype_reference_distributions.json"

classifier = PhenotypeClassifier(MODEL_PATH)
recommender = PhenotypeNextTestRecommender(classifier, REFERENCE_PATH)

features = {
    "Gram Stain": "Negative",
    "Shape": "Rods",
    "Catalase": "Positive",
    "Oxidase": "Positive",
    "Motility": "Positive",
    "Indole": "Negative",
    "Citrate": "Positive",
    "Urease": "Negative",
    "Growth Temperature": "20//37",
    "Media Grown On": "Blood Agar; MacConkey Agar",
}

prediction = classifier.predict(features, top_k=10)
print_prediction(prediction)

recommendations = recommender.recommend(
    features,
    n_recommendations=5,
    top_competing_genera=5,
)

print_next_test_report(recommendations)