"""Tests for flow dampening."""

import numpy as np
import pytest
from src.primitives.flow_dampening import (
    FlowTracker, update_flow, flow_at_time,
    time_to_decay_below_threshold,
    flow_penalty_multiplier,
    simulate_bank_run,
)


class TestFlowTracker:
    def test_initial_flow_zero(self):
        tracker = FlowTracker()
        assert tracker.current_flow == 0.0
        assert not tracker.is_above_threshold

    def test_update_adds_flow(self):
        tracker = FlowTracker(total_value_ref=1000.0)
        tracker = update_flow(tracker, 50.0, current_time=0)
        assert tracker.current_flow == 50.0

    def test_flow_decays_over_time(self):
        tracker = FlowTracker(memory=0.99, total_value_ref=1000.0)
        tracker = update_flow(tracker, 100.0, current_time=0)
        flow_later = flow_at_time(tracker, 100.0)
        assert flow_later < 100.0
        assert flow_later > 0

    def test_threshold_detection(self):
        tracker = FlowTracker(threshold=0.1, total_value_ref=1000.0)
        tracker = update_flow(tracker, 150.0, current_time=0)
        assert tracker.is_above_threshold  # 150/1000 = 0.15 > 0.1


class TestDecay:
    def test_decay_time_calculation(self):
        tracker = FlowTracker(
            memory=0.99, threshold=0.05,
            current_flow=100.0, total_value_ref=1000.0,
        )
        t = time_to_decay_below_threshold(tracker)
        assert t is not None
        assert t > 0
        # Verify: at time t, flow should be at threshold
        expected_flow = 100.0 * 0.99 ** t
        target = 0.05 * 1000.0
        assert abs(expected_flow - target) < 1.0

    def test_already_below_threshold(self):
        tracker = FlowTracker(
            threshold=0.1, current_flow=10.0, total_value_ref=1000.0,
        )
        assert time_to_decay_below_threshold(tracker) == 0.0


class TestPenalty:
    def test_no_penalty_below_half_threshold(self):
        tracker = FlowTracker(threshold=0.1, total_value_ref=1000.0)
        tracker = update_flow(tracker, 30.0, current_time=0)
        # 30/1000 = 0.03 < 0.05 (half of 0.1)
        assert flow_penalty_multiplier(tracker) == 1.0

    def test_penalty_at_threshold(self):
        tracker = FlowTracker(threshold=0.1, total_value_ref=1000.0)
        tracker = update_flow(tracker, 100.0, current_time=0)
        # 100/1000 = 0.1 = threshold
        penalty = flow_penalty_multiplier(tracker)
        assert penalty < 1.0
        assert penalty > 0.0

    def test_penalty_floor(self):
        tracker = FlowTracker(threshold=0.1, total_value_ref=1000.0)
        tracker = update_flow(tracker, 500.0, current_time=0)
        # Way above threshold
        assert flow_penalty_multiplier(tracker) >= 0.1


class TestBankRunSimulation:
    def test_simulation_runs(self):
        result = simulate_bank_run(
            initial_value=10000, supply=10000,
            memory=0.99, threshold=0.1,
            redemption_rate=0.05, n_steps=50,
        )
        assert len(result["times"]) == 50
        assert len(result["values"]) == 50
        assert all(result["values"] >= 0)

    def test_dampening_preserves_value(self):
        """With dampening, more value should remain vs without."""
        # With dampening (low threshold = aggressive)
        result_damped = simulate_bank_run(
            initial_value=10000, supply=10000,
            memory=0.99, threshold=0.05,
            redemption_rate=0.1, n_steps=50,
        )
        # Without dampening (high threshold = never triggers)
        result_free = simulate_bank_run(
            initial_value=10000, supply=10000,
            memory=0.99, threshold=10.0,
            redemption_rate=0.1, n_steps=50,
        )
        # Dampened system should preserve more value
        assert result_damped["values"][-1] > result_free["values"][-1]