# Filename: tutorial_timesfm.py

import yfinance as yf
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import timesfm
from os import path
from typing import Any, Sequence

import numpy as np
import torch
from huggingface_hub import snapshot_download


from timesfm.pytorch_patched_decoder import TimesFMConfig, PatchedTimeSeriesDecoder

import torch
import matplotlib.pyplot as plt

device = "cuda" if torch.cuda.is_available() else "cpu"


# --------------------------------------------------
# 1. Download stock data via yfinance
# --------------------------------------------------
def download_yfinance_data(ticker="AAPL", start="2020-01-01", end="2022-01-01"):
    """
    Download daily stock data for a given ticker from Yahoo Finance.
    Returns a pandas DataFrame with columns like 'Open', 'High', 'Low', 'Close', 'Volume'.
    """
    df = yf.download(ticker, start=start, end=end)
    df = df.dropna()
    return df["Close"].reset_index(drop=True)


# --------------------------------------------------
# 2. Create a dataset class for TimesFM
# --------------------------------------------------
class FinancialDataset(Dataset):
    def __init__(
        self,
        series: pd.Series,
        config: TimesFMConfig,
        context_length=128,  # how many past timesteps as input
        horizon_length=32,  # how many future steps to predict
    ):
        super().__init__()

        self.series = series.values.astype(np.float32)
        self.context_length = context_length
        self.horizon_length = horizon_length
        self.config = config

        self.samples = []
        # We want to ensure we have at least context_length + horizon_length points.
        for start_idx in range(0, len(self.series) - (context_length + horizon_length)):
            end_idx = start_idx + context_length
            # context slice
            x_context = self.series[start_idx:end_idx]
            # future/horizon slice
            x_future = self.series[end_idx : end_idx + horizon_length]
            self.samples.append((x_context, x_future))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        x_context, x_future = self.samples[index]
        # Convert to torch
        x_context = torch.tensor(x_context, dtype=torch.float32)
        x_future = torch.tensor(x_future, dtype=torch.float32)

        input_padding = torch.zeros_like(x_context)

        freq = torch.zeros(1, dtype=torch.long)

        return x_context, input_padding, freq, x_future


def collate_fn(batch):
    xs_context = [item[0] for item in batch]
    xs_padding = [item[1] for item in batch]
    freqs = [item[2] for item in batch]
    xs_future = [item[3] for item in batch]

    x_context = torch.stack(xs_context, dim=0)
    input_pad = torch.stack(xs_padding, dim=0)
    freq = torch.stack(freqs, dim=0)  # shape [B, 1]
    x_future = torch.stack(xs_future, dim=0)

    return x_context, input_pad, freq, x_future


def get_model(*, load_weights: bool = False):
    # standard model hack
    repo_id = "google/timesfm-2.0-500m-pytorch"
    tfm = timesfm.TimesFm(
        hparams=timesfm.TimesFmHparams(
            backend="cuda",
            per_core_batch_size=32,
            horizon_len=128,
            num_layers=50,
            use_positional_embedding=False,
            context_len=192,
        ),
        checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=repo_id),
    )

    model = PatchedTimeSeriesDecoder(tfm._model_config)

    if load_weights:
        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
        model.load_state_dict(loaded_checkpoint)
        print("After loading:")
        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
    model = model.to(device)

    # import sys
    # sys.exit(-1)
    # repo_id = "google/timesfm-1.0-200m"
    return model, tfm._model_config


def train_model(
    ticker="AAPL", start="2015-01-01", end="2022-01-01", train_split=0.8, batch_size=8, num_epochs=20, pretrained=False
):
    df_close = download_yfinance_data(ticker, start=start, end=end)
    model, config = get_model(load_weights=pretrained)

    total_len = len(df_close)
    train_size = int(total_len * train_split)
    val_size = total_len - train_size

    train_series = df_close.iloc[:train_size].reset_index(drop=True)
    val_series = df_close.iloc[train_size:].reset_index(drop=True)

    train_dataset = FinancialDataset(
        series=train_series, config=config, context_length=128, horizon_length=config.horizon_len
    )
    val_dataset = FinancialDataset(
        series=val_series, config=config, context_length=128, horizon_length=config.horizon_len
    )
    print("Train samples:", len(train_dataset))
    print("Val samples:", len(val_dataset))
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        model.train()
        total_train_loss = 0.0

        for x_context, x_padding, freq, x_future in train_dataloader:
            x_context, x_padding, freq, x_future = (
                x_context.to(device),
                x_padding.to(device),
                freq.to(device),
                x_future.to(device),
            )
            predictions = model(x_context, x_padding.float(), freq)
            # predictions shape => [B, N, horizon_len, (1 + #quantiles)]
            predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
            last_patch_pred = predictions_mean[:, -1, :]  # => [B, horizon_len]

            # x_future => [B, horizon_len]
            loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_train_loss += loss.item()

        avg_train_loss = total_train_loss / len(train_dataloader)

        # -------- Compute validation loss --------
        model.eval()
        total_val_loss = 0.0
        with torch.no_grad():
            for x_context, x_padding, freq, x_future in val_dataloader:
                x_context, x_padding, freq, x_future = (
                    x_context.to(device),
                    x_padding.to(device),
                    freq.to(device),
                    x_future.to(device),
                )
                predictions = model(x_context, x_padding.float(), freq)
                predictions_mean = predictions[..., 0]
                last_patch_pred = predictions_mean[:, -1, :]
                val_loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
                total_val_loss += val_loss.item()

        avg_val_loss = total_val_loss / max(len(val_dataloader), 1)

        print(f"[Epoch {epoch+1}] Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")

    torch.save(model.state_dict(), "timesfm_finetuned.pth")
    return model, train_dataloader, val_dataloader


def plot_predictions(model, dataloader):
    model.eval()
    with torch.no_grad():
        x_context, x_padding, freq, x_future = next(iter(dataloader))
        x_context, x_padding, freq, x_future = (
            x_context.to(device),
            x_padding.to(device),
            freq.to(device),
            x_future.to(device),
        )
        # Forward pass
        predictions = model(x_context, x_padding.float(), freq)
        # => [B, N, horizon_len, (1 + #quantiles)]
        predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
        last_patch_prediction = predictions_mean[:, -1, :]  # => [B, horizon_len]

        # We'll plot only the first sample in the batch
        i = 0
        pred_vals = last_patch_prediction[i].cpu().numpy()  # [horizon_len]
        context_vals = x_context[i].cpu().numpy()  # [context_len]
        future_vals = x_future[i].cpu().numpy()  # [horizon_len]

        horizon_len = future_vals.shape[0]
        context_len = context_vals.shape[0]

        plt.figure(figsize=(10, 5))

        # Plot context
        plt.plot(range(context_len), context_vals, label="Context (History)", color="blue")

        # Plot predicted future
        plt.plot(
            range(context_len, context_len + horizon_len),
            pred_vals,
            label="Predicted Future",
            color="orange",
        )

        # Plot ground truth future
        plt.plot(
            range(context_len, context_len + horizon_len),
            future_vals,
            label="Ground Truth Future",
            color="green",
            linestyle="--",
        )

        plt.xlabel("Time")
        plt.ylabel("Value")
        plt.title("Model Forecast vs. Ground Truth")
        plt.legend()
        plt.show()
        plt.savefig("pic_predictions.png")


if __name__ == "__main__":
    # Example usage
    model, train_dl, val_dl = train_model(
        ticker="AAPL",
        start="2012-01-01",
        end="2019-01-01",
        train_split=0.8,
        batch_size=256,
        num_epochs=50,
        pretrained=True,
    )

    plot_predictions(model, val_dl)