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. Данные для основной модели
    # Фильтруем датасет: для основной модели используются только те примеры, где 'labels' равны 0 или 1.
    main_model_df = dataset[dataset['labels'].isin([0, 1])].copy()
    
    X = main_model_df[feature_columns]
    y = main_model_df['labels'].astype('int16')

    # 2. Данные для мета-модели
    X_meta = dataset[feature_columns]
    
    # Модифицируем метки для мета-модели: если 'labels' содержит 1 или 0, то новая метка 1, если 2 - то 0.
    y_meta = dataset['labels'].apply(lambda label_val: 1 if label_val in [0, 1] else 0).astype('int16')

    # Для основной модели
    train_X, test_X, train_y, test_y = train_test_split(
        X, y, train_size=0.7, test_size=0.3, shuffle=True) 
    
    # Для мета-модели
    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)
    
    # Обучение основной модели
    model = CatBoostClassifier(iterations=1000,
                               custom_loss=['Accuracy'],
                               eval_metric='Accuracy',
                               verbose=False,
                               use_best_model=True,
                               task_type='CPU',
                               )
    
    # Проверка на случай, если после split выборки оказались пустыми (маловероятно при достаточном размере X)
    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: # Если тестовая выборка пуста, но обучающая есть
        print("Предупреждение: Тестовая выборка (test_X) для основной модели пуста. Модель обучается без eval_set.")
        model.fit(train_X, train_y, early_stopping_rounds=15, plot=False) # use_best_model может работать некорректно без eval_set
    else: # Если обучающая выборка пуста
        print("Ошибка: Обучающая выборка (train_X) для основной модели пуста. Модель не может быть обучена.")
        # В этом случае test_model далее, скорее всего, вызовет ошибку.
        # Возвращаем R2=-1 и необученную модель, мета-модель тоже не будет иметь смысла без основной.
        print("R2 зафиксирован как -1.0, модели не обучены.")
        return [-1.0, model, None] # model - инстанс, но не обученный

    # Обучение мета-модели
    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("Предупреждение: Тестовая выборка (test_X_m) для мета-модели пуста. Мета-модель обучается без eval_set.")
        meta_model.fit(train_X_m, train_y_m, early_stopping_rounds=25, plot=False)
    else:
        print("Ошибка: Обучающая выборка (train_X_m) для мета-модели пуста. Мета-модель не может быть обучена.")
        print("R2 зафиксирован как -1.0.")
        return [-1.0, model, meta_model] # meta_model - инстанс, но не обученный

    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 зафиксирован как -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'])


