from numba import jit
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt


@jit(nopython=True)
def process_data(close, labels, metalabels, stop, take, markup, forward, backward):
    last_deal = 2
    last_price = 0.0
    report = [0.0]
    chart = [0.0]
    line_f = 0
    line_b = 0

    for i in range(len(close)):
        line_f = len(report) if i <= forward else line_f
        line_b = len(report) if i <= backward else line_b

        pred = labels[i]
        pr = close[i]
        pred_meta = metalabels[i]  # 1 = allow trades

        if last_deal == 2 and pred_meta == 1:
            last_price = pr
            last_deal = 0 if pred < 0.5 else 1
            continue

        if last_deal == 0:
            if (-markup + (pr - last_price) >= take) or (
                -markup + (last_price - pr) >= stop
            ):
                last_deal = 2
                profit = -markup + (pr - last_price)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + profit)
                continue

        if last_deal == 1:
            if (-markup + (pr - last_price) >= stop) or (
                -markup + (last_price - pr) >= take
            ):
                last_deal = 2
                profit = -markup + (last_price - pr)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + (pr - last_price))
                continue

        # close deals by signals
        if last_deal == 0 and pred > 0.5:
            last_deal = 2
            profit = -markup + (pr - last_price)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + profit)
            continue

        if last_deal == 1 and pred < 0.5:
            last_deal = 2
            profit = -markup + (last_price - pr)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + (pr - last_price))
            continue

    return np.array(report), np.array(chart), line_f, line_b


@jit(nopython=True)
def process_data_one_direction(
    close, labels, metalabels, stop, take, markup, forward, backward, direction
):
    last_deal = 2
    last_price = 0.0
    report = [0.0]
    chart = [0.0]
    line_f = 0
    line_b = 0

    for i in range(len(close)):
        line_f = len(report) if i <= forward else line_f
        line_b = len(report) if i <= backward else line_b

        pred = labels[i]
        pr = close[i]
        pred_meta = metalabels[i]  # 1 = allow trades

        if last_deal == 2 and pred_meta == 1:
            last_price = pr
            last_deal = 2 if pred < 0.5 else 1
            continue

        if last_deal == 1 and direction == "buy":
            if (-markup + (pr - last_price) >= take) or (
                -markup + (last_price - pr) >= stop
            ):
                last_deal = 2
                profit = -markup + (pr - last_price)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + profit)
                continue

        if last_deal == 1 and direction == "sell":
            if (-markup + (pr - last_price) >= stop) or (
                -markup + (last_price - pr) >= take
            ):
                last_deal = 2
                profit = -markup + (last_price - pr)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + (pr - last_price))
                continue

        # close deals by signals
        if last_deal == 1 and pred < 0.5 and direction == "buy":
            last_deal = 2
            profit = -markup + (pr - last_price)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + profit)
            continue

        if last_deal == 1 and pred < 0.5 and direction == "sell":
            last_deal = 2
            profit = -markup + (last_price - pr)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + (pr - last_price))
            continue

    return np.array(report), np.array(chart), line_f, line_b


def tester(*args):
    """
    This is a fast strategy tester based on numba
    List of parameters:

    dataset: must contain first column as 'close' and last columns with "labels" and "meta_labels"

    stop: stop loss value

    take: take profit value

    forward: forward time interval

    backward: backward time interval

    markup: markup value

    plot: false/true
    """
    dataset, stop, take, forward, backward, markup, plot = args

    forw = dataset.index.get_indexer([forward], method="nearest")[0]
    backw = dataset.index.get_indexer([backward], method="nearest")[0]

    close = dataset["close"].to_numpy()
    labels = dataset["labels"].to_numpy()
    metalabels = dataset["meta_labels"].to_numpy()

    report, chart, line_f, line_b = process_data(
        close, labels, metalabels, stop, take, markup, forw, backw
    )

    y = report.reshape(-1, 1)
    X = np.arange(len(report)).reshape(-1, 1)
    lr = LinearRegression()
    lr.fit(X, y)

    l = 1 if lr.coef_[0][0] >= 0 else -1

    if plot:
        plt.plot(report)
        plt.plot(chart)
        plt.axvline(x=line_f, color="purple", ls=":", lw=1, label="OOS")
        plt.axvline(x=line_b, color="red", ls=":", lw=1, label="OOS2")
        plt.plot(lr.predict(X))
        plt.title("Strategy performance R^2 " + str(format(lr.score(X, y) * l, ".2f")))
        plt.xlabel("the number of trades")
        plt.ylabel("cumulative profit in pips")
        plt.show()

    return lr.score(X, y) * l


def tester_one_direction(*args):
    """
    This is a fast strategy tester based on numba
    List of parameters:

    dataset: must contain first column as 'close' and last columns with "labels" and "meta_labels"

    stop: stop loss value

    take: take profit value

    forward: forward time interval

    backward: backward time interval

    markup: markup value

    direction: buy/sell

    plot: false/true
    """
    dataset, stop, take, forward, backward, markup, direction, plot = args

    forw = dataset.index.get_indexer([forward], method="nearest")[0]
    backw = dataset.index.get_indexer([backward], method="nearest")[0]

    close = dataset["close"].to_numpy()
    labels = dataset["labels"].to_numpy()
    metalabels = dataset["meta_labels"].to_numpy()

    report, chart, line_f, line_b = process_data_one_direction(
        close, labels, metalabels, stop, take, markup, forw, backw, direction
    )

    y = report.reshape(-1, 1)
    X = np.arange(len(report)).reshape(-1, 1)
    lr = LinearRegression()
    lr.fit(X, y)

    l = 1 if lr.coef_[0][0] >= 0 else -1

    if plot:
        plt.plot(report)
        plt.axvline(x=line_f, color="purple", ls=":", lw=1, label="OOS")
        plt.axvline(x=line_b, color="red", ls=":", lw=1, label="OOS2")
        plt.plot(lr.predict(X))
        plt.title("Strategy performance R^2 " + str(format(lr.score(X, y) * l, ".2f")))
        plt.xlabel("the number of trades")
        plt.ylabel("cumulative profit in pips")
        plt.show()

    return lr.score(X, y) * l


def tester_slow(dataset, stop, take, markup, forward, plot=False):
    last_deal = 2
    last_price = 0.0
    report = [0.0]
    chart = [0.0]
    line_f = 0
    line_b = 0

    # Assuming variables are defined elsewhere
    forw = dataset.index.get_indexer([forward], method="nearest")
    backw = dataset.index.get_indexer([forward], method="nearest")

    close = dataset["close"].to_numpy()
    labels = dataset["labels"].to_numpy()
    metalabels = dataset["meta_labels"].to_numpy()

    for i in range(dataset.shape[0]):
        line_f = len(report) if i <= forw else line_f
        line_b = len(report) if i <= backw else line_b

        pred = labels[i]
        pr = close[i]
        pred_meta = metalabels[i]  # 1 = allow trades

        if last_deal == 2 and pred_meta == 1:
            last_price = pr
            last_deal = 0 if pred < 0.5 else 1
            continue

        if last_deal == 0:
            if (-markup + (pr - last_price) >= take) or (
                -markup + (last_price - pr) >= stop
            ):
                last_deal = 2
                profit = -markup + (pr - last_price)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + profit)
                continue

        if last_deal == 1:
            if (-markup + (pr - last_price) >= stop) or (
                -markup + (last_price - pr) >= take
            ):
                last_deal = 2
                profit = -markup + (last_price - pr)
                report.append(report[-1] + profit)
                chart.append(chart[-1] + (pr - last_price))
                continue

        # close deals by signals
        if last_deal == 0 and pred > 0.5 and pred_meta == 1:
            last_deal = 2
            profit = -markup + (pr - last_price)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + profit)
            continue

        if last_deal == 1 and pred < 0.5 and pred_meta == 1:
            last_deal = 2
            profit = -markup + (last_price - pr)
            report.append(report[-1] + profit)
            chart.append(chart[-1] + (pr - last_price))
            continue

    y = np.array(report).reshape(-1, 1)
    X = np.arange(len(report)).reshape(-1, 1)
    lr = LinearRegression()
    lr.fit(X, y)

    l = 1 if lr.coef_ >= 0 else -1

    if plot:
        plt.plot(report)
        plt.plot(chart)
        plt.axvline(x=line_f, color="purple", ls=":", lw=1, label="OOS")
        plt.axvline(x=line_b, color="red", ls=":", lw=1, label="OOS2")
        plt.plot(lr.predict(X))
        plt.title("Strategy performance R^2 " + str(format(lr.score(X, y) * l, ".2f")))
        plt.xlabel("the number of trades")
        plt.ylabel("cumulative profit in pips")
        plt.show()

    return lr.score(X, y) * l


def test_model(
    dataset: pd.DataFrame,
    result: list,
    stop: float,
    take: float,
    forward: float,
    backward: float,
    markup: float,
    plt=False,
):

    ext_dataset = dataset.copy()
    X = ext_dataset[ext_dataset.columns[1:]]

    ext_dataset["labels"] = result[0].predict_proba(X)[:, 1]
    ext_dataset["meta_labels"] = result[1].predict_proba(X)[:, 1]
    ext_dataset["labels"] = ext_dataset["labels"].apply(
        lambda x: 0.0 if x < 0.5 else 1.0
    )
    ext_dataset["meta_labels"] = ext_dataset["meta_labels"].apply(
        lambda x: 0.0 if x < 0.5 else 1.0
    )
    return tester(ext_dataset, stop, take, forward, backward, markup, plt)


def test_model_one_direction(
    dataset: pd.DataFrame,
    result: list,
    stop: float,
    take: float,
    forward: float,
    backward: float,
    markup: float,
    direction: str,
    plt=False,
):

    ext_dataset = dataset.copy()
    X = ext_dataset[ext_dataset.columns[1:]]

    ext_dataset["labels"] = result[0].predict_proba(X)[:, 1]
    ext_dataset["meta_labels"] = result[1].predict_proba(X)[:, 1]
    ext_dataset["labels"] = ext_dataset["labels"].apply(
        lambda x: 0.0 if x < 0.5 else 1.0
    )
    ext_dataset["meta_labels"] = ext_dataset["meta_labels"].apply(
        lambda x: 0.0 if x < 0.5 else 1.0
    )
    return tester_one_direction(
        ext_dataset, stop, take, forward, backward, markup, direction, plt
    )


def test_model_one_direction_meta(
    dataset: pd.DataFrame,
    result: list,
    result_meta: list,
    stop: float,
    take: float,
    forward: float,
    backward: float,
    markup: float,
    direction: str,
    plt=False,
):

    # Copy the source dataset
    ext_dataset = dataset.copy()

    # Check the set of features for the main model
    if hasattr(result[0], "feature_names_"):
        expected_features = result[0].feature_names_
        current_features = (
            ext_dataset.columns[1:-2].tolist()
            if len(ext_dataset.columns) > 3
            else ext_dataset.columns[1:].tolist()
        )

        # Add absent features
        for feature in expected_features:
            if feature not in current_features and feature.isdigit():
                ext_dataset[feature] = 0  # Fill the appropriate features with zeros
                print(f"Added the absent feature: {feature}")

    # Create X_main for the main model
    if hasattr(result[0], "feature_names_"):
        X_main = ext_dataset[result[0].feature_names_]
    else:
        X_main = (
            ext_dataset[ext_dataset.columns[1:-2]]
            if len(ext_dataset.columns) > 3
            else ext_dataset.columns[1:]
        )

    # Main model forecasts
    try:
        ext_dataset["labels"] = result[0].predict_proba(X_main)[:, 1]
        ext_dataset["labels"] = ext_dataset["labels"].apply(
            lambda x: 0.0 if x < 0.5 else 1.0
        )
    except Exception as e:
        print(f"Error when forecasting main model: {e}")
        if hasattr(result[0], "feature_names_"):
            print(f"Expected features: {result[0].feature_names_}")
        print(f"Provided features: {X_main.columns.tolist()}")
        raise

    # For meta model
    meta_data = ext_dataset.copy()

    # Retrieve submodels from result_meta
    # Make sure result_meta contains models, not lists
    submodels = []
    if result_meta:
        # If the first element has the predict_proba method, then result_meta features the model list
        if hasattr(result_meta[0], "predict_proba"):
            submodels = result_meta
        # Otherwise, it may be the list of lists, let's try to retrieve the first model from each element
        elif isinstance(result_meta[0], list) and len(result_meta[0]) > 0:
            for item in result_meta:
                if hasattr(item[0], "predict_proba"):
                    submodels.append(item[0])

    # Check if there are sumbodels to handle
    if submodels:
        try:
            count = 0
            for meta_model in submodels:
                # Create X for submodel
                if hasattr(meta_model, "feature_names_"):
                    X_for_metamodel = ext_dataset[meta_model.feature_names_]
                else:
                    X_for_metamodel = X_main

                meta_data[f"meta_feature{count}"] = meta_model.predict_proba(
                    X_for_metamodel
                )[:, 1]
                count += 1
        except Exception as e:
            print(f"Error when forecasting submodels: {e}")
            print(f"meta_model type: {type(meta_model)}")
            if hasattr(meta_model, "feature_names_"):
                print(f"Expected submodel features: {meta_model.feature_names_}")
            print(f"Provided features: {X_for_metamodel.columns.tolist()}")
            raise
    else:
        print("Warning: No models for handling found.")

    # Prepare data for the final meta model
    try:
        if hasattr(result[1], "feature_names_"):
            # Check and add absent features
            for feature in result[1].feature_names_:
                if feature not in meta_data.columns:
                    meta_data[feature] = 0
                    print(f"Added absent feature for meta model: {feature}")

            X_meta_final = meta_data[result[1].feature_names_]
        else:
            # Delete columns with labels if any
            drop_cols = ["close", "labels", "meta_labels"]
            drop_cols = [col for col in drop_cols if col in meta_data.columns]
            X_meta_final = meta_data.drop(columns=drop_cols if drop_cols else [])
    except Exception as e:
        print(f"Error when preparing data for meta model: {e}")
        raise

    # Meta model forecasts
    try:
        ext_dataset["meta_labels"] = result[1].predict_proba(X_meta_final)[:, 1]
        ext_dataset["meta_labels"] = ext_dataset["meta_labels"].apply(
            lambda x: 0.0 if x < 0.5 else 1.0
        )
    except Exception as e:
        print(f"Error forecasting meta model: {e}")
        if hasattr(result[1], "feature_names_"):
            print(f"Expected meta model features: {result[1].feature_names_}")
        print(f"Provided features: {X_meta_final.columns.tolist()}")
        raise

    # Apply final test
    return tester_one_direction(
        ext_dataset, stop, take, forward, backward, markup, direction, plt
    )
