import torch
import numpy as np
from cognitive_mapping_probe.llm_iface import get_or_load_model
from cognitive_mapping_probe.resonance_seismograph import run_silent_cogitation_seismic
from cognitive_mapping_probe.concepts import get_concept_vector, _get_last_token_hidden_state
from cognitive_mapping_probe.signal_analysis import analyze_cognitive_signal

def test_get_or_load_model_loads_correctly(model_id):
    """Testet, ob das Laden eines echten Modells funktioniert."""
    llm = get_or_load_model(model_id, seed=42)
    assert llm is not None
    assert llm.model_id == model_id
    assert llm.stable_config.hidden_dim > 0
    assert llm.stable_config.num_layers > 0

def test_run_silent_cogitation_seismic_output_shape_and_type(model_id):
    """Führt einen kurzen Lauf mit einem echten Modell durch und prüft die Datentypen."""
    num_steps = 10
    llm = get_or_load_model(model_id, seed=42)
    state_deltas = run_silent_cogitation_seismic(
        llm=llm, prompt_type="control_long_prose",
        num_steps=num_steps, temperature=0.1
    )
    assert isinstance(state_deltas, list)
    assert len(state_deltas) == num_steps
    assert all(isinstance(d, float) for d in state_deltas)

def test_get_last_token_hidden_state_robustness(model_id):
    """Testet die Helper-Funktion mit einem echten Modell."""
    llm = get_or_load_model(model_id, seed=42)
    hs = _get_last_token_hidden_state(llm, "test prompt")
    assert isinstance(hs, torch.Tensor)
    assert hs.shape == (llm.stable_config.hidden_dim,)

def test_get_concept_vector_logic(model_id):
    """Testet die Vektor-Extraktion mit einem echten Modell."""
    llm = get_or_load_model(model_id, seed=42)
    vector = get_concept_vector(llm, "love", baseline_words=["thing", "place"])
    assert isinstance(vector, torch.Tensor)
    assert vector.shape == (llm.stable_config.hidden_dim,)

def test_analyze_cognitive_signal_no_peaks():
    """
    Testet den Edge Case, dass ein Signal keine signifikanten Frequenz-Peaks hat.
    """
    flat_signal = np.linspace(0, 1, 100)
    results = analyze_cognitive_signal(flat_signal)
    assert results is not None
    assert results["dominant_periods_steps"] is None
    assert "spectral_entropy" in results

def test_analyze_cognitive_signal_with_peaks():
    """
    Testet den Normalfall, dass ein Signal Peaks hat, mit realistischerem Rauschen.
    """
    np.random.seed(42)
    steps = np.arange(200)
    # Signal mit einer starken Periode von 10 und einer schwächeren von 25
    signal_with_peak = (1.0 * np.sin(2 * np.pi * (1/10.0) * steps) + 
                          0.5 * np.sin(2 * np.pi * (1/25.0) * steps) +
                          np.random.randn(200) * 0.5) # Realistischeres Rauschen
    results = analyze_cognitive_signal(signal_with_peak)

    assert results["dominant_periods_steps"] is not None
    assert 10.0 in results["dominant_periods_steps"]
    assert 25.0 in results["dominant_periods_steps"]

def test_analyze_cognitive_signal_with_multiple_peaks():
    """
    Erweiterter Test, der die korrekte Identifizierung und Sortierung
    von drei Peaks verifiziert, mit realistischerem Rauschen.
    """
    np.random.seed(42)
    steps = np.arange(300)
    # Definiere drei Peaks mit unterschiedlicher Stärke (Amplitude)
    signal = (2.0 * np.sin(2 * np.pi * (1/10.0) * steps) + 
              1.5 * np.sin(2 * np.pi * (1/4.0)  * steps) +
              1.0 * np.sin(2 * np.pi * (1/30.0) * steps) +
              np.random.randn(300) * 0.5) # Realistischeres Rauschen
              
    results = analyze_cognitive_signal(signal, num_peaks=3)

    assert results["dominant_periods_steps"] is not None
    expected_periods = [10.0, 4.0, 30.0]
    assert results["dominant_periods_steps"] == expected_periods