import pytest
import pandas as pd
import numpy as np

from qmoms.examples import load_data
from qmoms import default_params, default_surf_dtype_mapping, filter_options, get_rate_for_maturity


@pytest.fixture(scope="session")
def raw_data():
    df_surf, df_rate, df_return = load_data()
    return df_surf, df_rate, df_return


@pytest.fixture(scope="session")
def filtered_surf(raw_data):
    df_surf, df_rate, _ = raw_data
    df = df_surf[default_surf_dtype_mapping.keys()].copy()
    df = filter_options(df, default_params['filter'])
    return df


@pytest.fixture(scope="session")
def first_group(filtered_surf):
    grouped = filtered_surf.groupby(['id', 'date', 'days'], group_keys=False)
    ids, group = next(iter(grouped))
    return ids, group


@pytest.fixture(scope="session")
def first_group_rate(raw_data, first_group):
    _, df_rate, _ = raw_data
    ids, _ = first_group
    _, date_val, days_val = ids
    rate = get_rate_for_maturity(df_rate, date=date_val, days=days_val)
    return rate


@pytest.fixture(scope="session")
def all_serial_results(filtered_surf, raw_data):
    from qmoms import qmoms_compute_bygroup
    _, 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)
    ret_list = [qmoms_compute_bygroup([g, default_params]) for _, g in grouped]
    return pd.DataFrame.from_records(ret_list)
