"""Tests for IV interpolation."""
import pytest
import numpy as np
from qmoms.qmoms import interpolate_iv_by_moneyness


class TestInterpolateIV:

    @pytest.fixture
    def sample_iv_data(self):
        mnes = np.array([0.85, 0.90, 0.95, 1.00, 1.05, 1.10, 1.15])
        vol = np.array([0.30, 0.27, 0.24, 0.22, 0.21, 0.20, 0.20])
        grid = np.array([500, 2])
        return mnes, vol, grid

    def test_output_shapes(self, sample_iv_data):
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert iv.shape == (1001, 1)
        assert ki.shape == (1001, 1)

    def test_grid_endpoints(self, sample_iv_data):
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert abs(ki[0, 0] - 1 / 3) < 1e-6      # 1/(1+k) = 1/3
        assert abs(ki[500, 0] - 1.0) < 1e-10      # midpoint = 1.0
        assert abs(ki[-1, 0] - 3.0) < 1e-6        # 1+k = 3

    def test_flat_extrapolation_left(self, sample_iv_data):
        """IV for very OTM puts (ki < min mnes) should be flat at the boundary IV."""
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        left_mask = ki[:, 0] < min(mnes)
        assert np.all(iv[left_mask, 0] == vol[0])  # vol[0] corresponds to lowest moneyness

    def test_flat_extrapolation_right(self, sample_iv_data):
        """IV for very OTM calls (ki > max mnes) should be flat at the boundary IV."""
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        right_mask = ki[:, 0] > max(mnes)
        assert np.all(iv[right_mask, 0] == vol[-1])  # vol[-1] corresponds to highest moneyness

    def test_atm_value(self, sample_iv_data):
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert abs(iv[500, 0] - 0.22) < 1e-6  # ATM vol = 0.22

    def test_iv_positive(self, sample_iv_data):
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert np.all(iv > 0)

    def test_grid_monotonic(self, sample_iv_data):
        mnes, vol, grid = sample_iv_data
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert np.all(np.diff(ki[:, 0]) > 0)

    def test_different_grid_size(self):
        mnes = np.array([0.9, 0.95, 1.0, 1.05, 1.1])
        vol = np.array([0.25, 0.22, 0.20, 0.19, 0.19])
        grid = np.array([200, 1])
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        assert iv.shape == (401, 1)
        assert ki.shape == (401, 1)
