"""Wrangle DEX market data to be better suitable for algorithmic trading.

- Wrangle is a process where we "massage" incoming price/liquidity data for the isseus we may have encountered during the data collection

- Common DEX data issues are absurd price high/low spikes due to MEV trades

- We also have some open/close values that are "broken" in a sense that they do not reflect the market price you would be able to trade,
  again likely due to MEV

- See :py:func:`fix_dex_price_data` for fixing

"""
import logging
import datetime
from itertools import islice

import pandas as pd
from IPython.testing.plugin.pytest_ipdoctest import ipdoctest_namespace
from pandas.core.groupby import DataFrameGroupBy
import numpy as np

from .time import naive_utcnow
from .forward_fill import forward_fill as _forward_fill
from ..pair import PandasPairUniverse

logger = logging.getLogger(__name__)


#: Floating point danger zone for price values
DEFAULT_MIN_MAX_RANGE = (0.00000001, 1_000_000.0)


def fix_bad_wicks(
    df: pd.DataFrame,
    threshold=(0.1, 1.9),
    too_slow_threshold=15,
    verbose=False,
    bad_open_close_threshold: float | None=3.0,
) -> pd.DataFrame:
    """Correct out bad high/low values in OHLC data.

    Applicable for both :term:`OHLCV` price feeds and liquidity feeds.

    On :term:`Uniswap` v2 and compatibles, Bad wicks are caused by e.g. very large flash loan, oracle price manipulation attacks,
    and misbheaving bots.

    This function removes bad high/low values and sets them to open/close if they seem to be wildly out of sample.

    :param threshold:
        How many pct % wicks are allowed through.

        Tuple (low threshold, high threshold) relative to close.

        Default to 50%. A high wick cannot be more than 50% of close.

    :param too_slow_threshold:
        Complain if this takes too long

    :param bad_open_close_threshold:
        How many X open must be above the high to be considered a broken data point.

        The open price will be replaced with high price.

        Do not set for liquidity processing.


    :param verbose:
        Make some debug logging when using the function for manual data diagnostics.

    """

    start = naive_utcnow()

    if len(df) == 0:
        return df

    # Optimised with np.where()
    # https://stackoverflow.com/a/65729035/315168
    if threshold is not None:
        df["high"] = np.where(df["high"] > df["close"] * threshold[1], df["close"], df["high"])
        df["low"] = np.where(df["low"] < df["close"] * threshold[0], df["close"], df["low"])

    # For manual diagnostics tracking down bad trading pair data
    if verbose and bad_open_close_threshold:
        bad_opens = df[df["open"] > df["high"] * bad_open_close_threshold]
        for idx, row in islice(bad_opens.iterrows(), 10):
            logger.warning(
                "Pair id %d, timestamp: %s, open: %s, high: %s, buy volume: %s sell volume: %s, volume: %s",
                row.pair_id,
                row.timestamp,
                row.open,
                row.high,
                row.get("buy_volume"),
                row.get("sell_volume"),
                row.get("volume"),
            )
        logger.warn("Total %d bad open price entries detected", len(bad_opens))

    # Issues in open price values with data point - open cannot be higher than high.
    # Not strickly "wicks" but we fix all data while we are at it.
    if bad_open_close_threshold:
        df["open"] = np.where(df["open"] > df["high"] * bad_open_close_threshold, df["high"], df["open"])
        df["close"] = np.where(df["close"] > df["high"] * bad_open_close_threshold, df["high"], df["close"])

    duration = naive_utcnow() - start

    if duration > datetime.timedelta(seconds=too_slow_threshold):
        logger.warning("Very slow fix_bad_wicks(): %s", duration)

    # The following code chokes
    # mask = (df["high"] > df["close"] * (1+threshold)) | (df["low"] < df["close"] * threshold)
    #df.loc[mask, "high"] = df["close"]
    #df.loc[mask, "low"] = df["close"]
    #df.loc[mask, "wick_filtered"] = True
    return df


def filter_bad_wicks(df: pd.DataFrame, threshold=(0.1, 1.9)) -> pd.DataFrame:
    """Mark the bad wicks.

    On :term:`Uniswap` v2 and compatibles, Bad wicks are caused by e.g. very large flash loan, oracle price manipulation attacks,
    and misbheaving bots.

    This function removes bad high/low values and sets them to open/close if they seem to be wildly out of sample.

    :param threshold:
        How many pct % wicks are allowed through as (low, high) tuple.

        This is a tuple (low threshold, high threshold).

        If low < close * threshold[0] ignore the value.

        If high > close * threshold[0] ignore the value.

    """

    df_matches = df.loc[
        (df["high"] > df["close"] * threshold[1]) | (df["low"] < df["close"] * threshold[0])
    ]

    return df_matches


def remove_zero_candles(
    df: pd.DataFrame,
) -> pd.DataFrame:
    """Remove any candle that has a zero value for OHLC

    :param df: Dataframe that may contain zero candles
    :return: pd.Dataframe
    """
    if len(df) > 0:
        filtered_df = df[(df['open'] != 0) & (df['high'] != 0) & (df['low'] != 0) & (df['close'] != 0)]
        return filtered_df
    return df


def remove_min_max_price(
    df: pd.DataFrame,
    min_max_price: tuple = DEFAULT_MIN_MAX_RANGE,
) -> pd.DataFrame:
    """Remove candles where open value is outside the floating point range detector.

    :param df:
        Dataframe that may contain open or close values for too funny precision.

    :param min_max_price:
        Min and max price allowed for open/close before dropping.

    :return: pd.Dataframe
    """
    if len(df) > 0:
        mask = (df["open"] < min_max_price[0]) | (df["close"] < min_max_price[0]) | \
               (df["open"] > min_max_price[1]) | (df["close"] > min_max_price[1])

        filtered_df = df[~mask]
        return filtered_df
    return df


def _replace_for_groups(group, replacement_dict):
    group_name = group.name
    if group_name in replacement_dict:
        for col, value in replacement_dict[group_name].items():
            group[col] = value
    return group


def fix_prices_in_between_time_frames(
    dfgb: pd.DataFrame | DataFrameGroupBy,
    fix_inbetween_threshold: tuple | None = (-0.99, 5.0),
    pair_id_column="pair_id",
):
    """Fix MEV bots breaking open/price.

    TODO: This needs to be later fixed at the data collection level, MEV transactions masked out.

    Example daily candle with a broken open price:

    .. code-block:: text

        COMP-WETH on Uniswap v2 on Ethereum

        2023-12-10 23:00:00	2023-12-10 23:00:00	56.253001	56.253001	55.948422	55.948422	348.915474	2373.049643	0.0	3.0	0.000000	348.915474	1015	18758972	18759006	NaN
        2023-12-11 00:00:00	2023-12-11 00:00:00	0.363204	55.181369	0.354456	55.181369	180421.431179	2357.352148	1.0	6.0	90162.684736	90258.746443	1015	18759541	18759541	NaN
        2023-12-11 02:00:00	2023-12-11 02:00:00	50.059992	50.612159	50.059992	50.612159	2859.011552	2241.581295	1.0	1.0	628.775186	2230.236366	1015	18759912	18760129	NaN

    `Underlying MEV TX causing the issue <https://etherscan.io/tx/0x1106418384414ed56cd7cbb9fedc66a02d39b663d580abc618f2d387348354ab>`__, ChatGPT formatted:

    .. code-block:: text

        +--------------------------------------+-------------------------------------+-------------------+-------------+-------------------+-------------------------------------+-------------------+-------------+-------------+
        |               Action                 |          Swap Amount (Token 1)      |   Value (Token 1)  |  Token 1    |    For (Token 2)   |          Swap Amount (Token 2)      |   Value (Token 2)  |  Token 2    |  Platform   |
        +--------------------------------------+-------------------------------------+-------------------+-------------+-------------------+-------------------------------------+-------------------+-------------+-------------+
        | Aggregated Swap of 3 Tokens           |                                     |                   |             |                   |                                     |                   |             |             |
        | Swap                                 | 20,096.81048332570788104            | $956,407.21       | COMP        | ETH               | 38.247439949510380362              | $101,436.63       | ETH         | Uniswap V2  |
        | Swap                                 | 65.583941528193262325               | $3,121.14         | COMP        | ETH               | 0.010044352430924858               | $26.64            | ETH         | Uniswap V2  |
        | Swap                                 | 0.010044352430924858                | $26.64            | ETH         | DAI               | 23.587216376218824815              | $23.58            | DAI         | Uniswap V2  |
        | Swap                                 | 65.780693352777842112               | $3,130.50         | COMP        | ETH               | 0.010014399362431825               | $26.56            | ETH         | Uniswap V2  |
        | Swap                                 | 0.010014399362431825                | $26.56            | ETH         | DAI               | 23.516703895990702452              | $23.51            | DAI         | Uniswap V2  |
        | Swap                                 | 65.978035432836175638               | $3,139.89         | COMP        | ETH               | 0.009984535616403163               | $26.48            | ETH         | Uniswap V2  |
        | Swap                                 | 0.009984535616403163                | $26.48            | ETH         | DAI               | 23.446402724807108015              | $23.44            | DAI         | Uniswap V2  |
        | Swap                                 | 66.175969539134684165               | $3,149.31         | COMP        | ETH               | 0.009954760926472085               | $26.40            | ETH         | Uniswap V2  |
        | Swap                                 | 0.009954760926472085                | $26.40            | ETH         | DAI               | 23.376312226345838338              | $23.37            | DAI         | Uniswap V2  |
        | Swap                                 | 5.013744524035853058                | $238.60           | COMP        | ETH               | 0.000751785040201655               | $1.99             | ETH         | Uniswap V2  |
        | Swap                                 | 0.000751785040201655                | $1.99             | ETH         | DAI               | 1.765375649239336674               | $1.76             | DAI         | Uniswap V2  |
        | Swap                                 | 38.247439949510380362               | $101,436.63       | ETH         | COMP              | 20,353.762314725950969279          | $968,635.55       | COMP        | Uniswap V2  |
        | Swap                                 | 0.149285130679667947                | $395.92           | ETH         | COMP              | 4,200                              | $199,878.00       | COMP        | Sushiswap   |
        +--------------------------------------+-------------------------------------+-------------------+-------------+-------------------+-------------------------------------+-------------------+-------------+-------------+

    :param dfgb:
        Assume grouped by pair_id and MultiLevel index (pair_id, timestamp).

    """
    assert isinstance(dfgb, DataFrameGroupBy), f"Currently only implemented for DataFrameGroupBy"

    replacements = {}

    for pair_id, price_df in dfgb:
        for column in ("open", "close",):
            price_series = price_df[column]
            fixed_price_series = examine_price_between_time_anomalies(
                price_series,
                low_diff=fix_inbetween_threshold[0],
                high_diff=fix_inbetween_threshold[1],
                heal=True,
            )
            if fixed_price_series is not None:
                replacements[pair_id] = replacements.get(pair_id, {})
                replacements[pair_id][column] = fixed_price_series

    healed = dfgb.apply(lambda x: _replace_for_groups(x, replacements))
    healed = healed.set_index("timestamp", drop=False)
    return healed.groupby(pair_id_column)


def fix_dex_price_data(
    df: pd.DataFrame | DataFrameGroupBy,
    freq: pd.DateOffset | str | None = None,
    forward_fill: bool = True,
    bad_open_close_threshold: float | None = 3.0,
    fix_wick_threshold: tuple | None = (0.1, 1.9),
    fix_inbetween_threshold: tuple | None = (-0.99, 5.0),
    min_max_price: tuple | None = DEFAULT_MIN_MAX_RANGE,
    remove_candles_with_zero: bool = True,
    pair_id_column="pair_id",
) -> pd.DataFrame:
    """Wrangle DEX price data for all known issues.

    - Fix broken open/high/low/close value so that they are less likely to cause problems for algorithms

    - Wrangle is a process where we massage incoming price/liquidity data for the isseus we may have encountered during the data collection

    - Common DEX data issues are absurd price high/low spikes due to MEV trades

    - We also have some open/close values that are "broken" in a sense that they do not reflect the market price you would be able to trade,
      again likely due to MEV

    Example:

    .. code-block:: python

          # After we know pair ids that fill the liquidity criteria,
          # we can build OHLCV dataset for these pairs
          print(f"Downloading/opening OHLCV dataset {time_bucket}")
          price_df = client.fetch_all_candles(time_bucket).to_pandas()
          print(f"Filtering out {len(top_liquid_pair_ids)} pairs")
          price_df = price_df.loc[price_df.pair_id.isin(top_liquid_pair_ids)]

          print("Wrangling DEX price data")
          price_df = price_df.set_index("timestamp", drop=False).groupby("pair_id")
          price_df = fix_dex_price_data(
              price_df,
              freq=time_bucket.to_frequency(),
              forward_fill=True,
          )

          print(f"Retrofitting OHLCV columns for human readability")
          price_df = price_df.obj
          price_df["pair_id"] = price_df.index.get_level_values(0)
          price_df["ticker"] = price_df.apply(lambda row: make_full_ticker(pair_metadata[row.pair_id]), axis=1)
          price_df["link"] = price_df.apply(lambda row: make_link(pair_metadata[row.pair_id]), axis=1)

          # Export data, make sure we got columns in an order we want
          print(f"Writing OHLCV CSV")
          del price_df["timestamp"]
          del price_df["pair_id"]
          price_df = price_df.reset_index()
          column_order = ('ticker', 'timestamp', 'open', 'high', 'low', 'close', 'volume', 'link', 'pair_id',)
          price_df = price_df.reindex(columns=column_order)  # Sort columns in a specific order
          price_df.to_csv(
            price_output_fname,
          )
          print(f"Wrote {price_output_fname}, {price_output_fname.stat().st_size:,} bytes")

    :param df:
        Price dataframe with OHLCV data.

        May contain columns named open, close, high, low, volume and timestamp.

        For multipair data this must be `DataFrameGroupBy`.

    :param freq:
        The incoming Pandas frequency of the data, e.g. "d" for daily.

        If the incoming data frequency and `freq` parameter do not match, the data is resampled o the given frequency.

    :param fix_wick_threshold:
        Apply abnormal high/low wick fix filter.

        Percent value of maximum allowed high/low wick relative to close.
        By default fix values where low is 90% lower than close and high is 90% higher than close.

        See :py:func:`~tradingstrategy.utils.groupeduniverse.fix_bad_wicks` for more information.

    :param bad_open_close_threshold:
        See :py:func:`fix_bad_wicks`.

    :param pair_id_column:
        The pair/reserve id column name in the dataframe.

    :param remove_candles_with_zero:
        Remove candles with zero values for OHLC.

        To deal with abnormal data.

    :param min_max_price:
        Remove candles where open value is outside the floating point range detector.

        See :py:func:`remove_min_max_price`.

    :param forward_fill:
        Forward-will gaps in the data.

        Forward-filling data will delete any unknown columns,
        see :py:func:`tradingstrategy.utils.forward_fill.forward_fill` details.

    :return:
        Fixed data frame.

        If forward fill is used, all other columns outside OHLCV are dropped.
    """

    assert isinstance(df, (pd.DataFrame, DataFrameGroupBy)), f"Got: {df.__class__}"

    if isinstance(df, DataFrameGroupBy):
        raw_df = df.obj
    else:
        raw_df = df

    if min_max_price:
        logger.info("Removing open/close prices outside our min-max tolerance range: %s", min_max_price)
        raw_df = remove_min_max_price(
            raw_df,
            min_max_price=min_max_price,
        )
        raw_df = raw_df.copy()  # Need to mutate in fix_bad_wicks(), cannot be view

    if fix_wick_threshold or bad_open_close_threshold:
        logger.info("Fixing bad wicks, fix_wick_threshold:%s, bad_open_close_threshold: %s", fix_wick_threshold, bad_open_close_threshold)
        raw_df = fix_bad_wicks(
            raw_df,
            fix_wick_threshold,
            bad_open_close_threshold=bad_open_close_threshold,
        )

    if remove_candles_with_zero:
        logger.info("Fixing zero volume candles")
        raw_df = remove_zero_candles(raw_df)

    # For the further cleanup, need to regroup the DataFrame
    if isinstance(df, DataFrameGroupBy):

        logger.info("Regrouping fixed price data")

        # Need to group here
        # TODO: Make this smarter, but how? Read index data in groupby instance?
        regrouped = raw_df.set_index("timestamp", drop=False).groupby(pair_id_column, group_keys=True)

        logger.info("Fixing prices having bad open/close values between timeframes: %s", fix_inbetween_threshold)

        ff_df = fix_prices_in_between_time_frames(
            regrouped,
            fix_inbetween_threshold=fix_inbetween_threshold,
            pair_id_column=pair_id_column,
        )

    else:
        assert not fix_inbetween_threshold, "fix_inbetween_threshold() only works for DataFrameGroupBy input. Set fix_inbetween_threshold == None if you really want to call this"
        ff_df = raw_df

    if forward_fill:
        logger.info("Forward filling price data")
        assert freq, "freq argument must be given if forward_fill=True"
        df = _forward_fill(ff_df, freq)
        return df
    else:
        return raw_df


def examine_price_between_time_anomalies(
    price_series: pd.Series,
    high_diff=5.00,
    low_diff=-0.99,
    column: str = "close",
    heal=False,
) -> pd.DataFrame | pd.Series:
    """Find bad open/close prices where the open price is very different from the value of previous and next day.

    - Must be done pair-by-pair

    TODO: Work in progress.

    :param price_series:
        Incoming OHLCV or single price series.

    :param high_diff:
        A price between days cannot be higher than this multiplier.

    :param low_diff:
        A price between days cannot be lower than this multiplier.

    :param column:
        Column name to check.

        Only relevant if input is DataFrame.

    :param heal:
        IF set return the healed price data witih avg values replacing the outliers.

    :return:
        Diagnostics results as DataFrame.

        If `heal` is set, return healed price data or None if no healing needed.
    """
    assert isinstance(price_series, (pd.Series, pd.DataFrame)), f"Got: {price_series.__class__}"

    if isinstance(price_series, pd.Series):
        df = pd.DataFrame({
            "price": price_series,
        })
    else:
        df = price_series
        df["price"] = df[column]

    df = df.copy()

    # Calculate surrounding price
    df["price_before"] = df['price'].shift(1)
    df["price_after"] = df['price'].shift(-1)
    df['surrounding_avg'] = (df["price_after"] + df["price_before"]) / 2

    # How much the current price is different from the surrounding avg price
    df["price_diff"] = ((df["surrounding_avg"] - df['price']) / df['price'])

    # Create mask for rows where we detect anomalies
    df["high_anomaly"] = df["price_diff"] > high_diff
    df["low_anomaly"] = df["price_diff"] < low_diff
    df["anomaly"] =  df["low_anomaly"] | df["high_anomaly"]

    if heal:
        # Only execute if we have data to heal
        if df["anomaly"].sum() > 0:  # Count of True values of anomalies
            df["fixed_price"] = np.where(df["anomaly"], df["surrounding_avg"], df["price"])
            return df["fixed_price"]
        else:
            return None
    else:
        # Get the rows that are anomalies
        # Can't have value at start and end,
        # drop NaNs
        df = df.dropna(subset=["surrounding_avg"])

        return df.loc[df["anomaly"]]


def examine_anomalies(
    pair_universe: PandasPairUniverse | None,
    price_df: pd.DataFrame,
    printer=lambda x: print(x),
    max_print=2,
    pair_id_column: str | None="pair_id",
    open_close_max_diff=5.00,
    open_close_min_diff=-0.99,
    between_high_diff: float | None=5.00,
    between_low_diff: float | None=-0.99,
):
    """Check the price dataframe for data issues.

    - Print out to consoles bad rows in the OHLCV candle price data

    Perform

    - Open/close diff check

    - In between timeframes diff check

    TODO: This is a work in progress helper.

    See also:

    - :py:func:`examine_price_between_time_anomalies`

    :param price_df:
        OHLCV data for multiple trading pairs.

        Can be grouped by `pair_id_column`.

    :param pair_id_column:
        Fix column identifies the pair name in the data.

    :param max_print:
        How many entries print per each anomaly check

    :param open_close_max_diff:
        Abnormal price increase X

    :param open_close_min_diff:
        Abnormal price decrease X
    """

    assert isinstance(price_df, pd.DataFrame)

    issues_found = False

    # Find zero prices
    zero_prices = price_df.loc[price_df["open"] <= 0]

    if pair_id_column:
        zero_prices = zero_prices.drop_duplicates(subset=pair_id_column, keep='first')

    if len(zero_prices) > 0:
        printer(f"Total {len(zero_prices)} rows with zero price entry gap")

    for zero_price_entry in zero_prices.iloc[0:max_print].iterrows():
        printer(f"Found zero price entry {zero_price_entry}")
        issues_found = True

    # Find abnormal price jumps within intraday open->close is very different
    open_close_mask = ((price_df["close"] - price_df["open"]) / price_df["open"]) >= open_close_max_diff
    open_close_mask = open_close_mask | (((price_df["close"] - price_df["open"]) / price_df["open"]) <= open_close_min_diff)

    open_close_gap = price_df.loc[open_close_mask]

    if len(open_close_gap) > 0:
        printer(f"Total {len(open_close_gap)} rows with open/close price gap")

    if pair_id_column:
        open_close_gap = open_close_gap.drop_duplicates(subset=pair_id_column, keep='first')

    if len(open_close_gap) > 0:
        for idx, open_close_entry in open_close_gap.iloc[0:max_print].iterrows():
            diff = (open_close_entry["close"] - open_close_entry["open"]) / open_close_entry["open"]
            printer(f"Found abnormal open/close price diff {diff} at\n{open_close_entry}")
            issues_found = True
    else:
        printer("No open/close price extreme diff anomalies found")


    if between_high_diff and between_low_diff:
        grouped = price_df.groupby(pair_id_column)
        anomalies = []

        between_anomaly_count = 0
        for pair, group in grouped:
            for column in ("open", "close"):
                between_anomalies = examine_price_between_time_anomalies(
                    group,
                    high_diff=between_high_diff,
                    low_diff=between_low_diff,
                    column=column,
                )

                between_anomaly_count += len(between_anomalies)

                # Only record last one per pair
                if len(between_anomalies) > 0:
                    row = between_anomalies.iloc[-1]
                    anomalies.append((pair, row))

        if anomalies:
            print(f"Found {between_anomaly_count} price entries that greatly differ from one before and after")
            for entry in anomalies[0:max_print]:
                pair, diagnostics_row = entry
                timestamp = diagnostics_row["timestamp"]
                printer(f"Found abnormal between price on {pair} at {timestamp}:\n{diagnostics_row}")
            issues_found = True
        else:
            printer("No price time frame in-between extreme diff anomalies found")

    if not issues_found:
        printer(f"No data issues found, {len(price_df)} rows analysed")

    return issues_found

