"""
dependence on stochvolmodels for model calibration
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import qis
from typing import Dict, Tuple, Optional
from numba.typed import List
from enum import Enum

# analytics
from stochvolmodels.data.option_chain import OptionChain

# chain
from option_chain_analytics import OptionsDataDFs, create_chain_from_from_options_dfs
from option_chain_analytics.option_chain import SliceColumn, SlicesChain


pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 500)
pd.set_option('display.width', 1000)


def generate_vol_chain_np(chain: SlicesChain,
                          value_time: pd.Timestamp,
                          days_map: Dict[str, int] = {'1w': 7, '1m': 21},
                          delta_bounds: Tuple[Optional[float], Optional[float]] = (-0.1, 0.1),
                          is_filtered: bool = True
                          ) -> OptionChain:
    # generate table snapshot of vol data: index= days, columns = deltas

    ttms, future_prices, discfactors = List(), List(), List()
    optiontypes_ttms, strikes_ttms = List(), List()
    bid_ivs, ask_ivs = List(), List()
    bid_prices, ask_prices = List(), List()
    slice_ids = []
    for label, day in days_map.items():
        next_date = value_time + pd.DateOffset(days=day)  # if overlapping next date will be last avilable maturity
        slice_date = chain.get_next_slice_after_date(mat_date=next_date)
        slice_t = chain.expiry_slices[slice_date]
        df = slice_t.get_joint_slice(delta_bounds=delta_bounds, is_filtered=is_filtered)

        slice_ids.append(f"{label}: {slice_t.expiry_id}")
        ttms.append(slice_t.get_ttm())
        future_prices.append(slice_t.get_future_price())
        discfactors.append(1.0)
        strikes_ttms.append(df.index.to_numpy())
        optiontypes_ttms.append(df[SliceColumn.OPTION_TYPE].to_numpy(dtype=str))
        bid_ivs.append(df[SliceColumn.BID_IV].to_numpy())
        ask_ivs.append(df[SliceColumn.ASK_IV].to_numpy())
        bid_prices.append(df[SliceColumn.BID_PRICE].to_numpy())
        ask_prices.append(df[SliceColumn.ASK_PRICE].to_numpy())

    out = OptionChain(ttms=np.array(ttms),
                      forwards=np.array(future_prices),
                      discfactors=np.array(discfactors),
                      ids=np.array(slice_ids),
                      strikes_ttms=strikes_ttms,
                      optiontypes_ttms=optiontypes_ttms,
                      bid_ivs=bid_ivs,
                      ask_ivs=ask_ivs,
                      bid_prices=bid_prices,
                      ask_prices=ask_prices)
    return out


def load_option_chain(options_data_dfs: OptionsDataDFs,
                      value_time: pd.Timestamp = pd.Timestamp('2023-02-06 08:00:00+00:00'),
                      days_map: Dict[str, int] = {'1w': 7, '1m': 21},
                      delta_bounds: Tuple[Optional[float], Optional[float]] = (-0.1, 0.1),
                      is_filtered: bool = True
                      ) -> Optional[OptionChain]:
    chain = create_chain_from_from_options_dfs(options_data_dfs=options_data_dfs, value_time=value_time)
    if chain is not None:
        option_chain = generate_vol_chain_np(chain=chain,
                                             value_time=value_time,
                                             days_map=days_map,
                                             delta_bounds=delta_bounds,
                                             is_filtered=is_filtered)
    else:
        option_chain = None

    return option_chain


def sample_option_chain_at_times(options_data_dfs: OptionsDataDFs,
                                 time_period: qis.TimePeriod,
                                 freq: str = 'W-FRI',
                                 days_map: Dict[str, int] = {'1w': 7, '1m': 21},
                                 hour_offset: int = 8,
                                 delta_bounds: Tuple[Optional[float], Optional[float]] = (-0.1, 0.1),
                                 ) -> Dict[pd.Timestamp, OptionChain]:
    value_times = qis.generate_dates_schedule(time_period=time_period,
                                              freq=freq,
                                              hour_offset=hour_offset)
    option_chains = {}
    for value_time in value_times:
        option_chains[value_time] = load_option_chain(options_data_dfs=options_data_dfs,
                                                      value_time=value_time,
                                                      days_map=days_map,
                                                      delta_bounds=delta_bounds,
                                                      is_filtered=True)
    return option_chains


class UnitTests(Enum):
    PRINT_CHAIN_DATA = 1
    GENERATE_VOL_CHAIN_NP = 2
    SAMPLE_CHAIN_AT_TIMES = 3


def run_unit_test(unit_test: UnitTests):

    import local_path as local_path

    ticker = 'BTC'  # BTC, ETH
    value_time = pd.Timestamp('2021-10-21 08:00:00+00:00')
    value_time = pd.Timestamp('2023-10-06 08:00:00+00:00')

    from option_chain_analytics.ts_loaders import ts_data_loader_wrapper
    options_data_dfs = OptionsDataDFs(**ts_data_loader_wrapper(ticker=ticker))
    options_data_dfs.get_start_end_date().print()
    chain = create_chain_from_from_options_dfs(options_data_dfs=options_data_dfs, value_time=value_time)

    if unit_test == UnitTests.PRINT_CHAIN_DATA:
        for expiry, eslice in chain.expiry_slices.items():
            eslice.print()

    elif unit_test == UnitTests.GENERATE_VOL_CHAIN_NP:
        option_chain = generate_vol_chain_np(chain=chain,
                                             value_time=value_time,
                                             days_map={'1w': 7},
                                             delta_bounds=(-0.1, 0.1),
                                             is_filtered=True)
        option_chain.print()
        skews = option_chain.get_chain_skews(delta=0.35)
        print(skews)

    elif unit_test == UnitTests.SAMPLE_CHAIN_AT_TIMES:
        time_period = qis.TimePeriod('01Jan2023', '31Jan2023', tz='UTC')
        option_chains = sample_option_chain_at_times(options_data_dfs=options_data_dfs,
                                                     time_period=time_period,
                                                     freq='W-FRI',
                                                     hour_offset=9
                                                     )
        for key, chain in option_chains.items():
            print(f"{key}")
            print(chain)

    plt.show()


if __name__ == '__main__':

    unit_test = UnitTests.SAMPLE_CHAIN_AT_TIMES

    is_run_all_tests = False
    if is_run_all_tests:
        for unit_test in UnitTests:
            run_unit_test(unit_test=unit_test)
    else:
        run_unit_test(unit_test=unit_test)
