import torch

import numpy as np

def predict_auto_regressive(model: torch.nn.Module, features: torch.Tensor, sequence_length: int = 96):
    """
    Predicts the next value in a sequence using an autoregressive approach.

    Args:
    - model: The trained PyTorch model.
    - features: Initial sequence of data with length equal to sequence_length.
    - sequence_length: The length of the sequence.

    Returns:
    - A tensor containing the predicted values.
    """

    model.eval()
    predictions = []

    with torch.no_grad():
        for _ in range(sequence_length):

            output = model(features.unsqueeze(0))
            predictions.append(output.item())
            
            features = torch.cat((features[1:], output), dim=0)
    
    
    return torch.tensor(predictions).unsqueeze(-1)

def predict_auto_regressive_quantile(model: torch.nn.Module, sample_from_dist, features: torch.Tensor, quantiles: torch.Tensor, sequence_length: int = 96):
    """
    Predicts the next value in a sequence using an autoregressive approach.

    Args:
    - model: The trained PyTorch model.
    - features: Initial sequence of data with length equal to sequence_length.
    - sequence_length: The length of the sequence.

    Returns:
    - A tensor containing the predicted values.
    """

    model.eval()
    predictions = []

    with torch.no_grad():
        for _ in range(sequence_length):
            output = model(features.unsqueeze(0))

            predictions.append(output.squeeze(0))

            sample = sample_from_dist(quantiles.cpu().numpy(), output.squeeze(0).cpu().numpy())

            features = torch.cat((features[1:], torch.tensor(sample,).to("cuda").unsqueeze(0).unsqueeze(0)), dim=0)
            features = features.float()

    return torch.stack(predictions)
            