"""Tests for qmoms_compute and qmoms_compute_bygroup — regression tests with known baselines."""
import pytest
import pandas as pd
import numpy as np

from qmoms import default_params, qmoms_compute, qmoms_compute_bygroup, get_rate_for_maturity
from qmoms import prepare_integration_inputs_surf

# Baseline values from first group (id=12490, date=2023-01-03, days=30)
BASELINE = {
    'nopt': 17.0,
    'smfiv': 0.09397582654321562,
    'mfiv_bkm': 0.09720121456064883,
    'mfiv_bjn': 0.09583697053671561,
    'smfivd': 0.046490939953513825,
    'mfivd_bkm': 0.05525188752160318,
    'mfivd_bjn': 0.052134290918487425,
    'mfis': -0.5311837403375936,
    'mfik': 3.5657934265068563,
    'cvix_sigma2': 0.09538130621117445,
    'cvix_mnes20': 0.09620316629444423,
    'rix': 0.003117596603115752,
    'rixnorm': 0.05642516017026186,
    'tlm_sigma2': 35.49780847988205,
    'tlm_delta20': 34.31435729208929,
    'slopedn': 0.17173997757901813,
    'slopeup': -0.052627040632152865,
}

TOL = 1e-8


class TestPrepareIntegrationInputs:

    def test_status_ok(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        mnes = group.mnes
        iv = group.impl_volatility
        out = prepare_integration_inputs_surf(mnes, iv, days, first_group_rate, default_params)
        assert out['status'] == 1

    def test_output_keys(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        out = prepare_integration_inputs_surf(group.mnes, group.impl_volatility,
                                              days, first_group_rate, default_params)
        expected_keys = {'m', 'k', 'u', 'ki', 'iv', 'ic', 'ip',
                         'currcalls', 'currputs', 'Q', 'dKi', 'er', 'mat', 'status'}
        assert set(out.keys()) == expected_keys

    def test_grid_size(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        out = prepare_integration_inputs_surf(group.mnes, group.impl_volatility,
                                              days, first_group_rate, default_params)
        assert len(out['ki']) == 1001
        assert len(out['iv']) == 1001
        assert out['m'] == 500

    def test_too_few_points_returns_status_zero(self):
        mnes = np.array([0.95, 1.0, 1.05])
        vol = np.array([0.25, 0.22, 0.20])
        out = prepare_integration_inputs_surf(mnes, vol, 30, 0.05, default_params)
        assert out['status'] == 0

    def test_duplicate_moneyness_returns_status_zero(self):
        mnes = np.array([0.90, 0.95, 0.95, 1.00, 1.05])
        vol = np.array([0.28, 0.25, 0.25, 0.22, 0.20])
        out = prepare_integration_inputs_surf(mnes, vol, 30, 0.05, default_params)
        assert out['status'] == 0

    def test_mat_and_er(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        out = prepare_integration_inputs_surf(group.mnes, group.impl_volatility,
                                              days, first_group_rate, default_params)
        assert abs(out['mat'] - days / 365) < 1e-10
        assert abs(out['er'] - np.exp(out['mat'] * first_group_rate)) < 1e-10


class TestQmomsCompute:

    def test_regression_all_moments(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params)
        for key, expected in BASELINE.items():
            actual = qout[key]
            assert abs(actual - expected) < TOL, (
                f"{key}: expected {expected}, got {actual}"
            )

    def test_output_is_series(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='pandas')
        assert isinstance(qout, pd.Series)

    def test_output_dict(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert isinstance(qout, dict)

    def test_variances_positive(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['smfiv'] > 0
        assert qout['mfiv_bkm'] > 0
        assert qout['mfiv_bjn'] > 0

    def test_semivariances_sum(self, first_group, first_group_rate):
        """Up + Down semivariance should approximately equal total variance."""
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        smfiv_total = qout['smfiv']
        smfivd = qout['smfivd']
        smfivu = smfiv_total - smfivd
        assert smfivu > 0
        assert smfivd > 0

    def test_rix_positive(self, first_group, first_group_rate):
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['rix'] > 0

    def test_negative_skewness(self, first_group, first_group_rate):
        """Typical equity option surface should produce negative implied skewness."""
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['mfis'] < 0

    def test_kurtosis_above_three(self, first_group, first_group_rate):
        """Typical equity option surface should produce excess kurtosis > 3."""
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['mfik'] > 3

    def test_bjn_leq_bkm(self, first_group, first_group_rate):
        """MFIV_BJN <= MFIV_BKM by construction (log vs log-modified contract)."""
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['mfiv_bjn'] <= qout['mfiv_bkm'] + 1e-10

    def test_cvix_leq_mfiv(self, first_group, first_group_rate):
        """Corridor VIX should be <= total MFIV."""
        ids, group = first_group
        _, _, days = ids
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=default_params,
                             output='dict')
        assert qout['cvix_sigma2'] <= qout['mfiv_bkm'] + 1e-10

    def test_disable_optional_outputs(self, first_group, first_group_rate):
        """When optional computations disabled, their keys should be absent."""
        ids, group = first_group
        _, _, days = ids
        params = default_params.copy()
        params['semivars'] = {'compute': False}
        params['mfismfik'] = {'compute': False}
        params['cvix'] = {'compute': False}
        params['rix'] = {'compute': False}
        params['tlm'] = {'compute': False}
        params['slope'] = {'compute': False}
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=params,
                             output='dict')
        assert 'smfiv' in qout  # always computed
        assert 'smfivd' not in qout
        assert 'cvix_sigma2' not in qout
        assert 'rix' not in qout
        assert 'tlm_sigma2' not in qout
        assert 'slopedn' not in qout

    def test_too_few_points_returns_nopt_only(self):
        """When there are too few data points, output should still return nopt."""
        mnes = np.array([0.95, 1.0, 1.05])
        vol = np.array([0.25, 0.22, 0.20])
        qout = qmoms_compute(mnes=mnes, vol=vol, days=30, rate=0.05,
                             params=default_params, output='dict')
        assert isinstance(qout, dict)
        assert qout['nopt'] == 3
        assert 'smfiv' not in qout

    def test_atmfwd_mode(self, first_group, first_group_rate):
        """Test that atmfwd=True runs without error and produces valid output."""
        ids, group = first_group
        _, _, days = ids
        params = default_params.copy()
        params['atmfwd'] = True
        qout = qmoms_compute(mnes=group.mnes, vol=group.impl_volatility,
                             days=days, rate=first_group_rate, params=params,
                             output='dict')
        assert qout['smfiv'] > 0
        assert qout['mfiv_bkm'] > 0


class TestQmomsComputeBygroup:

    def test_regression(self, raw_data, filtered_surf):
        _, df_rate, _ = raw_data
        df = get_rate_for_maturity(df_rate, df_surf=filtered_surf.copy())
        grouped = df.groupby(['id', 'date', 'days'], group_keys=False)
        first_name, first_group = next(iter(grouped))
        ret = qmoms_compute_bygroup([first_group, default_params])
        assert isinstance(ret, pd.Series)
        assert ret['id'] == 12490
        assert abs(ret['smfiv'] - BASELINE['smfiv']) < TOL

    def test_output_has_id_date_days(self, raw_data, filtered_surf):
        _, df_rate, _ = raw_data
        df = get_rate_for_maturity(df_rate, df_surf=filtered_surf.copy())
        grouped = df.groupby(['id', 'date', 'days'], group_keys=False)
        _, first_group = next(iter(grouped))
        ret = qmoms_compute_bygroup([first_group, default_params])
        assert 'id' in ret.index
        assert 'date' in ret.index
        assert 'days' in ret.index


class TestFullSerialRun:

    def test_serial_output_shape(self, all_serial_results):
        assert all_serial_results.shape == (3750, 20)

    def test_serial_output_columns(self, all_serial_results):
        expected = ['id', 'date', 'days', 'nopt', 'smfiv', 'mfiv_bkm', 'mfiv_bjn',
                    'smfivd', 'mfivd_bkm', 'mfivd_bjn', 'mfis', 'mfik',
                    'cvix_sigma2', 'cvix_mnes20', 'rix', 'rixnorm',
                    'tlm_sigma2', 'tlm_delta20', 'slopedn', 'slopeup']
        assert list(all_serial_results.columns) == expected

    def test_smfiv_positive(self, all_serial_results):
        assert (all_serial_results['smfiv'] > 0).all()

    def test_mfiv_bkm_positive(self, all_serial_results):
        assert (all_serial_results['mfiv_bkm'] > 0).all()

    def test_no_nan_in_core_moments(self, all_serial_results):
        for col in ['smfiv', 'mfiv_bkm', 'mfiv_bjn']:
            assert all_serial_results[col].notna().all(), f"NaN in {col}"

    def test_mean_smfiv(self, all_serial_results):
        assert abs(all_serial_results['smfiv'].mean() - 0.120205) < 0.001

    def test_mean_mfis(self, all_serial_results):
        assert abs(all_serial_results['mfis'].mean() - (-0.641184)) < 0.001

    def test_unique_ids(self, all_serial_results):
        ids = all_serial_results['id'].unique()
        assert len(ids) == 5
        assert set(ids) == {12490, 14541, 14593, 18542, 93436}
