from qmoms import default_surf_dtype_mapping
default_surf_dtype_mapping = {'id': 'int64', 
                              'date': 'str', 
                              'days': 'int64', 
                              'mnes': 'float64', 
                              'prem': 'float64', 
                              'impl_volatility': 'float64', 
                              'delta': 'float64'}

from qmoms import default_rate_dtype_mapping
default_rate_dtype_mapping = {'date': 'str',  
                              'days': 'int64',
                              'rate': 'float64'}

from qmoms.examples import load_data
df_surf, df_rate = load_data()

def qmoms_compute(mnes, vol, days, rate, params, output='pandas'):
    """
    Computes implied moments for option pricing using interpolated volatility and various moment formulas.

    Parameters:
    - mnes (array-like): moneyness (K/S) of the options.
    - vol (array-like): implied volatilities corresponding to the moneyness.
    - days (int): Days until the options' expiration.
    - rate (float): Risk-free interest rate.
    - params (dict): Configuration parameters containing grid settings and other options.
    - output (str): Specifies the output format ('pandas' for a pandas.Series, else returns a dictionary).

    Returns:
    - pandas.Series or dict: Calculated moments and optionally other metrics based on the provided parameters.

    The function performs the following:
    - Sorts inputs by moneyness.
    - Computes forward prices and interpolation of implied volatilities.
    - Uses Black or Black-Scholes formulas to compute option prices.
    - Calculates implied variances and semi-variances using specified methods.
    - Optionally computes advanced metrics like MFIS/MFIK, Corridor VIX, RIX, Tail Loss Measure,
    and Slopes based on user configurations in `params`.

    #######################################################################################################
    default_params = {'atmfwd': False,  # if True, use FWD as ATM level and Black model for option valuation
                  'grid': {'number_points': 500,  # points*2 + 1
                           'grid_limit': 2},  # 1/(1+k) to 1+k; typically, k = 2
                  'filter': {'mnes_lim': [0, 1000],  # [0.7, 1.3],
                             'delta_put_lim': [-0.5 + 1e-3, 0],  # typically, only OTM
                             'delta_call_lim': [0, 0.5],  # typically, only OTM
                             'best_bid_zero': -1,  # if negative, no restriction on zero bid;
                             'open_int_zero': -1,  # if negative, no restriction on zero open interest;
                             'min_price': 0},  # a limit on the min mid-price
                  'semivars': {'compute': True},  # save semivariances or not, they are computed anyway
                  'mfismfik': {'compute': True},  # compute skewness/ kurtosis
                  'cvix': {'compute': True,  # compute corridor VIX
                           'abs_dev': [0.2],  # absolute deviation from ATM level of 1
                           'vol_dev': [2]},  # deviation in terms of own sigmas from ATM level of 1
                  'rix': {'compute': True},  # compute RIX (Gao,Gao,Song RFS) or not,
                  'tlm': {'compute': True,  # compute Tail Loss Measure (Xiao/ Vilkov) or not
                          'delta_lim': [20],  # OTM put delta for the threshold (pos->neg or neg)
                          'vol_lim': [2]},  # deviation to the left in terms of sigmas from ATM level of 1
                  'slope': {'compute': True,
                            'deltaP_lim': [-0.5, -0.05],
                            # limits to compute the slopedn as Slope from  IV = const + Slope * Delta
                            'deltaC_lim': [0.05, 0.5]
                            # limits to compute the slopeup as -Slope from IV = const + Slope * Delta
                            }
                  }
    #######################################################################################################
    To modify parameters, update the default_params dictionary with the dictionary with desired parameters
    for example, in case you do not want to compute the tlm, and want to change slope parameters, use
    new_params0 = {'tlm': {'compute': False}}
    new_params1 = {'slope': {'compute': True,
                                'deltaP_lim': [-0.4, -0.1],
                                # limits to compute the slopedn as Slope from  IV = const + Slope * Delta
                                'deltaC_lim': [0.05, 0.5]
                                # limits to compute the slopeup as -Slope from IV = const + Slope * Delta
                                }}
    params = default_params | new_params0 | new_params1
    #######################################################################################################
    """

def qmoms_compute_bygroup(groupparams,
                          id=None, rate=None, days=None, date=None,
                          cols_map={'id': 'id',
                                    'date': 'date',
                                    'days': 'days',
                                    'rate': 'rate',
                                    'mnes': 'mnes',
                                    'impl_volatility': 'impl_volatility'}):
    """
    Computes implied moments for grouped option data using specified parameters and column mappings.

    Parameters:
    - groupparams (tuple, list, or dict): If a tuple or list, this should contain the group data and parameters as the
    first and second elements, respectively. If a dict, it is used directly as parameters.
    - id (int, optional): Identifier for the data group, defaults to the first 'id' in the group data as specified in
    cols_map if not provided.
    - rate (float, optional): Risk-free interest rate, defaults to the first 'rate' in the group data as specified in
    cols_map if not provided.
    - days (int, optional): Days until expiration, defaults to the first 'days' in the group data as specified in cols_
    map if not provided.
    - date (any, optional): Date of the data, defaults to the first 'date' in the group data as specified in cols_map
    if not provided.
    - cols_map (dict, optional): A dictionary mapping the expected columns in the group data to the actual column names.
    Default is {'id': 'id', 'date': 'date', 'days': 'days', 'rate': 'rate',
                'mnes': 'mnes', 'impl_volatility': 'impl_volatility'}.

    Returns:
    - pandas.Series: Contains the computed moments along with initial group identifiers such as id, date, and days.

    The function sorts the group data by moneyness, removes duplicates, and computes moments using interpolation.
    It then merges the computed moments with initial data identifiers and returns a pandas Series, using the column
    mappings specified in cols_map for accessing data fields.
    """
    cols_map = cols_map | {'id': 'id',
                           'date': 'date',
                           'days': 'days',
                           'rate': 'rate',
                           'mnes': 'mnes',
                           'impl_volatility': 'impl_volatility'}

    if (isinstance(groupparams, tuple) or isinstance(groupparams, list)) and len(groupparams) == 2:
        group = groupparams[0]
        params = groupparams[1]
    else:
        params = groupparams

    # scalars
    id = id or group[cols_map['id']].iloc[0]
    date = date or group[cols_map['date']].iloc[0]
    days = days or group[cols_map['days']].iloc[0]
    rate = rate or group[cols_map['rate']].iloc[0]

    # the surface is given in two columns
    group = group.sort_values(by=[cols_map['mnes']])
    mnes = group[cols_map['mnes']]
    vol = group[cols_map['impl_volatility']]

    # remove duplicated moneyness points
    goods = ~mnes.duplicated()
    mnes = mnes[goods]
    vol = vol[goods]

    # pre-define the output dict
    res = {'id': id, 'date': date, 'days': days}

    # compute the moments with interpolation
    res_computed = qmoms_compute(mnes, vol, days, rate, params, output='dict')

    # update the output
    res = res | res_computed

    return pd.Series(res)

'''the script with examples is provided in qmoms/examples/qmoms_test.py'''
# import packages and variables
import pandas as pd
from tqdm import *
from multiprocessing import cpu_count, Pool
from qmoms.examples import load_data
from qmoms import default_params, default_surf_dtype_mapping, filter_options, get_rate_for_maturity
from qmoms import qmoms_compute, qmoms_compute_bygroup

# set number of used cores for parallel computation
if cpu_count() > 30:
    CPUUsed = 20
else:
    CPUUsed = cpu_count()

if __name__ == '__main__':
    # load the provided data 
    df_surf, df_rate = load_data()

    #######################################################################
    # FILTER THE DATA USING DEFAULT PARAMS
    # note that you can easily define your own filter -- make sure you select
    # only OTM options with non-repeating moneyness levels
    # we use OTM calls identified by mnes >=1

    # select only relevant columns and order them
    df_surf = df_surf[default_surf_dtype_mapping.keys()]
    # Filter the whole dataset
    df_surf = filter_options(df_surf, default_params['filter'])
    print("loading and filter done")
    #######################################################################

    #######################################################################
    # GROUP THE DATA FOR COMPUTATIONS, AND TEST THE OUTPUT ON ONE GROUP
    df_surf_grouped = df_surf.groupby(['id', 'date', 'days'], group_keys=False)
    # create an iterator object
    group = iter(df_surf_grouped)
    # select the next group
    group_next = next(group)
    ids, group_now = group_next
    # extract group identifiers
    id, date, days = ids
    # extract data on moneyness and implied volatility
    mnes = group_now.mnes
    iv = group_now.impl_volatility
    # interpolate the rate to the exact maturity
    rate = get_rate_for_maturity(df_rate, date=date, days=days)
    # feed to the function computing implied moments
    qout = qmoms_compute(mnes=mnes, vol=iv, days=days, rate=rate, params=default_params)
    print(ids, qout)

    #######################################################################
    # GROUP THE DATA, AND RUN COMPUTATIONS IN A LOOP SEQUENTIALLY
    qmoms_all = {}
    df_surf_grouped = df_surf.groupby(['id', 'date', 'days'], group_keys=False)
    for ids, group_now in tqdm(df_surf_grouped):
        id, date, days = ids
        # extract data on moneyness and implied volatility
        mnes = group_now.mnes
        iv = group_now.impl_volatility
        # interpolate the rate to the exact maturity
        rate = get_rate_for_maturity(df_rate, date=date, days=days)
        # feed to the function computing implied moments
        qout = qmoms_compute(mnes=mnes, vol=iv, days=days, rate=rate, params=default_params)
        qmoms_all[ids] = qout
        pass

    qmoms_all = pd.DataFrame(qmoms_all).T
    qmoms_all.index.names = ['id', 'date', 'days']
    print(qmoms_all)

    #######################################################################
    # GROUP THE DATA, AND RUN COMPUTATIONS FOR THE WHOLE DATAFRAME
    # first, add the interest rate for each maturity to the surface dataframe
    df_surf = get_rate_for_maturity(df_rate, df_surf=df_surf)
    # group dataset by id,date,days
    grouped = df_surf.groupby(['id', 'date', 'days'], group_keys=False)

    # PARALLEL
    with Pool(CPUUsed) as p:
        ret_list = p.map(qmoms_compute_bygroup, [[group_indv, default_params] for name, group_indv in grouped])
        out_p = pd.DataFrame.from_records(ret_list)
        pass
    print(out_p)

    # SERIAL
    ret_list = [qmoms_compute_bygroup([group_indv, default_params]) for name, group_indv in grouped]
    out_s = pd.DataFrame.from_records(ret_list)
    print(out_s)