import pandas as pd
import torch
import pytorch_lightning as pl
from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer, Baseline
from pytorch_forecasting.data import GroupNormalizer
from pytorch_forecasting.metrics import SMAPE, MAE
from pytorch_lightning.callbacks import EarlyStopping

# Чтение данных из CSV файла
data = pd.read_csv("BTC_USDT_1m_OKX_4.csv", parse_dates=["Time"])
data["Time"] = pd.to_datetime(data["Time"])
data = data.sort_values("Time", ascending=True)
data = data.drop(['Unnamed: 0'], axis=1)

# Подготовка данных (предполагая, что дата колонка называется ‘date’)
data["time_idx"] = (data["Time"] - data["Time"].min()).dt.total_seconds() // 60  # Нормализация временной оси
data["time_idx"] = data["time_idx"].astype(int)

# Создание дополнительных требуемых колонок: например, ‘group’ если нужны группы (для TFT рекомендуется)
data["group"] ="0"  # Если у вас нет разных групп, вы можете использовать одно и то же значение

# Убедитесь, что данные отсортированы по времени
data = data.sort_values(by=["group", "time_idx"])

# Создаем TimeSeriesDataSet
max_prediction_length = 24  # Примеры длины прогнозирования (дополните вашими значениями)
max_encoder_length = 72  # Примеры длины входов для обучения (дополните вашими значениями)
training_cutoff = data["time_idx"].max() - max_prediction_length  # До какого момента тренируем

dataset = TimeSeriesDataSet(
    data[lambda x: x.time_idx <= training_cutoff],
    time_idx="time_idx",
    target="Close",
    group_ids=["group"],
    min_encoder_length=max_encoder_length // 2,  # Это обрезание для примера
    max_encoder_length=max_encoder_length,
    max_prediction_length=max_prediction_length,
    static_categoricals=["group"],
    time_varying_known_reals=["time_idx", "Open", "Highest", "Lowest", "Vol", "VolCcy"],
    time_varying_unknown_reals=["Close"],
    target_normalizer=GroupNormalizer(groups=["group"], transformation="softplus"),
    allow_missing_timesteps=True
)

train_dataloader = dataset.to_dataloader(train=True, batch_size=128)
val_dataloader = dataset.to_dataloader(train=False, batch_size=128)

tft = TemporalFusionTransformer.from_dataset(
    dataset,
    learning_rate=0.03,
    hidden_size=16,
    attention_head_size=1,
    dropout=0.1,
    hidden_continuous_size=8,
    output_size=7,  # Для примеров часов/времени
    loss=MAE(),
    log_interval=10,
    reduce_on_plateau_patience=4,
)

# Использование Pytorch Lightning для организации тренировки
trainer = pl.Trainer(
    max_epochs=30,
    #gpus=1 if torch.cuda.is_available() else 0,
    gradient_clip_val=0.1,
    callbacks=[EarlyStopping(monitor="val_loss", patience=10, mode="min")],
)

# Тренировка модели
trainer.fit(
    tft,
    train_dataloaders=train_dataloader,
    val_dataloaders=val_dataloader
)

INFO:pytorch_lightning.utilities.rank_zero:GPU available: False, used: False
INFO:pytorch_lightning.utilities.rank_zero:TPU available: False, using: 0 TPU cores
INFO:pytorch_lightning.utilities.rank_zero:HPU available: False, using: 0 HPUs
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-91-8106ea4f2d9e> in <cell line: 23>()
     21 
     22 # Тренировка модели
---> 23 trainer.fit(
     24     tft,
     25     train_dataloaders=train_dataloader,

1 frames
/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/compile.py in _maybe_unwrap_optimized(model)
    114         return model
    115     _check_mixed_imports(model)
--> 116     raise TypeError(
    117         f"`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`"
    118     )

TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `TemporalFusionTransformer`