# -*- coding: utf-8 -*-
"""
Created on Tue Apr  7 13:24:11 2020

@author: dingq

## this script is largely learnt from @hugou2046 on www.joinquant.com, with some modifications
"""

import pickle
import json

import numpy as np
import pandas as pd
import tushare as ts

from tqdm import *
from jqdatasdk import *
## to use this script, you will need to register an account with www.joinquant.com
auth(#your id,# your password)
from datetime import timedelta
from scipy.interpolate import interp1d

import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.dates as mdate

import os
os.chdir(r'...\replication package')


%matplotlib inline

# to process, you will need to register an account with the website: tushare.pro
pro = ts.pro_api(#your tushare account token)

class CH_VIX(object):
    
    def __init__(self,symbol,start_date,end_date,read=False):
        
        self.symbol = symbol
        self.start_date = start_date
        self.end_date = end_date
        
        self.read = read
        self.OptData = pd.DataFrame()
        
        
    def GetCBOE_VIX(self) -> pd.Series:
        
        if self.read:
            
            self.OptData = pd.read_csv('OptData.csv',index_col=0,parse_dates=['date'])
        
        else:
            
            self.PrepareData()
        
        
        start_dt = min(self.OptData['date'])
        end_dt = max(self.OptData['date'])
        
        df_rate = self.GetShiBor(start_dt, end_dt)

        date_index = [] # to save index
        vix_value = []  # to save vix
        skew_value = [] # to save skew

        for trade_date, slice_df in tqdm(self.OptData.groupby('date'), desc='计算中..',leave=True):

            maturity, df_jy, df_cjy = self.filter_contract(slice_df)
            
            # to obtain riskless rate             
            rf_rate_jy = df_rate.loc[trade_date, int(maturity['jy'] * 365)]             
            rf_rate_cjy = df_rate.loc[trade_date, int(maturity['cjy'] * 365)]

            # to calculate forward price
            fp_jy = self.cal_forward_price(maturity['jy'], rf_rate=rf_rate_jy, df=df_jy)
            fp_cjy = self.cal_forward_price(maturity['cjy'], rf_rate=rf_rate_cjy, df=df_cjy)

            # to calculate mid price
            df_mp_jy = self.cal_mid_price(maturity['jy'], df_jy, fp_jy)
            df_mp_cjy = self.cal_mid_price(maturity['cjy'], df_cjy, fp_cjy)

            # to calculate excercise spread
            df_diff_k_jy = self.cal_k_diff(df_jy)
            df_diff_k_cjy = self.cal_k_diff(df_cjy)

            # to calculate VIX
            df_tovix_jy = pd.concat([df_mp_jy, df_diff_k_jy], axis=1).reset_index()
            df_tovix_cjy = pd.concat([df_mp_cjy, df_diff_k_cjy],
                                     axis=1).reset_index()

            nearest_k_jy = self._nearest_k(df_jy, fp_jy)
            nearest_k_cjy = self._nearest_k(df_cjy, fp_cjy)


            vix = self.cal_vix(df_tovix_jy, fp_jy, rf_rate_jy, maturity['jy'],
                          nearest_k_jy, df_tovix_cjy, fp_cjy, rf_rate_cjy,
                          maturity['cjy'], nearest_k_cjy)

            skew = self.cal_skew(df_tovix_jy, fp_jy, rf_rate_jy, maturity['jy'],
                            nearest_k_jy, df_tovix_cjy, fp_cjy, rf_rate_cjy,
                            maturity['cjy'], nearest_k_cjy)

            date_index.append(trade_date)
            vix_value.append(vix)
            skew_value.append(skew)

        data = pd.DataFrame({
            "CH_VIX": vix_value,
            "CH_SKEW": skew_value
        },
                            index=date_index)
        
        data.fillna(method='pad', inplace=True)
        
        data.index = pd.DatetimeIndex(data.index)
        

        return data

    # to obtain the option contract basic information
    def GetOptContractBasicInfo(self) -> pd.DataFrame:

        # CO-call option，PO-put option
        OptionContractBasicInfo = opt.run_query(
            query(opt.OPT_CONTRACT_INFO.list_date,
                  opt.OPT_CONTRACT_INFO.exercise_date,
                  opt.OPT_CONTRACT_INFO.exercise_price,
                  opt.OPT_CONTRACT_INFO.contract_type,
                  opt.OPT_CONTRACT_INFO.code).filter(
                      opt.OPT_CONTRACT_INFO.underlying_symbol == self.symbol,
                opt.OPT_CONTRACT_INFO.last_trade_date >= self.start_date,
            opt.OPT_CONTRACT_INFO.list_date <= self.end_date))

        return OptionContractBasicInfo


    # to obtain option daily price
    @staticmethod
    def GetOptDailyPrice(code_list:list) -> pd.DataFrame:
        
        price_temp = []  # 储存数据

        # to obtain the data for a certain option
        for CODE in tqdm(code_list, desc='DownLoad Daily Price',leave=True):

            q = query(
                opt.OPT_DAILY_PRICE.date, opt.OPT_DAILY_PRICE.close,
                opt.OPT_DAILY_PRICE.code).filter(opt.OPT_DAILY_PRICE.code == CODE)

            price_temp.append(opt.run_query(q))
        
        return pd.concat(price_temp)


    # to prepare data
    def PrepareData(self):
        
        start = self.start_date
        end = self.end_date
        
        # to obtain basic option information
        OptContractBasicInfo = self.GetOptContractBasicInfo()
        # to obtain the code for each option
        code_list = OptContractBasicInfo['code'].unique().tolist()
        # to obetain daily price
        OptDailyPrice = self.GetOptDailyPrice(code_list)
        
        # to aquire the daily data
        date_ = pd.to_datetime(OptDailyPrice['date'])
        OptDailyPrice = OptDailyPrice[(date_>=start)&(date_<=end)]
        
        OptData = pd.merge(OptDailyPrice, OptContractBasicInfo, on='code')

        # to calculate duratoin to maturity
        ## T
        OptData['maturity'] = (pd.to_datetime(OptData['exercise_date']) -
                               pd.to_datetime(OptData['date'])) / timedelta(365)

        use_col = 'date,close,contract_type,exercise_price,maturity'.split(',')

        df_use = OptData[use_col].copy()

        df_use['contract_type'] = df_use['contract_type'].map({
            "CO": "call",
            "PO": "put"
        })
        df_use = df_use.sort_values('date')

        print('data saved..OptData.csv')
        df_use.to_csv('OptData.csv')
        
        self.OptData = df_use

    # fetch data pf SHIBOR from tushare
    @staticmethod
    def GetShiBor(startDate: str, endDate: str) -> pd.DataFrame:
        '''
        startDate,endDate the date format: yyyy-mm-dd
        '''
       
        # the max rows requirement is 2000

        limit = 2000
        dates = [x.strftime('%Y%m%d') for x in get_trade_days(startDate, endDate)]
        n_days = len(dates)

        if n_days > limit:

            n = n_days // limit
            df_list = []
            i = 0
            pos1, pos2 = n * i, n * (i + 1) - 1

            while pos2 < n_days:
                
                df = pro.shibor(
                    start_date=dates[pos1], end_date=dates[pos2]).set_index('date')

                df_list.append(shibor_df)
                i += 1
                pos1, pos2 = n * i, n * (i + 1) - 1

            if pos1 < n_days:

                df = pro.shibor(
                    start_date=dates[pos1], end_date=dates[-1]).set_index('date')

                df_list.append(df)

            shibor_df = pd.concat(df_list, axis=0)

        else:

            shibor_df = pro.shibor(
                start_date=startDate.strftime('%Y%m%d'),
                end_date=endDate.strftime('%Y%m%d')).set_index('date')

        shibor_df.index = pd.DatetimeIndex(shibor_df.index)
        shibor_df.sort_index(inplace=True)
        
        # interpolation
        def _interpld_fun(r):

            y_vals = r.values / 100

            daily_range = np.arange(1, 361)
            periods = [1, 7, 14, 30, 90, 180, 270, 360]

           
            f = interp1d(periods, y_vals, kind='cubic')
            t_ser = pd.Series(data=f(daily_range), index=daily_range)

            return t_ser

        shibor_df = shibor_df.apply(lambda x: _interpld_fun(x), axis=1)

        shibor_df.index = pd.DatetimeIndex(shibor_df.index)

        return shibor_df


    # select the contract of near future of that day (duration> 1 week)
    def filter_contract(self,cur_df: pd.DataFrame):

        # the maturity date of the current contract
        ex_t = cur_df['maturity'].unique()
        # choose the expiration date >=5
        ex_t = ex_t[ex_t >= 5. / 365]
        
        # sort the expiration date
        try:
            jy_dt, cjy_dt = np.sort(ex_t)[:2]
        
        except ValueError:
            
            print(ex_t,np.sort(ex_t)[:2])
        
        maturity_dict = dict(zip(['jy', 'cjy'], [jy_dt, cjy_dt]))

        cur_df = cur_df[cur_df['maturity'].isin([jy_dt, cjy_dt])]

        keep_cols = ['close', 'contract_type', 'exercise_price']

        cur_df_jy = cur_df.query('maturity==@jy_dt')[keep_cols]
        cur_df_cjy = cur_df.query('maturity==@cjy_dt')[keep_cols]

        cur_df_jy = cur_df_jy.pivot_table(
            index='exercise_price', columns='contract_type', values='close')

        cur_df_cjy = cur_df_cjy.pivot_table(
            index='exercise_price', columns='contract_type', values='close')

 
        cur_df_jy = self._check_fields(cur_df_jy)
        cur_df_cjy = self._check_fields(cur_df_cjy)
        
      
        cur_df_jy['diff'] = np.abs(cur_df_jy['call'] - cur_df_jy['put'])
        cur_df_cjy['diff'] = np.abs(cur_df_cjy['call'] - cur_df_cjy['put'])

        return maturity_dict, cur_df_jy, cur_df_cjy


   
    @staticmethod
    def _check_fields(x_df: pd.DataFrame) -> pd.DataFrame:

       
        target_fields = ['call', 'put']

        for col in target_fields:

            if col not in x_df.columns:
                print("%snone" % col)
                df[col] = 0

        return x_df


    
    @staticmethod
    def cal_forward_price(maturity: dict, rf_rate: float,
                          df: pd.DataFrame) -> float:

        
        min_con = df.sort_values('diff').iloc[0]

        
        k_min = min_con.name

        # F = Strike Price + e^RT x (Call Price - Put Price)
        f_price = k_min + np.exp(maturity * rf_rate) * (
            min_con['call'] - min_con['put'])

        return f_price

    def cal_mid_price(self,maturity: dict, df: pd.DataFrame,
                      forward_price: float) -> pd.DataFrame:

        def _cal_mid_fun(x, val: float):
            res = None
            if x['exercise_price'] < val:
                res = x['put']
            elif x['exercise_price'] > val:
                res = x['call']
            else:
                res = (x['put'] + x['call']) / 2
            return res

        m_k = self._nearest_k(df, forward_price)

        ret = pd.DataFrame(index=df.index)
      
        m_p_lst = df.reset_index().apply(lambda x: _cal_mid_fun(x, val=m_k), axis=1)

        ret['mid_p'] = m_p_lst.values

        return ret



    @staticmethod
    def _nearest_k(df: pd.DataFrame, forward_price: float) -> float:

        temp_df = df[df.index <= forward_price]
        if temp_df.empty:

            temp_df = df

        m_k = temp_df.sort_values('diff').index[0]

        return m_k


    @staticmethod
    def cal_k_diff(df: pd.DataFrame) -> pd.DataFrame:

        arr_k = df.index.values
        ret = pd.DataFrame(index=df.index)

        res = []
        res.append(arr_k[1] - arr_k[0])
        res.extend(0.5 * (arr_k[2:] - arr_k[0:-2]))
        res.append(arr_k[-1] - arr_k[-2])
        ret['diff_k'] = res
        return ret


    @staticmethod
    def cal_vix_sub(df: pd.DataFrame, forward_price: float, rf_rate: float,
                    maturity: float, nearest_k: float):

        def _vix_sub_fun(x):
            ret = x['diff_k'] * np.exp(rf_rate * maturity) * x['mid_p'] / np.square(
                x['exercise_price'])
            return ret

        temp_var = df.apply(lambda x: _vix_sub_fun(x), axis=1)

        sigma = 2 * temp_var.sum() / maturity - np.square(forward_price /
                                                          nearest_k - 1) / maturity

        return sigma


    def cal_vix(self,df_jy: pd.DataFrame, forward_price_jy: float, rf_rate_jy: float,
                maturity_jy: float, nearest_k_jy: float, df_cjy: pd.DataFrame,
                forward_price_cjy: float, rf_rate_cjy: float, maturity_cjy: float,
                nearest_k_cjy: float):

        sigma_jy = self.cal_vix_sub(df_jy, forward_price_jy, rf_rate_jy, maturity_jy,
                               nearest_k_jy)

        sigma_cjy = self.cal_vix_sub(df_cjy, forward_price_cjy, rf_rate_cjy,
                                maturity_cjy, nearest_k_cjy)

        w = (maturity_cjy - 30.0 / 365) / (maturity_cjy - maturity_jy)

        to_sqrt = maturity_jy * sigma_jy * w + maturity_cjy * sigma_cjy * (1 - w)

        if to_sqrt >= 0:

            vix = 100 * np.sqrt(to_sqrt * 365.0 / 30)

        else:

            vix = np.nan

        return vix



    @staticmethod
    def cal_epsilon(forward_price:float, nearest_k:float)->tuple:

        e1 = -(1 + np.log(forward_price / nearest_k) - forward_price / nearest_k)

        e2 = 2 * np.log(nearest_k / forward_price) * (
            nearest_k / forward_price - 1) + np.square(
                np.log(nearest_k / forward_price)) * 0.5

        e3 = 3 * np.square(np.log(nearest_k / forward_price)) * (
            np.log(nearest_k / forward_price) / 3 - 1 + forward_price / nearest_k)

        return e1, e2, e3


    def cal_moments_sub(self,df:pd.DataFrame, maturity:float, rf_rate:float, forward_price:float, nearest_k:float)->float:

        e1, e2, e3 = self.cal_epsilon(forward_price, nearest_k)

        temp_p1 = -np.sum(
            df['mid_p'] * df['diff_k'] / np.square(df['exercise_price']))

        p1 = np.exp(maturity * rf_rate) * (temp_p1) + e1

        temp_p2 = np.sum(df['mid_p'] * df['diff_k'] * 2 *
                         (1 - np.log(df['exercise_price'] / forward_price)) /
                         np.square(df['exercise_price']))
        p2 = np.exp(maturity * rf_rate) * (temp_p2) + e2

        temp_p3 = temp_p3 = np.sum(
            df['mid_p'] * df['diff_k'] * 3 *
            (2 * np.log(df['exercise_price'] / forward_price) -
             np.square(np.log(df['exercise_price'] / forward_price))) /
            np.square(df['exercise_price']))

        p3 = np.exp(maturity * rf_rate) * (temp_p3) + e3

        s = (p3 - 3 * p1 * p2 + 2 * p1**3) / (p2 - p1**2)**(3 / 2)

        return s    
    

#510050 is SSE 50
VIX = CH_VIX('510050.XSHG','2019-12-26','2020-05-31')
VIX_ser = VIX.GetCBOE_VIX()

VIX_ser.to_csv(r'CH_VIX.csv')

VIX_0 = VIX.cal_vix_sub()
