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):
    
    # Копируем исходный датасет
    ext_dataset = dataset.copy()
    
    # Проверяем набор признаков для основной модели
    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()
        
        # Добавляем отсутствующие признаки
        for feature in expected_features:
            if feature not in current_features and feature.isdigit():
                ext_dataset[feature] = 0  # Заполняем нулями отсутствующие признаки
                print(f"Добавлен отсутствующий признак: {feature}")
    
    # Создаем X_main для основной модели
    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:]
    
    # Предсказания основной модели
    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"Ошибка при предсказании основной модели: {e}")
        if hasattr(result[0], 'feature_names_'):
            print(f"Ожидаемые признаки: {result[0].feature_names_}")
        print(f"Предоставленные признаки: {X_main.columns.tolist()}")
        raise
    
    # Для мета-модели
    meta_data = ext_dataset.copy()
    
    # Извлекаем подмодели из result_meta
    # Проверяем, что result_meta содержит модели, а не списки
    submodels = []
    if result_meta:
        # Если первый элемент имеет метод predict_proba, то result_meta - список моделей
        if hasattr(result_meta[0], 'predict_proba'):
            submodels = result_meta
        # Иначе это может быть список списков, попробуем извлечь первую модель из каждого элемента
        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])
    
    # Проверяем, есть ли подмодели для обработки
    if submodels:
        try:
            count = 0
            for meta_model in submodels:
                # Создаем X для подмодели
                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"Ошибка при предсказании под-моделей: {e}")
            print(f"Тип meta_model: {type(meta_model)}")
            if hasattr(meta_model, 'feature_names_'):
                print(f"Ожидаемые признаки подмодели: {meta_model.feature_names_}")
            print(f"Предоставленные признаки: {X_for_metamodel.columns.tolist()}")
            raise
    else:
        print("Предупреждение: Не найдено подмоделей для обработки.")
    
    # Подготавливаем данные для финальной мета-модели
    try:
        if hasattr(result[1], 'feature_names_'):
            # Проверяем и добавляем отсутствующие признаки
            for feature in result[1].feature_names_:
                if feature not in meta_data.columns:
                    meta_data[feature] = 0
                    print(f"Добавлен отсутствующий признак для мета-модели: {feature}")
            
            X_meta_final = meta_data[result[1].feature_names_]
        else:
            # Удаляем колонки с метками, если они есть
            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"Ошибка при подготовке данных для мета-модели: {e}")
        raise
    
    # Предсказания мета-модели
    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"Ошибка при предсказании мета-модели: {e}")
        if hasattr(result[1], 'feature_names_'):
            print(f"Ожидаемые признаки мета-модели: {result[1].feature_names_}")
        print(f"Предоставленные признаки: {X_meta_final.columns.tolist()}")
        raise
    
    # Применяем финальное тестирование
    return tester_one_direction(ext_dataset, stop, take, forward, backward, markup, direction, plt)