"""
Tests for GRU model input/output shapes.

Ref: revisions.txt Task 3.6
"""
import pytest
import numpy as np
import sys, os

# --- Add path for src imports --- #
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
src_path = os.path.join(project_root, 'src')
if src_path not in sys.path:
    sys.path.insert(0, src_path)
# --- End Add path --- #

# Import the v3 model builder
from model_gru_v3 import build_gru_model_v3
# TODO: Import v2 model builder if needed for comparison tests
# from model_gru import build_gru_model

# --- Constants for Testing --- #
LOOKBACK = 60
N_FEATURES = 25
BATCH_SIZE = 4

# --- Tests --- #

def test_gru_v3_output_shapes():
    """Verify the output shapes of the GRU v3 model heads."""
    print(f"\nBuilding GRU v3 model for shape test...")
    # Build the v3 model with default parameters
    model = build_gru_model_v3(lookback=LOOKBACK, n_features=N_FEATURES)
    assert model is not None, "Failed to build GRU v3 model"
    
    # Check number of outputs
    assert len(model.outputs) == 2, f"Expected 2 outputs, got {len(model.outputs)}"
    
    # Check output names and shapes
    # Output order in the model definition was [mu, dir3]
    mu_output_shape = model.outputs[0].shape.as_list()
    dir3_output_shape = model.outputs[1].shape.as_list()
    
    # Assert shapes (ignoring batch size None)
    # mu head should be (None, 1)
    assert mu_output_shape == [None, 1], f"Expected mu shape [None, 1], got {mu_output_shape}"
    # dir3 head should be (None, 3)
    assert dir3_output_shape == [None, 3], f"Expected dir3 shape [None, 3], got {dir3_output_shape}"

    print("GRU v3 output shapes test passed.")

def test_gru_v3_prediction_shapes():
    """Verify the prediction shapes match the output shapes for a sample batch."""
    model = build_gru_model_v3(lookback=LOOKBACK, n_features=N_FEATURES)
    assert model is not None, "Failed to build GRU v3 model"

    # Create dummy input data
    dummy_input = np.random.rand(BATCH_SIZE, LOOKBACK, N_FEATURES)

    # Generate predictions
    predictions = model.predict(dummy_input)

    # Check prediction structure and shapes
    assert isinstance(predictions, list), "Predictions should be a list for multi-output model"
    assert len(predictions) == 2, f"Expected 2 prediction arrays, got {len(predictions)}"

    # Predictions order should match model.outputs order [mu, dir3]
    mu_preds = predictions[0]
    dir3_preds = predictions[1]

    # Assert prediction shapes match expected batch size
    assert mu_preds.shape == (BATCH_SIZE, 1), f"Expected mu prediction shape ({BATCH_SIZE}, 1), got {mu_preds.shape}"
    assert dir3_preds.shape == (BATCH_SIZE, 3), f"Expected dir3 prediction shape ({BATCH_SIZE}, 3), got {dir3_preds.shape}"
    
    print("GRU v3 prediction shapes test passed.")

# TODO: Add tests for GRU v2 model shapes if it's still relevant.

def test_logits_view_shapes():
    """Test that softmax applied to predict_logits output matches predict output."""
    print(f"\nBuilding GRU v3 model for logits view test...")
    model = build_gru_model_v3(lookback=LOOKBACK, n_features=N_FEATURES)
    assert model is not None, "Failed to build GRU v3 model"

    # --- Requires GRUModelHandler to run predict_logits --- #
    # We need to instantiate the handler to test its methods.
    # Mock config and directories needed for handler init.
    mock_config = {
        'control': {'use_v3': True},
        'gru_v3': {} # Use defaults for building
    }
    mock_run_id = "test_logits_run"
    mock_models_dir = "./mock_models/test_logits_run"
    os.makedirs(mock_models_dir, exist_ok=True) # Create mock dir
    
    # Import handler locally for test setup
    from gru_model_handler import GRUModelHandler
    handler = GRUModelHandler(run_id=mock_run_id, models_dir=mock_models_dir, config=mock_config)
    handler.model = model # Assign the already built model to the handler
    handler.model_version_used = 'v3' # Set version manually
    # --- End Handler Setup --- #

    # Create dummy input data
    dummy_input = np.random.rand(BATCH_SIZE, LOOKBACK, N_FEATURES).astype(np.float32)

    # Generate predictions using both methods
    logits = handler.predict_logits(dummy_input)
    predictions = handler.predict(dummy_input)

    assert logits is not None, "predict_logits returned None"
    assert predictions is not None, "predict returned None"
    assert isinstance(predictions, list) and len(predictions) == 2, "predict output structure incorrect"
    
    probs_from_predict = predictions[1] # dir3 is the second output

    # Apply softmax to logits
    # Use tf.nn.softmax for consistency with Keras backend
    import tensorflow as tf
    probs_from_logits = tf.nn.softmax(logits).numpy()

    # Assert shapes match first
    assert probs_from_logits.shape == probs_from_predict.shape, \
        f"Shape mismatch: softmax(logits)={probs_from_logits.shape}, predict_probs={probs_from_predict.shape}"

    # Assert values are close
    np.testing.assert_allclose(
        probs_from_logits, 
        probs_from_predict, 
        rtol=1e-6, 
        atol=1e-6, # Use tighter tolerance for numerical precision check
        err_msg="Softmax applied to logits does not match probability output from model.predict()"
    )

    print("Logits view test passed.")
    # Clean up mock directory
    import shutil
    if os.path.exists("./mock_models"):
         shutil.rmtree("./mock_models") 