# -*- coding: utf-8 -*-
"""
MIT License

Copyright (c) 2008-2024 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.
_________________________________________________________________________________



Data required:
1. Zero-cd rates
-- one can use zero cd rates from OptionMetrics, or any other provider, e.g., short-term t-bill rates from FRED

2. Surface data from OptionMetrics or any other provider
-- see attached files for 1996 for all SP500 components and SP itself (ID = 99991) 
    for formats and names of the columns
-- note that the ID gives the PERMNO of the company and is based on CRSP data

Output: csv file with all the required option-based measures (the file with specified measures for 1996 is attached)

Notes:
-- the code is written in a very simple form for a purpose -- this way it is easier to debug and adopt;
Please write me with any comments on vilkov@vilkov.net

"""
# %% Load Variables and Dataframes
# import packages and variables
import pandas as pd
import numpy as np
from tqdm import *
from multiprocessing import cpu_count, Pool
from qmoms.examples import load_data
from qmoms import default_params, default_surf_dtype_mapping, default_rate_dtype_mapping, default_date_format
from qmoms import filter_options, get_rate_for_maturity
from qmoms import compute_glb_betas, prepare_integration_inputs_surf

import sympy
from sympy import Symbol
from sympy.solvers.inequalities import reduce_rational_inequalities

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

def get_glb_theta_boundary(betas_coef, taylor_order):
    Bmj, Bm2j, Bj2j, Bmjj, Bm3j, Bj3j, Bm2jj, Bmj2j, gm, rho, kp, Rf = betas_coef

    # define symbol for symbolic computation in sympy
    th = Symbol('th', real=True)

    # define equation
    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

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

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

    if len(eq_res) == 4:  # for the case where we have 2 intervals
        Minval1, Maxval1, Minval2, Maxval2 = eq_res
        Minval1, Maxval1, Minval2, Maxval2 = (float(Minval1), float(Maxval1),
                                              float(Minval2), float(Maxval2))

        res = [Minval1, Maxval1, Minval2, Maxval2]

    elif len(eq_res) == 2:  # for the case where only 1 interval
        Minval1, Maxval1 = eq_res
        Minval1, Maxval1 = float(Minval1), float(Maxval1)

        Minval2, Maxval2 = np.nan, np.nan
        res = [Minval1, Maxval1, Minval2, Maxval2]

    return res




# %% ALL FUNCTIONS THAT ARE USED IN COMPUTATIONS ARE IMPORTED FROM fn_defs_qmoments
# *****************************************************************************
# *****************************************************************************
if __name__ == '__main__':

    df_surf, df_rate, df_return = load_data(surf_dtype_mapping=default_surf_dtype_mapping,
                                            rate_dtype_mapping=default_rate_dtype_mapping,
                                            surf_date_format=default_date_format,
                                            rate_date_format=default_date_format,
                                            return_date_format=default_date_format,
                                            rate_factor=0.01)

    #######################################################################
    # 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")
    #######################################################################

    print('started betas')
    beta_df = compute_glb_betas(df_return, parallel=False, cpuused=CPUUsed)
    print(beta_df)

    #######################################################################
    # 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

    inputs_surf = prepare_integration_inputs_surf(mnes=mnes, vol=iv, days=days, rate=rate, params=default_params)
    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']
    mat = inputs_surf['mat']
    er = inputs_surf['er']

    # params_ind = [er, ki, dKi, Q, mat, betas_coef]

    betas_sel = ['Bmj', 'Bm2j', 'Bj2j', 'Bmjj', 'Bm3j', 'Bj3j', 'Bm2jj', 'Bmj2j']
    gm, rho, kp = 2, 1, 4

    betas_coef = list(beta_df.loc[beta_df.id == '12490', betas_sel].loc[0].values) + [gm, rho, kp, er]
    taylor_order = 2


    bb = get_glb_theta_boundary(betas_coef, taylor_order=taylor_order)
    print(bb)


    def Optimize_LB_theta(theta, params_opt):
        def compute_other_mom(er, ki, dKi, Q, theta):
            m_temp = er ** (theta + 1) + theta * (theta + 1) * er * np.sum(np.power(ki, theta - 1) * Q * dKi)
            return m_temp

        theta = theta[0]
        params_tup = params_opt[0]
        er, ki, dKi, Q = params_tup

        nume = compute_other_mom(er, ki, dKi, Q, theta)
        theta2 = theta - 1
        denom = compute_other_mom(er, ki, dKi, Q, theta2)

        obj = nume / denom - er

        return -obj


    def optimization_wrapper(params_ind):
        er, ki, dKi, Q, mat, betas_coef = params_ind

        params_opt = er, ki, dKi, Q
        LB_UB_Dic = {}

        ##############################################################
        ''' Simple return bounds for 3rd order Taylor expansion'''
        # get theta associated with maixmium lower bound based on betas expected return
        bnds_t = get_glb_theta_boundary(betas_coef, taylor_order=3)
        Minval1, Maxval1, Minval2, Maxval2 = bnds_t
        bnds = [(Minval1, Maxval1)]
        theta0 = [0.1 * Maxval1] if 0.1 * Maxval1 > Minval1 else [(Minval1 + Maxval1) / 2]

        try:
            res1 = minimize(Optimize_LB_theta, theta0, args=[params_opt],
                            method='SLSQP',  # L-BFGS-B, SLSQP
                            bounds=bnds,
                            options={'disp': False, 'ftol': 1e-20, 'maxiter': 100})

            opt_theta = res1.x[0]
            Eri = -res1.fun  # negate it because the function is negated for maximization
        except:
            opt_theta = np.nan
            Eri = np.nan

        LB_UB_Dic['theta_GLB_TS3'] = opt_theta
        LB_UB_Dic['GLB_TS3'] = Eri / mat

        ##############################################################
        ''' Simple return bounds for 2nd order Taylor expansion'''
        # get theta associated with maixmium lower bound based on betas expected return
        bnds_t = get_glb_theta_boundary(betas_coef, taylor_order=2)
        Minval1, Maxval1, Minval2, Maxval2 = bnds_t
        bnds = [(Minval1, Maxval1)]
        theta0 = [0.1 * Maxval1] if 0.1 * Maxval1 > Minval1 else [(Minval1 + Maxval1) / 2]
        try:
            res1 = minimize(Optimize_LB_theta, theta0, args=[params_opt],
                            method='SLSQP',  # L-BFGS-B, SLSQP
                            bounds=bnds,
                            options={'disp': False, 'ftol': 1e-20, 'maxiter': 100})

            opt_theta = res1.x[0]
            Eri = -res1.fun  # negate it because the function is negated for maximization
        except:
            opt_theta = np.nan
            Eri = np.nan

        LB_UB_Dic['theta_GLB_TS2'] = opt_theta
        LB_UB_Dic['GLB_TS2'] = Eri / mat

        return LB_UB_Dic





