"""Tests for the GLB (Generalized Lower Bounds) procedure."""
import pytest
import pandas as pd
import numpy as np

from qmoms import (default_params, default_surf_dtype_mapping, filter_options,
                    get_rate_for_maturity, compute_glb_betas, prepare_integration_inputs_surf,
                    compute_glb, get_glb_theta_boundary)


class TestGetGLBThetaBoundary:

    @pytest.fixture(scope="class")
    def sample_betas_coef(self, raw_data):
        _, _, df_return = raw_data
        beta_df = compute_glb_betas(df_return, parallel=False)
        betas_sel = ['Bmj', 'Bm2j', 'Bj2j', 'Bmjj', 'Bm3j', 'Bj3j', 'Bm2jj', 'Bmj2j']
        gm, rho, kp = 2, 1, 4
        # Get first row of betas for id='12490'
        row = beta_df.loc[beta_df.id == '12490'].iloc[0]
        betas = list(row[betas_sel].values)

        # er for the first group
        df_surf, df_rate, _ = raw_data
        df_surf_f = df_surf[default_surf_dtype_mapping.keys()].copy()
        df_surf_f = filter_options(df_surf_f, default_params['filter'])
        grouped = df_surf_f.groupby(['id', 'date', 'days'], group_keys=False)
        ids, group_now = next(iter(grouped))
        _, date_val, days_val = ids
        rate = get_rate_for_maturity(df_rate, date=date_val, days=days_val)
        mat = days_val / 365
        er = np.exp(mat * rate)

        return betas + [gm, rho, kp, er]

    def test_taylor_order_2_returns_four_values(self, sample_betas_coef):
        res = get_glb_theta_boundary(sample_betas_coef, taylor_order=2)
        assert len(res) == 4

    def test_taylor_order_3_returns_four_values(self, sample_betas_coef):
        res = get_glb_theta_boundary(sample_betas_coef, taylor_order=3)
        assert len(res) == 4

    def test_taylor_order_2_has_valid_interval(self, sample_betas_coef):
        res = get_glb_theta_boundary(sample_betas_coef, taylor_order=2)
        Minval1, Maxval1 = res[0], res[1]
        assert not np.isnan(Minval1)
        assert not np.isnan(Maxval1)
        assert Minval1 < Maxval1

    def test_invalid_taylor_order_raises(self, sample_betas_coef):
        with pytest.raises(ValueError):
            get_glb_theta_boundary(sample_betas_coef, taylor_order=4)


class TestComputeGLB:

    @pytest.fixture(scope="class")
    def glb_inputs(self, raw_data):
        df_surf, df_rate, df_return = raw_data
        df_surf_f = df_surf[default_surf_dtype_mapping.keys()].copy()
        df_surf_f = filter_options(df_surf_f, default_params['filter'])

        beta_df = compute_glb_betas(df_return, parallel=False)

        grouped = df_surf_f.groupby(['id', 'date', 'days'], group_keys=False)
        ids, group_now = next(iter(grouped))
        _, date_val, days_val = ids
        rate = get_rate_for_maturity(df_rate, date=date_val, days=days_val)

        inputs = prepare_integration_inputs_surf(
            group_now.mnes, group_now.impl_volatility, days_val, rate, default_params)

        betas_sel = ['Bmj', 'Bm2j', 'Bj2j', 'Bmjj', 'Bm3j', 'Bj3j', 'Bm2jj', 'Bmj2j']
        gm, rho, kp = 2, 1, 4
        row = beta_df.loc[beta_df.id == '12490'].iloc[0]
        betas = list(row[betas_sel].values) + [gm, rho, kp, inputs['er']]

        return inputs, betas

    def test_compute_glb_returns_dict(self, glb_inputs):
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        assert isinstance(result, dict)

    def test_compute_glb_has_expected_keys(self, glb_inputs):
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        expected_keys = {'theta_GLB_TS2', 'GLB_TS2', 'theta_GLB_TS3', 'GLB_TS3'}
        assert set(result.keys()) == expected_keys

    def test_glb_values_finite(self, glb_inputs):
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        for key, val in result.items():
            assert np.isfinite(val), f"{key} is not finite: {val}"

    def test_glb_ts3_reasonable_range(self, glb_inputs):
        """GLB should be a positive lower bound, typically in [0, 1] range annualized."""
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        assert -0.5 < result['GLB_TS3'] < 2.0

    def test_glb_ts2_reasonable_range(self, glb_inputs):
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        assert -0.5 < result['GLB_TS2'] < 2.0

    def test_theta_positive(self, glb_inputs):
        inputs, betas = glb_inputs
        result = compute_glb(inputs['er'], inputs['ki'], inputs['dKi'],
                             inputs['Q'], inputs['mat'], betas)
        assert result['theta_GLB_TS2'] > 0
        assert result['theta_GLB_TS3'] > 0
