"""
A set of functions for performing backtesting
"""
import numpy as np
import bilbystats as bs
from typing import Literal

from datetime import timedelta
import pandas as pd

# Define your “enum” of allowed string values:
InstrumentType = Literal[
    "equity",
    "commodity",
]

import nasdaqdatalink

SYMBOL_KEY_DICT = {
    "equity": "ticker",
    "commodity": "commodity",
}

TABLE_NAME_DICT = {
    "equity": "DY/SPA",
    "commodity": "DY/FUA",
}

def get_nasdaq_data(
    symbol: str = None,
    start_date: str  = None,
    end_date: str = None,
    security_type: InstrumentType = "equity"
):
    """
    Retrieve NASDAQ data for a given security and date range.

    This function queries the NASDAQ Data Link API for historical data based 
    on the specified ticker symbol(s), date range, and security type. It 
    returns a DataFrame containing the relevant OHLC data if available.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    symbol : str or list of str, optional
        The equity ticker symbol or list of symbols to retrieve (e.g., '000001' 
        or ['000001', '000002']). If None, no symbol filtering is applied.
    start_date : str, optional
        The start date in 'YYYY-MM-DD' format to begin filtering data.
    end_date : str, optional
        The end date in 'YYYY-MM-DD' format to stop filtering data.
    security_type : InstrumentType, optional
        The type of financial instrument to query. Defaults to "equity".

    ---------------------------------------------------------------------------
    OUTPUT:
    data : pandas.DataFrame or None
        A DataFrame containing the OHLC data for the given query parameters. 
        Returns None if no data is found or an error occurs.
    ---------------------------------------------------------------------------
    AUTHOR: Steve Enwright-Ward
    ---------------------------------------------------------------------------
    """
    nasdaqdatalink.ApiConfig.api_key = bs.read_api_key('NASDAQ_DATALINK')

    # Build our query parameters
    query_params = {}

    # Add ticker filter - for China equities, the column might be 'ticker' or 'symbol'
    if symbol:
        symbol_key = SYMBOL_KEY_DICT[security_type]
        query_params[symbol_key] = symbol  # Based on your commodities pattern

    # Add date range filter if provided
    if start_date or end_date:
        date_filter = {}

        if start_date:
            date_filter["gte"] = start_date

        if end_date:
            date_filter["lte"] = end_date

        if date_filter:
            query_params["date"] = date_filter

    try:
        # Get the data using the proper parameter format with paginate=True
        table_name = TABLE_NAME_DICT[security_type]

        data = nasdaqdatalink.get_table(
            table_name,
            paginate=True,
            **query_params,
        )

        # Check if data is empty
        if data.empty:
            print(f"No data found for ticker {symbol} in the specified date range.")
            return None

        data = data.sort_values(by='date', ascending=True).reset_index(drop=True)
        
        # List of columns to exclude
        exclude_cols = ['ticker', 'date']

        # Identify columns to convert
        cols_to_convert = data.columns.difference(exclude_cols)

        # Apply conversion using pd.to_numeric with errors='coerce'
        data[cols_to_convert] = data[cols_to_convert].apply(pd.to_numeric, errors='coerce')

        data.index = data['date']

        return data

    except Exception as e:
        print(f"Error querying equities data: {e}")
        return None
    


def ma(df, col_name, ndays, date_col_name):
    first_date =  min(df[date_col_name])
    last_date = max(df[date_col_name])
    dates = pd.date_range(start=first_date, end=last_date, freq="D")
    #unique_dates = df[date_col_name].unique()

    ma_avg = np.zeros(len(dates))

    # Loop over each row and calculate average over the past 7 calendar days
    for i in range(len(dates)):
        current_date = dates[i]
        window_start = current_date - timedelta(days=ndays)
        window_end = current_date - timedelta(days=1)
        window_df = df[(df[date_col_name] >= window_start) & (df[date_col_name] <= window_end)]
    
        # Compute mean of 'Price' in this window
        if len(window_df) > 0:
            ma_avg[i] = window_df[col_name].mean()
        else:
            ma_avg[i] = np.nan

    av_name = str(ndays) + '_day_average_' + col_name

    data = {
        'date': dates,
        av_name: ma_avg
    }
    
    # Create the DataFrame
    df_out = pd.DataFrame(data)
    df_out.index = dates
    
    return df_out
    
def signal_trade(signal_df, signal_cov, stock_price_df, stock_cov, threshold = 0, lag = 1, trade_type = 'buy'):
    revenue_difference = np.zeros(len(stock_price_df) - lag)
    revenue_ratio = np.zeros(len(stock_price_df) - lag)
    ntrades = 0
    price_data = stock_price_df[stock_cov]
    for i in range(len(stock_price_df) - lag):
        todays_price = price_data.iloc[i]
        date = pd.Timestamp(stock_price_df.index[i])
        todays_signal = signal_df[signal_cov][date]
        
        if trade_type == 'buy':
            if todays_signal > threshold:
                ntrades += 1
                revenue_difference[i] = price_data.iloc[i + lag] - todays_price
                revenue_ratio[i] = price_data.iloc[i + lag]/todays_price
        elif trade_type == 'sell':
            if todays_signal < threshold:
                ntrades += 1
                revenue_difference[i] = todays_price - price_data.iloc[i + lag] 
                revenue_ratio[i] = todays_price/price_data.iloc[i + lag]
                
    revenue_dict = {
        'revenue_difference': revenue_difference,
        'revenue_ratio': revenue_ratio,
    }
    revenue_df = pd.DataFrame(revenue_dict)
    revenue_df.index = stock_price_df.index[:(-lag)]
    
    average_difference_revenue = np.sum(revenue_difference)/ntrades
    total_difference_revenue = np.sum(revenue_difference)
    average_ratio_revenue = np.sum(revenue_ratio)/ntrades - 1
    total_ratio_revenue = ntrades*average_ratio_revenue            
            
    av_revenue_dict = {
        'average_revenue_difference': average_difference_revenue,
        'average_revenue_ratio': average_ratio_revenue,
        'total_revenue_difference': total_difference_revenue,
        'total_revenue_ratio': total_ratio_revenue
    }
    
    av_revenue_df = pd.DataFrame(av_revenue_dict, index=[0])
            
    return revenue_df, av_revenue_df

def signal_trade2(signal_df, signal_cov, stock_price_df, stock_cov, threshold = 0, lag = 1, trade_type = 'buy'):
    #revenue_difference = np.zeros(len(stock_price_df) - lag)
    revenue_difference = []
    revenue_ratio = []
    new_dates = []
    #revenue_ratio = np.zeros(len(stock_price_df) - lag)
    ntrades = 0
    price_data = stock_price_df[stock_cov]
    for i in range(len(stock_price_df) - lag):
        todays_price = price_data.iloc[i]
        date = pd.Timestamp(stock_price_df.index[i])
        todays_signal = signal_df[signal_cov][date]
        
        if trade_type == 'buy':
            if todays_signal > threshold:
                ntrades += 1
                revenue_difference.append(price_data[i + lag] - todays_price)
                revenue_ratio.append(price_data[i + lag]/todays_price)
                new_dates.append(date)
                #revenue_difference[i] = price_data[i + lag] - todays_price
                #revenue_ratio[i] = price_data[i + lag]/todays_price
        elif trade_type == 'sell':
            if todays_signal < threshold:
                ntrades += 1
                revenue_difference[i] = todays_price - price_data[i + lag] 
                revenue_ratio[i] = todays_price/price_data[i + lag]
                
    revenue_difference = np.array(revenue_difference)
    revenue_ratio = np.array(revenue_ratio)

    revenue_dict = {
        'revenue_difference': revenue_difference,
        'revenue_ratio': revenue_ratio,
    }
    revenue_df = pd.DataFrame(revenue_dict)
    revenue_df.index = new_dates
    #revenue_df.index = stock_price_df.index[:(-lag)]
    
    average_difference_revenue = np.sum(revenue_difference)/ntrades
    total_difference_revenue = np.sum(revenue_difference)
    average_ratio_revenue = np.sum(revenue_ratio)/ntrades - 1
    total_ratio_revenue = ntrades*average_ratio_revenue            
            
    av_revenue_dict = {
        'average_revenue_difference': average_difference_revenue,
        'average_revenue_ratio': average_ratio_revenue,
        'total_revenue_difference': total_difference_revenue,
        'total_revenue_ratio': total_ratio_revenue
    }
    
    av_revenue_df = pd.DataFrame(av_revenue_dict, index=[0])
            
    return revenue_df, av_revenue_df