import numpy as np
import pandas as pd
import pytest

from fastbt.src.fastbt.simulate_volatility import generate_volatility_scenarios


@pytest.fixture
def sample_series():
    # Simulate a geometric Brownian motion price series
    np.random.seed(42)
    n = 100
    init_price = 100.0
    mu = 0.0002
    sigma = 0.01
    dt = 1
    log_returns = np.random.normal(mu, sigma, n - 1)
    prices = [init_price]
    for ret in log_returns:
        prices.append(prices[-1] * np.exp(ret))
    idx = pd.date_range("2020-01-01", periods=n, freq="B")
    return pd.Series(prices, index=idx)


@pytest.fixture
def flat_series():
    # Flat price series (zero volatility, like a fixed-rate bond)
    idx = pd.date_range("2020-01-01", periods=50, freq="B")
    return pd.Series([100.0] * 50, index=idx)


def test_output_shape_and_columns(sample_series):
    df = generate_volatility_scenarios(sample_series, random_seed=123)
    assert df.shape[0] == len(sample_series)
    assert "original" in df.columns
    assert any(col.startswith("vol_") for col in df.columns)
    assert len(df.columns) == 10


def test_sigmas_argument(sample_series):
    orig = sample_series.copy()
    custom_sigmas = [0.005, 0.01]
    df = generate_volatility_scenarios(orig, sigmas=custom_sigmas, random_seed=123)
    assert len(df.columns) == 1 + len(custom_sigmas)
    for s in custom_sigmas:
        col = f"vol_{s:.4f}"
        assert col in df.columns


def test_match_endpoints_true(sample_series):
    df = generate_volatility_scenarios(
        sample_series, sigmas=[0.01], match_endpoints=True, random_seed=123
    )
    original = df["original"]
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    np.testing.assert_almost_equal(original.iloc[0], synthetic.iloc[0], decimal=8)
    rel_err = abs(original.iloc[-1] - synthetic.iloc[-1]) / abs(original.iloc[-1])
    assert rel_err < 1e-6


def test_match_endpoints_false(sample_series):
    df = generate_volatility_scenarios(
        sample_series, sigmas=[0.01], match_endpoints=False, random_seed=123
    )
    original = df["original"]
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    np.testing.assert_almost_equal(original.iloc[0], synthetic.iloc[0], decimal=8)
    diff = abs(original.iloc[-1] - synthetic.iloc[-1])
    assert diff > 0.01


def test_too_short_input():
    with pytest.raises(ValueError):
        generate_volatility_scenarios(pd.Series([100.0]))


def test_zero_volatility_flat_series(flat_series):
    # Flat series should have zero volatility
    df = generate_volatility_scenarios(flat_series, sigmas=[0.0], random_seed=123)
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    # All prices should be the same as the starting price
    np.testing.assert_array_almost_equal(synthetic.values, [100.0] * len(flat_series))


def test_very_low_volatility(flat_series):
    # Very low volatility, prices should barely move
    df = generate_volatility_scenarios(flat_series, sigmas=[1e-6], random_seed=456)
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    # All prices should be extremely close to 100
    assert np.allclose(synthetic.values, 100.0, atol=1e-3)


def test_custom_sigma_on_flat_series(flat_series):
    # Custom sigma with match_endpoints True should still match start/end
    df = generate_volatility_scenarios(
        flat_series, sigmas=[0.02], match_endpoints=True, random_seed=42
    )
    original = df["original"]
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    # Start and end prices must match
    np.testing.assert_almost_equal(original.iloc[0], synthetic.iloc[0], decimal=8)
    np.testing.assert_almost_equal(original.iloc[-1], synthetic.iloc[-1], decimal=6)
    # Path should have some movement (not flat)
    assert np.any(np.abs(np.diff(synthetic.values)) > 1e-4)


def test_zero_sigma_with_match_endpoints(sample_series):
    # Setting sigma=0 with match_endpoints should result in a deterministic straight line (log-space)
    df = generate_volatility_scenarios(
        sample_series, sigmas=[0.0], match_endpoints=True
    )
    synthetic = df[[col for col in df.columns if col.startswith("vol_")][0]]
    # The synthetic series should be a geometric progression from start to end
    expected = np.exp(
        np.linspace(
            np.log(sample_series.iloc[0]),
            np.log(sample_series.iloc[-1]),
            len(sample_series),
        )
    )
    np.testing.assert_allclose(synthetic.values, expected, rtol=1e-8)
