"""Test all NeuronLens frontend functions with model cleanup."""
import os
import sys
from pathlib import Path

# Import neuronlens package (should be installed via pip install -e .)
import neuronlens

env_path = Path(__file__).parent.parent / ".env"
if env_path.exists():
    try:
        from dotenv import load_dotenv
        load_dotenv(env_path, override=True)
    except ImportError:
        pass

engine = neuronlens.Engine(api_key=os.getenv("API_KEY") or os.getenv("NEURONLENS_API_KEY"))

print("=" * 60)
print("Testing All NeuronLens Frontend Functions")
print("=" * 60)

# 1. Sentiment (SLM)
print("\n[1/7] Testing analyze_sentiment...")
result = engine.analyze_sentiment("Strong quarterly earnings", top_n=5)
print(f"  Result: {result.get('sentiment', {}).get('predicted_label', 'N/A')}")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 2. Agent
print("\n[2/7] Testing analyze_agent...")
result = engine.analyze_agent("Get Tesla stock price")
print(f"  Result: {result.get('alignment_status', 'N/A')}")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 3. Hallucination
print("\n[3/7] Testing analyze_hallucination...")
result = engine.analyze_hallucination("Write a summary about Apple", max_tokens=50)
print(f"  Result: {len(result.get('token_details', []))} tokens analyzed")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 4. Trading
print("\n[4/7] Testing analyze_trading...")
result = engine.analyze_trading("Apple reported strong earnings", "AAPL")
print(f"  Result: {len(result.get('signals', []))} signals")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 5. Reasoning
print("\n[5/7] Testing analyze_reasoning...")
result = engine.analyze_reasoning("What is 2+2? Show reasoning.")
print(f"  Result: {len(result.get('features', []))} features")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 6. Search & Steer
print("\n[6/7] Testing search_features...")
result = engine.search_features("financial analysis", k=5, layer=16)
print(f"  Result: {len(result)} features found")
engine.unload_all_models()
print("  ✓ Model unloaded")

# 7. Extract Top Features (Common)
print("\n[7/7] Testing extract_top_features...")
sae_base = os.getenv("SAE_MODELS_BASE", "")
sae_path = f"{sae_base}/finbert/ae.pt" if sae_base else "/path/to/sae"
result = engine.extract_top_features("Strong earnings", "ProsusAI/finbert", sae_path, 10, top_n=5)
print(f"  Result: {len(result.get('top_features', []))} features")
engine.unload_all_models()
print("  ✓ Model unloaded")

print("\n" + "=" * 60)
print("✓ All 7 tests completed")
print("=" * 60)
