import sys
from unittest.mock import MagicMock

# Mock tensorflow BEFORE importing loka to avoid ImportError in environment without tf
mock_tf = MagicMock()
sys.modules["tensorflow"] = mock_tf
sys.modules["tensorflow.keras"] = mock_tf.keras
sys.modules["tensorflow.keras.models"] = mock_tf.keras.models

import numpy as np
import pytest
from unittest.mock import patch
from grahana import predict

@pytest.fixture
def mock_model():
    """Mock the Keras model to avoid needing the actual .h5 file and tensorflow."""
    with patch("grahana.model_loader.load_model") as mock_load:
        model_instance = MagicMock()
        # Mock prediction: return [0.1, 0.9] -> candidate
        model_instance.predict.return_value = np.array([[0.1, 0.9]])
        mock_load.return_value = model_instance
        yield model_instance

@pytest.fixture
def mock_path_exists():
    """Mock Path.exists to always return True for the model file."""
    with patch("pathlib.Path.exists") as mock_exists:
        mock_exists.return_value = True
        yield mock_exists

def test_predict_label(mock_model, mock_path_exists):
    # Reset singleton if it was already loaded (hacky but needed for tests if running in same process)
    import grahana.model_loader
    grahana.model_loader._model = None
    
    label = predict(3.0, 1000.0, 2.0, 10.0, 1.0, return_label=True)
    assert isinstance(label, str)
    # With probs [0.1, 0.9], argmax=1 -> encoder classes ['candidate', 'false_positive'] -> 'false_positive'
    assert label == "false_positive"

def test_predict_probs(mock_model, mock_path_exists):
    import grahana.model_loader
    grahana.model_loader._model = None
    
    probs = predict(3.0, 1000.0, 2.0, 10.0, 1.0, return_label=False)
    assert isinstance(probs, np.ndarray)
    assert probs.shape == (2,)
    assert np.allclose(probs, [0.1, 0.9])
