'''
MIT License

Copyright (c) 2008-2026 Grigory Vilkov

Contact: vilkov@vilkov.net www.vilkov.net

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
'''

# import packages and variables
import pandas as pd
import numpy as np
import scipy.stats as ss
from scipy.interpolate import PchipInterpolator
from scipy.optimize import minimize
import sympy
from sympy import Symbol
from sympy.solvers.inequalities import reduce_rational_inequalities
import warnings
from multiprocessing import cpu_count, Pool

# Ignore RuntimeWarning
warnings.filterwarnings("ignore", category=RuntimeWarning)


# *****************************************************************************
# compute the moments from a dataframe-based group
# *****************************************************************************
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 = {'id': 'id',
                'date': 'date',
                'days': 'days',
                'rate': 'rate',
                'mnes': 'mnes',
                'impl_volatility': 'impl_volatility'} | cols_map

    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)

# *****************************************************************************
# % compute the moments for one id/date/days combo
# *****************************************************************************
def prepare_integration_inputs_surf(mnes, vol, days, rate, params):
    """
    Prepares inputs for integration in surface-based implied volatility calculations
    by processing and interpolating option market data.

    Parameters:
    - mnes (array-like): Moneyness (K/S) of the options, where K is the strike price and S is the stock price.
    - vol (array-like): Implied volatilities corresponding to the moneyness values provided.
    - days (int): Days until the options' expiration.
    - rate (float): Risk-free interest rate.
    - params (dict): Configuration parameters that define grid settings.

    Returns:
    - dict: A dictionary containing the computed integration inputs such as the interpolated
      implied volatilities, computed Black-Scholes or Black prices for OTM calls and puts,
      forward prices, and other intermediate calculations, depending on the 'params' configuration.
      The function returns detailed status in the output dict indicating success or failure.

    Detailed Process:
    - Inputs are sorted by moneyness.
    - Forward prices are determined and implied volatilities are interpolated.
    - Depending on the model specified in 'params', uses either the Black or Black-Scholes formula
      to compute prices for OTM options.
    - Calculates additional metrics based on the configuration provided in 'params', such as
      implied variances and semi-variances, and potentially complex metrics like Corridor VIX,
      using specific sub-parameters for calculation.

    Example of modifying parameters:
    - To adjust parameters for integration, update the default_params dictionary. For instance,
      to disable the calculation of the Tail Loss Measure and modify slope parameters:
        new_params0 = {'tlm': {'compute': False}}
        new_params1 = {'slope': {'compute': True,
                                 'deltaP_lim': [-0.4, -0.1],
                                 'deltaC_lim': [0.05, 0.5]}}
        params = default_params | new_params0 | new_params1
    """
    mnes = np.array(mnes)
    vol = np.array(vol)
    ii = np.argsort(mnes)
    mnes = mnes[ii]
    vol = vol[ii]

    # assign grid parameters to variables
    gridparams = params.get('grid', {})
    gridparams_m = gridparams.get('number_points', 500)
    gridparams_k = gridparams.get('grid_limit', 2)
    grid = np.array([gridparams_m, gridparams_k])
    ##########################################
    # define maturity and forward price (for S0=1)
    ##########################################
    mat = days / 365
    er = np.exp(mat * rate)
    ##########################################
    # compute the moments WITH INTERPOLATION
    ##########################################
    nopt = len(mnes)
    # dicct = {'nopt': nopt}

    if (nopt >= 4) & (len(np.unique(mnes)) == len(mnes)):
        #######################################################################
        # interpolate
        iv, ki = interpolate_iv_by_moneyness(mnes, vol, grid)
        # calculate BS prices - OTM calls - mnes >=1
        otmcalls = ki >= 1
        otmputs = ~otmcalls
        kicalls = ki[otmcalls]
        kiputs = ki[otmputs]

        if params.get('atmfwd', False):  # use Black formula
            # OTM calls
            currcalls, itmputs = Black(1, kicalls, rate, iv[otmcalls], mat)
            # OTM puts
            itmcalls, currputs = Black(1, kiputs, rate, iv[otmputs], mat)
        else:
            # OTM calls
            currcalls, itmputs = BlackScholes(1, kicalls, rate, iv[otmcalls], mat)
            # OTM puts
            itmcalls, currputs = BlackScholes(1, kiputs, rate, iv[otmputs], mat)

        #######################################################################
        # use the gridparams to define some variables
        m = gridparams_m  # use points -500:500, i.e. 1001 points for integral approximation
        k = gridparams_k  # use moneyness 1/(1+2) to 1+2 - should be enough as claimed by JT
        u = np.power((1 + k), (1 / m))

        mi = np.arange(-m, m + 1)  # exponent for the grid
        mi.shape = otmcalls.shape
        ic = mi[otmcalls]  # exponent for OTM calls
        ip = mi[otmputs]

        ### compute some general input to implied variances
        Q = np.append(currputs, currcalls)  # all option prices
        Q.shape = (len(Q), 1)
        dKi = np.empty((len(Q), 1))  # dK
        dKi[:] = np.nan
        x = (ki[2:, ] - ki[0:-2]) / 2
        x = np.reshape(x, [ki.shape[0] - 2, 1])
        dKi[1:-1] = x
        dKi[0] = ki[1] - ki[0]
        dKi[-1] = ki[-1] - ki[-2]
        dKi = abs(dKi)

        out = dict(m=m, k=k, u=u, ki=ki.flatten(), iv=iv.flatten(), ic=ic.flatten(), ip=ip.flatten(),
                   currcalls=currcalls.flatten(), currputs=currputs.flatten(), Q=Q.flatten(), dKi=dKi.flatten(),
                   er=er, mat=mat, status=1)
    else:
        out = dict(status=0)

    return out



# *****************************************************************************
# % compute the moments for one id/date/days combo
# *****************************************************************************
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
    #######################################################################################################
    """
    mnes = np.array(mnes)
    vol = np.array(vol)
    ii = np.argsort(mnes)
    mnes = mnes[ii]
    vol = vol[ii]

    # assign grid parameters to variables
    # gridparams = params.get('grid', {})
    # gridparams_m = gridparams.get('number_points', 500)
    # gridparams_k = gridparams.get('grid_limit', 2)
    # grid = np.array([gridparams_m, gridparams_k])
    ##########################################
    # define maturity and forward price (for S0=1)
    ##########################################

    # initialize the output dictionaty
    dicct = {'nopt': len(mnes)}

    # inputs: mnes, vol, days, rate, params
    inputs_surf = prepare_integration_inputs_surf(mnes, vol, days, rate, params)
    # Unpacking the dictionary into variables
    status = inputs_surf['status']
    ##########################################
    # compute the moments WITH INTERPOLATION
    ##########################################
    if status:
        m = inputs_surf['m']
        k = inputs_surf['k']
        ki = inputs_surf['ki']
        iv = inputs_surf['iv']
        u = inputs_surf['u']
        ic = inputs_surf['ic']
        ip = inputs_surf['ip']
        currcalls = inputs_surf['currcalls']
        currputs = inputs_surf['currputs']
        Q = inputs_surf['Q']
        dKi = inputs_surf['dKi']
        er = inputs_surf['er']
        mat = inputs_surf['mat']

        #
        otmcalls = ki >= 1
        otmputs = ~otmcalls
        kicalls = ki[otmcalls]
        kiputs = ki[otmputs]

        #######################################################################
        ### compute SMFIV -- Ian Martin
        K0sq = 1
        # inputs for semivars:
        svar_ingredients = np.divide(np.multiply(dKi, Q), K0sq)
        svar_multiplier = 2 * er
        # semivariance
        moments_smfivu = svar_multiplier * np.sum(svar_ingredients[otmcalls]) / mat
        moments_smfivd = svar_multiplier * np.sum(svar_ingredients[otmputs]) / mat
        # total variance
        moments_smfiv = moments_smfivu + moments_smfivd

        #######################################################################
        ### compute MFIV_BJN - Britten-Jones and Neuberger (2002)
        Ksq = ki ** 2  # denominator
        # inputs for semivars:
        bjn_nom = np.multiply(dKi, Q)
        bjn_ingredients = np.divide(bjn_nom, Ksq)
        bjn_multiplier = 2 * er
        # semivariance
        moments_mfivu_bjn = bjn_multiplier * np.sum(bjn_ingredients[otmcalls]) / mat
        moments_mfivd_bjn = bjn_multiplier * np.sum(bjn_ingredients[otmputs]) / mat
        # total variance
        moments_mfiv_bjn = moments_mfivu_bjn + moments_mfivd_bjn

        #######################################################################
        ### compute MFIV_BKM - Bakshi, Kapadia and Madan (2003)
        Ksq = ki ** 2
        # inputs for semivars:
        bkm_nom = np.multiply(dKi, np.multiply(1 - np.log(ki), Q))
        bkm_ingredients = np.divide(bkm_nom, Ksq)
        bkm_multiplier = 2 * er
        # semivariance
        moments_mfivu_bkm = bkm_multiplier * np.sum(bkm_ingredients[otmcalls]) / mat
        moments_mfivd_bkm = bkm_multiplier * np.sum(bkm_ingredients[otmputs]) / mat
        # total variance
        moments_mfiv_bkm = moments_mfivu_bkm + moments_mfivd_bkm

        # OUTPUT 1
        dicct.update({'smfiv': moments_smfiv,
                      'mfiv_bkm': moments_mfiv_bkm,
                      'mfiv_bjn': moments_mfiv_bjn})

        if params.get('semivars', {}).get('compute', False):
            dicct.update({'smfivd': moments_smfivd,
                          'mfivd_bkm': moments_mfivd_bkm,
                          'mfivd_bjn': moments_mfivd_bjn})

        #######################################################################
        ### compute some inputs for MFIS/MFIK based on BKM
        a = 2 * (u - 1)
        b1 = 1 - (np.log(1 + k) / m) * ic
        b1.shape = currcalls.shape
        b1 = np.multiply(b1, currcalls)
        h = np.power(u, ic)
        h.shape = currcalls.shape
        b1 = np.divide(b1, h)
        b2 = 1 - (np.log(1 + k) / m) * ip
        b2.shape = currputs.shape
        b2 = np.multiply(b2, currputs)
        g = np.power(u, ip)
        g.shape = currputs.shape
        b2 = np.divide(b2, g)
        b_all = np.concatenate((b2, b1))  # note: puts first because ki has moneyness range from puts to calls

        # MFIS/ MFIK
        if params.get('mfismfik', {}).get('compute', False):
            mfismfik_dic = compute_skew_kurtosis(m, k, u, ic, ip, currcalls, currputs, er,
                                                 moments_mfiv_bkm * mat / er)  # adjustment /er on 2023-03-30
            dicct.update(mfismfik_dic)

        #######################################################################
        # CORRIDOR VIX / Bondarenko et al
        if params.get('cvix', {}).get('compute', False):
            cvix_dic = cvix_func(b_all, params['cvix'], ki, moments_mfiv_bkm, mat, a)
            dicct.update(cvix_dic)

        #######################################################################
        # RIX - Gao Gao Song RFS
        if params.get('rix', {}).get('compute', False):
            rix = moments_mfivd_bkm - moments_mfivd_bjn
            rixn = rix / moments_mfivd_bkm
            dicct.update({'rix': rix, 'rixnorm': rixn})

        #######################################################################
        # TAIL LOSS MEASURE Hamidieh / Vilkov and Xiao
        if params.get('tlm', {}).get('compute', False):
            if params.get('atmfwd', False):
                _, currputd = Black_delta(1, kiputs, rate, iv[~otmcalls], mat)
            else:
                _, currputd = BlackScholes_delta(1, kiputs, rate, iv[~otmcalls], mat)
            tlm_dic = tlm_func(currputs, kiputs, currputd, params['tlm'], moments_mfiv_bkm, mat)
            dicct.update(tlm_dic)

        #######################################################################
        # SLOPE Kelly et al + applications in Carbon Tail Risk and Pricing Climate Change Exposure
        if params.get('slope', {}).get('compute', False):
            delta_lim = params['slope'].get('deltaP_lim', [-0.5, -0.05])
            ivputs = iv[otmputs]
            if params.get('atmfwd', False):
                _, currputd = Black_delta(1, kiputs, rate, ivputs, mat)
            else:
                _, currputd = BlackScholes_delta(1, kiputs, rate, ivputs, mat)
            selputs = (currputd >= min(delta_lim)) & (currputd <= max(delta_lim))
            dnow = currputd[selputs]
            ivnow = ivputs[selputs]

            if sum(selputs) > 3:
                slope, intercept, r_value, p_value, std_err = ss.linregress(dnow, ivnow)
            else:
                slope = np.nan
            dicct.update({'slopedn': slope})

            # same for calls
            delta_lim = params['slope'].get('deltaC_lim', [0.05, 0.5])
            ivcalls = iv[otmcalls]
            if params.get('atmfwd', False):
                currcalld, _ = Black_delta(1, kicalls, rate, ivcalls, mat)
            else:
                currcalld, _ = BlackScholes_delta(1, kicalls, rate, ivcalls, mat)
            selcalls = (currcalld >= min(delta_lim)) & (currcalld <= max(delta_lim))
            dnow = currcalld[selcalls]
            ivnow = ivcalls[selcalls]

            if sum(selcalls) > 3:
                slope, intercept, r_value, p_value, std_err = ss.linregress(dnow, ivnow)
            else:
                slope = np.nan
            dicct.update({'slopeup': -slope})

        #######################################################################
    out = dicct
    if output == 'pandas':
        out = pd.Series(dicct)
    return out


### compute MFIS/MFIK
def compute_skew_kurtosis(m, k, u, ic, ip, currcalls, currputs, er, V):
    u = u.ravel()
    ic = ic.ravel()
    ip = ip.ravel()
    currcalls = currcalls.ravel()
    currputs = currputs.ravel()

    result_dict = {}

    a = 3 * (u - 1) * np.log(1 + k) / m
    b1 = np.sum(ic * (2 - (np.log(1 + k) / m) * ic) * currcalls / u ** ic)
    b2 = np.sum(ip * (2 - (np.log(1 + k) / m) * ip) * currputs / u ** ip)
    W = a * (b1 + b2)

    a = 4 * (u - 1) * (np.log(1 + k) / m) ** 2
    b1 = np.sum(ic ** 2 * (3 - (np.log(1 + k) / m) * ic) * currcalls / u ** ic)
    b2 = np.sum(ip ** 2 * (3 - (np.log(1 + k) / m) * ip) * currputs / u ** ip)
    X = a * (b1 + b2)

    mu = er - 1 - er / 2 * V - er / 6 * W - er / 24 * X
    c = (er * V - mu ** 2)

    mfis = (er * W - 3 * mu * er * V + 2 * mu ** 3) / (c ** (3 / 2))
    mfik = (er * X - 4 * mu * er * W + 6 * er * mu ** 2 * V - 3 * mu ** 4) / c ** 2

    result_dict['mfis'] = mfis.item()
    result_dict['mfik'] = mfik.item()

    return result_dict


### compute TAIL LOSS MEASURE
def tlm_func(currputs, kiputs, currputd, tlm_params, iv, mat):
    currputs = currputs.ravel()
    kiputs = kiputs.ravel()
    currputd = currputd.ravel()
    result_dict = {}

    sd_ls = tlm_params.get('vol_lim', [])
    absD_ls = tlm_params.get('delta_lim', [])

    # define time varying threshold
    iv_m = np.sqrt(iv * mat)  # convert to horizon = days
    for r in sd_ls:
        tlm = np.nan
        k0 = 1.0 - r * iv_m
        goods = kiputs <= k0
        if np.sum(goods) > 1:
            Pt = currputs[goods]
            K = kiputs[goods]
            P0 = np.max(Pt)
            ind = np.argmax(Pt)
            K0 = K[ind]
            try:
                res = minimize(f_tlm, np.array([2, 0.01]), args=(P0, K0, Pt, K,),
                               method='SLSQP',
                               options={'ftol': 1e-32, 'disp': False})
                tlm = res.x[1] / (1 - res.x[0])
            except (ValueError, RuntimeError):
                pass
        result_dict['tlm_sigma' + str(r)] = tlm

    for r in absD_ls:
        tlm = np.nan
        if abs(r) > 1:
            r = r / 100
        goods = abs(currputd) <= abs(r)
        if np.sum(goods) > 1:
            Pt = currputs[goods]
            K = kiputs[goods]
            P0 = np.max(Pt)
            ind = np.argmax(Pt)
            K0 = K[ind]
            try:
                res = minimize(f_tlm, np.array([2, 0.01]), args=(P0, K0, Pt, K,),
                               method='SLSQP', options={'ftol': 1e-32, 'disp': False})
                tlm = res.x[1] / (1 - res.x[0])
            except (ValueError, RuntimeError):
                pass
        result_dict['tlm_delta' + str(int(abs(r * 100)))] = tlm
    return result_dict


###############################################################################
def f_tlm(X, P0, K0, Pt, K):
    # prices = Pt
    # strikes = K
    xi = X[0]
    beta = X[1]
    checkterms = 1 + (xi / beta) * (K0 - K)
    term2 = np.multiply(P0, np.divide(np.power(checkterms, (1 - 1 / xi)), Pt))
    return np.sum((1 - term2) ** 2)


### compute CORRIDOR VIX
def cvix_func(b_all, cvix_params, ki, iv, mat, a):
    '''
    Computes corridor variance based on r relative (sigma) deviations
    or specified absolute deviations from ATM moneyness of 1
    '''
    # change a to adjust for maturity
    a = a / mat

    sd_ls = cvix_params.get('vol_dev', [])
    absD_ls = cvix_params.get('abs_dev', [])

    result_dict = {}

    # define time varying threshold for semi-variances
    iv_m = np.sqrt(iv * mat)  # convert to horizon = days

    b_all.shape = ki.shape

    # compute CVIX based on r standard deviations (mfiv**0.5)
    for r in sd_ls:
        kl = 1 - r * iv_m
        ku = 1 + r * iv_m
        sel_ki = (ki >= kl) & (ki <= ku)  # range for cvix
        result_dict['cvix_sigma' + str(r)] = a * np.sum(b_all[sel_ki])

    # compute CVIX based on absolute deviation
    for r in absD_ls:
        kl = 1 - r
        ku = 1 + r
        sel_ki = (ki >= kl) & (ki <= ku)  # range for CVIX
        result_dict['cvix_mnes' + str(int(r * 100))] = a * np.sum(b_all[sel_ki])

    return result_dict


# *****************************************************************************
# COMPUTE BS AND BS DELTA
# *****************************************************************************
def BlackScholes(S0, K, r, sigma, T):
    d1 = (np.log(S0 / K) + (r + np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    d2 = (np.log(S0 / K) + (r - np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    c = S0 * ss.norm.cdf(d1) - K * np.exp(-r * T) * ss.norm.cdf(d2)
    p = K * np.exp(-r * T) * ss.norm.cdf(-d2) - S0 * ss.norm.cdf(-d1)
    return c, p


def Black(F, K, r, sigma, T):
    d1 = (np.log(F / K) + (np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    d2 = (np.log(F / K) + (-np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    c = (F * ss.norm.cdf(d1) - K * ss.norm.cdf(d2)) * np.exp(-r * T)
    p = (K * ss.norm.cdf(-d2) - F * ss.norm.cdf(-d1)) * np.exp(-r * T)
    return c, p


def BlackScholes_delta(S0, K, r, sigma, T):
    d1 = (np.log(S0 / K) + (r + np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    delta_c = ss.norm.cdf(d1)
    delta_p = delta_c - 1
    return delta_c, delta_p


def Black_delta(F, K, r, sigma, T):
    d1 = (np.log(F / K) + (np.square(sigma) / 2) * T) / (sigma * np.sqrt(T))
    delta_c = ss.norm.cdf(d1) * np.exp(-r * T)
    delta_p = -ss.norm.cdf(-d1) * np.exp(-r * T)
    return delta_c, delta_p


# *****************************************************************************
# % Function OM_Interpolate_IV
# *****************************************************************************
def interpolate_iv_by_moneyness(mnes, vol, grid):
    # set the grid in terms of moneyness that we want to use to compute the
    # MFIV/ MFIS and other integrals as needed
    m = grid[0]  # use points -500:500, i.e. 1001 points for integral approximation
    k = grid[1]  # use moneyness 1/(1+2) to 1+2 - should be enough as claimed by JT
    u = np.power((1 + k), (1 / m))

    # create strikes using s=1, i.e. get k/s = moneyness
    mi = np.arange(-m, m + 1)
    ki = np.power(u, mi)
    iv = np.empty((len(ki), 1))
    iv[:] = np.nan
    ki.shape = iv.shape

    currspline = PchipInterpolator(mnes, vol, axis=0)

    k_s_max = max(mnes)  # OTM calls
    k_s_min = min(mnes)  # OTM puts
    iv_max = vol[0]  # for more OTM puts i.e we have iv_max for min k/s, i.e. for OTM put option
    iv_min = vol[-1]  # for more OTM calls  i.e. we have iv_min for OTM call option

    # calculate the interpolated/ extrapolated IV for these k/s
    ks_larger_ind = ki > k_s_max  # more OTM call
    ks_smaller_ind = ki < k_s_min  # more OTM put
    ks_between_ind = (ki >= k_s_min) & (ki <= k_s_max)

    if sum(ks_larger_ind) > 0:
        iv[ks_larger_ind] = iv_min

    if sum(ks_smaller_ind) > 0:
        iv[ks_smaller_ind] = iv_max

    # evaluate the spline at ki[ks_between_ind]
    if sum(ks_between_ind) > 0:
        s = currspline(ki[ks_between_ind], nu=0, extrapolate=None)
        s.shape = iv[ks_between_ind].shape
        iv[ks_between_ind] = s
    return iv, ki


# *****************************************************************************
# % Function to compute betas for Generalized Lower Bounds
# *****************************************************************************
def compute_glb_betas_id(data_temp=pd.DataFrame(), id='ret', win_obs=252, min_obs=126):
    """
        Computes inputs to the GLB procedure for a given asset over specified rolling windows using
        historical return data, market factor and riskfree rate.

        This function processes a DataFrame of returns for a single
        asset identified by 'id', market factor identifiied by 'mktrf' and riskfree rate 'rf'

        Parameters:
        data_temp (pandas.DataFrame, optional): DataFrame containing the asset return data
                                                along with market and risk-free rate information.
                                                It must include columns for 'rf' and 'mktrf' and
                                                a column as specified by 'id'.
                                                Normally, daily data is used.
        id (str, optional): The column name in data_temp that identifies the asset returns.
                            Default is 'ret'.
        win_obs (int, optional): The size of the rolling window (in days) for the beta calculation.
                                 Default is 252.
        min_obs (int, optional): The minimum number of observations required within the window to
                                 perform the calculations. Default is 126.

        Returns:
        pandas.DataFrame: A DataFrame containing computed beta values for each window. The columns
                          include 'id', 'date', and various beta metrics like 'Bmj', 'Bm2j', etc.

        Raises:
        ValueError: If the required columns are not present in the input DataFrame.

        Example:
        >>> data = pd.DataFrame({
            'ret': np.random.randn(300),
            'mktrf': np.random.randn(300),
            'Rf': np.random.rand(300)
        })
        >>> betas = compute_glb_betas_id(data_temp=data, id='ret', win_obs=250, min_obs=120)
        >>> print(betas.head())

        Note:
        The function assumes the input DataFrame includes necessary market and risk-free rate data
        and is indexed by datetime. It processes data to compute betas based on variations of market
        data manipulations and asset returns adjusted for the risk-free rate. The calculations are
        sensitive to the presence of required data and the configuration of the DataFrame index.
        """

    if data_temp.shape[0] == 0:
        return pd.DataFrame()

    beta_cols = ['id', 'date', 'Bmj', 'Bm2j', 'Bj2j', 'Bmjj', 'Bm3j', 'Bj3j', 'Bm2jj', 'Bmj2j', 'Bcapm']
    # add some columns to the data
    cols_added = ['ret_rf', 'ret_rf2', 'ret_rf_mktrf', 'ret_rf3', 'mktrf2_ret_rf', 'mktrf_ret_rf2']
    data_temp[cols_added] = np.nan
    data_temp['ret_rf'] = data_temp[id] - data_temp['Rf']
    data_temp['ret_rf2'] = data_temp['ret_rf'] ** 2
    data_temp['ret_rf_mktrf'] = data_temp['ret_rf'] * data_temp['mktrf']
    data_temp['ret_rf3'] = data_temp['ret_rf'] ** 3
    data_temp['mktrf2_ret_rf'] = data_temp['mktrf2'] * data_temp['ret_rf']
    data_temp['mktrf_ret_rf2'] = data_temp['mktrf'] * data_temp['ret_rf2']

    vars_for_cov = ['mktrf',  # 0
                    'mktrf2',  # 1
                    'ret_rf2',  # 2
                    'ret_rf_mktrf',  # 3
                    'mktrf3',  # 4
                    'ret_rf3',  # 5
                    'mktrf2_ret_rf',  # 6
                    'mktrf_ret_rf2']  # 7

    betas_all = []
    for t in data_temp.index[min_obs:]:
        data_reg = data_temp.loc[t - pd.DateOffset(days=win_obs):t, :].dropna()
        if data_reg.shape[0] >= min_obs:
            ttt = np.cov(data_reg[[*vars_for_cov, id]], rowvar=False)
            var_ind = ttt[-1, -1]
            var_mkt = ttt[0, 0]
            covs_all = ttt[:, -1]
            # compute all betas except for Bcapm
            Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j = covs_all[:8] / var_ind
            # compute Bcapm
            ttt = np.cov(data_reg[['mktrf', 'ret_rf']], rowvar=False)
            cov = ttt[0, 1]
            Bcapm = cov / var_mkt
            # collect
            temp_betas = [id, t, Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j, Bcapm]
            betas_all.append(temp_betas)
    out = pd.DataFrame(betas_all, columns=beta_cols)
    return out

def compute_glb_betas_wrap(kwargs):
    """
        Wraps the compute_glb_betas_id function, forwarding keyword arguments.

        This function serves as a wrapper that takes a dictionary of arguments and unpacks
        them into the compute_glb_betas_id function. It is particularly useful for scenarios
        where keyword arguments need to be dynamically passed to compute_glb_betas_id, such
        as when using multiprocessing where the function and its arguments need to be pickled.

        Parameters:
        kwargs (dict): A dictionary of keyword arguments that are passed to the
                       compute_glb_betas_id function. Keys should match the expected
                       parameter names of the compute_glb_betas_id function.

        Returns:
        The return value from compute_glb_betas_id, which depends on its implementation.
        Typically, this might include computational results such as beta coefficients or
        other statistics as determined by the function's logic.

        Example:
        >>> kwargs = {'data_temp': df_sample, 'id': 'sample_id', 'win_obs': 250, 'min_obs': 125}
        >>> result = compute_glb_betas_wrap(kwargs)
        >>> print(result)

        Note:
        Ensure that all keys in the kwargs dictionary correctly correspond to the parameter
        names expected by compute_glb_betas_id to avoid TypeError exceptions.
        """
    return compute_glb_betas_id(**kwargs)

def compute_glb_betas(df_ret, win_obs=251, min_obs=126, ids=[], parallel=False, cpuused=1):
    """
        Computes a set of inputs for the GLB computation procedure .

        This function adjusts return data, computes necessary return transformations,
        and calculates beta coefficients in a parallelized or non-parallelized manner
        based on the specified parameters. Beta calculations are performed using
        the compute_glb_betas_id function, potentially in parallel across multiple CPUs.

        Parameters:
        df_ret_in (pandas.DataFrame): DataFrame containing data
                                    with asset returns (column name = asset ID),
                                    market returns ('mktrf') and risk-free rate ('rf').
        Note that the index must be in datetime format!
        win_obs (int, optional): Window of observations to consider for beta calculation.
                                 Default is 251.
        min_obs (int, optional): Minimum number of observations required to perform
                                 the calculation. Default is 126.
        ids (list, optional): List of identifiers for which betas are to be calculated.
                              If empty, betas are calculated for all columns except 'mktrf' and 'rf'.
        parallel (bool, optional): Flag to determine if the computation should be
                                   done in parallel. Default is False.
        cpuused (int, optional): Number of CPUs to use for parallel processing.
                                 If 1 or less, it uses all available CPUs (with parallel=True). Default is 1.

        Returns:
        pandas.DataFrame: DataFrame containing beta coefficients for each identifier.

        Raises:
        ValueError: If input data frame df_ret_in does not contain necessary columns.

        Notes:
        - The function assumes that 'mktrf' and 'rf' columns are present in the input DataFrame.
        - Beta calculations might be resource-intensive depending on the size of the data and
          the number of CPUs used.
        """

    df_ret_in = df_ret.copy()

    cols_present = [z for z in df_ret_in.columns if z not in ['mktrf','rf']]
    ids_ls = ids if len(ids)>0 else cols_present

    df_ret_in[ids_ls] = df_ret_in[ids_ls] + 1
    df_ret_in['mkt'] = df_ret_in['mktrf'] + df_ret_in['rf']
    df_ret_in['mktrf2'] = df_ret_in['mktrf'] ** 2
    df_ret_in['mktrf3'] = df_ret_in['mktrf'] ** 3

    # make risk free gross and market gross
    df_ret_in[['mkt', 'Rf']] = df_ret_in[['mkt', 'rf']] + 1
    # ret_df = ret_df.rename(columns={'rf': 'Rf'})

    # define some parameters for beta computation
    # beta_cols = ['id', 'date', 'Bmj', 'Bm2j', 'Bj2j', 'Bmjj', 'Bm3j', 'Bj3j', 'Bm2jj', 'Bmj2j', 'Bcapm']
    cols_rel = ['mkt', 'mktrf', 'mktrf2', 'mktrf3', 'Rf']

    if parallel:
        cpuused = cpu_count() if cpuused<=1 else cpuused
        with Pool(cpuused) as p:
            beta_df = p.map(compute_glb_betas_wrap,
                            [{'data_temp': df_ret_in[[id, *cols_rel]].copy(),
                              'id':id, 'win_obs':win_obs, 'min_obs':min_obs} for id in ids_ls])
    else:
        beta_df = [compute_glb_betas_wrap(kwargs) for kwargs in
                   [{'data_temp': df_ret_in[[id, *cols_rel]].copy(),
                     'id': id, 'win_obs': win_obs, 'min_obs': min_obs} for id in ids_ls]]

    beta_df = pd.concat(beta_df)
    return beta_df


# *****************************************************************************
# GLB (Generalized Lower Bounds) - Chabi-Yo, Dim, Vilkov (MS, 2023)
# *****************************************************************************
def get_glb_theta_boundary(betas_coef, taylor_order):
    """
    Solves the polynomial inequality for the feasible theta range in the GLB procedure.

    Parameters:
    - betas_coef (list): [Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j, gm, rho, kp, Rf]
      where gm=gamma, rho, kp=kappa are risk-aversion parameters and Rf is the gross risk-free rate.
    - taylor_order (int): 2 or 3 for the order of Taylor expansion.

    Returns:
    - list: [Minval1, Maxval1, Minval2, Maxval2] boundaries of the feasible theta intervals.
      If only one interval exists, Minval2 and Maxval2 are NaN.
    """
    Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j, gm, rho, kp, Rf = betas_coef

    th = Symbol('th', real=True)

    if taylor_order == 2:
        equation = th + (th - 1) * th * Bj2j / (2 * Rf) <= gm * (Bmj + th * Bmjj / Rf) - rho * gm * Bm2j / Rf
    elif taylor_order == 3:
        equation = th * (1 - gm * Bmj / th + gm * rho * Bm2j / (th * Rf) + (th - 1) * Bj2j / (2 * Rf) - gm * Bmjj / Rf -
                         kp * rho * gm * Bm3j / (th * Rf ** 2) + (th - 1) * (th - 2) * Bj3j / (6 * Rf ** 2) +
                         gm * rho * Bm2jj / (Rf ** 2) - (th - 1) * gm * Bmj2j / (2 * Rf ** 2)) <= 0
    else:
        raise ValueError(f"taylor_order must be 2 or 3, got {taylor_order}")

    solution = reduce_rational_inequalities([[equation]], th, relational=False)

    eq_res = list(solution.atoms())
    eq_res = [val for val in eq_res if val not in [sympy.false, sympy.true]]
    eq_res.sort()
    res = [np.nan] * 4

    if len(eq_res) == 4:
        Minval1, Maxval1, Minval2, Maxval2 = [float(v) for v in eq_res]
        res = [Minval1, Maxval1, Minval2, Maxval2]
    elif len(eq_res) == 2:
        Minval1, Maxval1 = [float(v) for v in eq_res]
        res = [Minval1, Maxval1, np.nan, np.nan]

    return res


def _compute_glb_moment(er, ki, dKi, Q, theta):
    """Computes the theta-power moment from option data for GLB."""
    return er ** (theta + 1) + theta * (theta + 1) * er * np.sum(np.power(ki, theta - 1) * Q * dKi)


def _optimize_glb_theta(theta, params_opt):
    """Objective function for GLB theta optimization (negated for maximization)."""
    theta = theta[0]
    er, ki, dKi, Q = params_opt[0]
    nume = _compute_glb_moment(er, ki, dKi, Q, theta)
    denom = _compute_glb_moment(er, ki, dKi, Q, theta - 1)
    obj = nume / denom - er
    return -obj


def compute_glb(er, ki, dKi, Q, mat, betas_coef):
    """
    Computes Generalized Lower Bounds on expected returns for one asset/date.

    Based on Chabi-Yo, Dim, and Vilkov (Management Science, 2023):
    https://ssrn.com/abstract=3565130

    Parameters:
    - er (float): exp(rate * mat), the gross risk-free return.
    - ki (array): Moneyness grid from prepare_integration_inputs_surf.
    - dKi (array): Grid spacing from prepare_integration_inputs_surf.
    - Q (array): OTM option prices from prepare_integration_inputs_surf.
    - mat (float): Time to maturity in years.
    - betas_coef (list): [Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j, gm, rho, kp, Rf]

    Returns:
    - dict: {'theta_GLB_TS2': ..., 'GLB_TS2': ..., 'theta_GLB_TS3': ..., 'GLB_TS3': ...}
      where GLB values are annualized lower bounds on expected excess returns.
    """
    params_opt = (er, ki, dKi, Q)
    result = {}

    for taylor_order, suffix in [(3, 'TS3'), (2, 'TS2')]:
        try:
            bnds_t = get_glb_theta_boundary(betas_coef, taylor_order=taylor_order)
            Minval1, Maxval1 = bnds_t[0], bnds_t[1]
            bnds = [(Minval1, Maxval1)]
            theta0 = [0.1 * Maxval1] if 0.1 * Maxval1 > Minval1 else [(Minval1 + Maxval1) / 2]

            res = minimize(_optimize_glb_theta, theta0, args=[params_opt],
                           method='SLSQP',
                           bounds=bnds,
                           options={'disp': False, 'ftol': 1e-20, 'maxiter': 100})
            result[f'theta_GLB_{suffix}'] = res.x[0]
            result[f'GLB_{suffix}'] = -res.fun / mat
        except (ValueError, RuntimeError, TypeError):
            result[f'theta_GLB_{suffix}'] = np.nan
            result[f'GLB_{suffix}'] = np.nan

    return result
