import pandas as pd
import math
from datetime import datetime
from catboost import CatBoostClassifier
from sklearn.model_selection import train_test_split
from bots.botlibs.labeling_lib import *
from bots.botlibs.tester_lib import test_model
from bots.botlibs.export_lib import export_model_to_ONNX


def get_prices() -> pd.DataFrame:
    p = pd.read_csv("files/" + hyper_params["symbol"] + ".csv", sep="\s+")
    pFixed = pd.DataFrame(columns=["time", "close"])
    pFixed["time"] = p["<DATE>"] + " " + p["<TIME>"]
    pFixed["time"] = pd.to_datetime(pFixed["time"], format="mixed")
    pFixed["close"] = p["<CLOSE>"]
    pFixed.set_index("time", inplace=True)
    pFixed.index = pd.to_datetime(pFixed.index, unit="s")
    return pFixed.dropna()


def get_features(data: pd.DataFrame) -> pd.DataFrame:
    pFixed = data.copy()
    pFixedC = data.copy()
    count = 0

    for i in hyper_params["periods"]:
        pFixed[str(count)] = pFixedC.rolling(i).std()
        count += 1

    return pFixed.dropna()


def fit_final_models(dataset: pd.DataFrame) -> list:
    feature_columns = dataset.columns[1:-1]

    # 1. Data for the main model
    # Sort the dataset: only examples where 'labels; are equal to 0 or 1 are used for the main model.
    main_model_df = dataset[dataset["labels"].isin([0, 1])].copy()

    X = main_model_df[feature_columns]
    y = main_model_df["labels"].astype("int16")

    # 2. Meta model data
    X_meta = dataset[feature_columns]

    # Modify labels for meta model: if 'labels' contains 1 or 0, the new label is 1, if 2 - then 0.
    y_meta = (
        dataset["labels"]
        .apply(lambda label_val: 1 if label_val in [0, 1] else 0)
        .astype("int16")
    )

    # For the main model
    train_X, test_X, train_y, test_y = train_test_split(
        X, y, train_size=0.7, test_size=0.3, shuffle=True
    )

    # For meta model
    train_X_m, test_X_m, train_y_m, test_y_m = train_test_split(
        X_meta, y_meta, train_size=0.7, test_size=0.3, shuffle=True
    )

    # Train the main model
    model = CatBoostClassifier(
        iterations=1000,
        custom_loss=["Accuracy"],
        eval_metric="Accuracy",
        verbose=False,
        use_best_model=True,
        task_type="CPU",
    )

    # Check in case samples turned out to be empty after 'split' (the probability is low if the X size is sufficient)
    if not train_X.empty and not test_X.empty:
        model.fit(
            train_X,
            train_y,
            eval_set=(test_X, test_y),
            early_stopping_rounds=25,
            plot=False,
        )
    elif not train_X.empty:  # If the test sample is empty, but the training one exists
        print(
            "Warning: The test sample (test_X) for the main model is empty. The model is trained without eval_set."
        )
        model.fit(
            train_X, train_y, early_stopping_rounds=15, plot=False
        )  # use_best_model may work incorrectly without eval_set
    else:  # If the training set is empty
        print(
            "Error: The training set (train_X) for the main model is empty. The model cannot be trained."
        )
        # In this case, test_model will most likely throw an error later.
        # Return R2=-1 and the untrained model, the meta model will also not make sense without the main one.
        print("R2 is fixed at -1.0, models are not trained.")
        return [-1.0, model, None]  # model - instance, but not trained

    # Meta model training
    meta_model = CatBoostClassifier(
        iterations=1000,
        custom_loss=["F1"],
        eval_metric="F1",
        verbose=False,
        use_best_model=True,
        task_type="CPU",
    )

    if not train_X_m.empty and not test_X_m.empty:
        meta_model.fit(
            train_X_m,
            train_y_m,
            eval_set=(test_X_m, test_y_m),
            early_stopping_rounds=25,
            plot=False,
        )
    elif not train_X_m.empty:
        print(
            "Warning: The test sample (test_X_m) for the meta model is empty. The meta model is trained without eval_set."
        )
        meta_model.fit(train_X_m, train_y_m, early_stopping_rounds=25, plot=False)
    else:
        print(
            "Error: The training set (train_X_m) for the meta model is empty. The meta model cannot be trained."
        )
        print("R2 fixed as -1.0.")
        return [-1.0, model, meta_model]  # meta_model - instance, but not trained

    data_for_test = get_features(get_prices())
    R2 = test_model(
        data_for_test,
        [model, meta_model],
        hyper_params["stop_loss"],
        hyper_params["take_profit"],
        hyper_params["forward"],
        hyper_params["backward"],
        hyper_params["markup"],
        plt=False,
    )

    if math.isnan(R2):
        R2 = -1.0
        print("R2 fixed as -1.0")
    print("R2: " + str(R2))
    result = [R2, model, meta_model]
    return result


# set hyper parameters
hyper_params = {
    "symbol": "EURUSD_H1",
    "export_path": "/Users/dmitrievsky/Library/Containers/com.isaacmarovitz.Whisky/Bottles/54CFA88F-36A3-47F7-915A-D09B24E89192/drive_c/Program Files/MetaTrader 5/MQL5/Include/Trend following/",
    "model_number": 0,
    "markup": 0.00010,
    "stop_loss": 0.00500,
    "take_profit": 0.00500,
    "periods": [i for i in range(15, 300, 30)],
    "backward": datetime(2015, 1, 1),
    "forward": datetime(2024, 1, 1),
}

# plot the patterns
# dataset = get_prices()
# plot_best_n_patterns(dataset['close'].values, 50, 50, 5)

# fit the models
models = []
for i in range(1):
    print("Learn " + str(i) + " model")
    dataset = get_features(get_prices())
    data = dataset[
        (dataset.index < hyper_params["forward"])
        & (dataset.index > hyper_params["backward"])
    ].copy()
    data = get_fractal_pattern_labels_from_future_outcome(
        data, 100, 100, 0.9, 5, 25, 0.00010
    )
    models.append(fit_final_models(data))

# test the models
models.sort(key=lambda x: x[0])
data = get_features(get_prices())
test_model(
    data,
    models[-1][1:],
    hyper_params["stop_loss"],
    hyper_params["take_profit"],
    hyper_params["forward"],
    hyper_params["backward"],
    hyper_params["markup"],
    plt=True,
)

models[-1][2].get_best_score()["validation"]

export_model_to_ONNX(
    model=models[-1],
    symbol=hyper_params["symbol"],
    periods=hyper_params["periods"],
    periods_meta=hyper_params["periods"],
    model_number=hyper_params["model_number"],
    export_path=hyper_params["export_path"],
)
