import numpy as np
import torch
import wandb
import plotly.graph_objects as go
from .model import get_model, get_loaders, fix_date, load_data, Config
from .prediction import fix_points
from wandb.sdk.wandb_config import Config as WandbConfig

def predict(model, loader, scaler, device, no_log=False):
    model.eval()  # Set model to evaluation mode
    predictions = []
    actuals = []
    originalInputs = []

    with torch.no_grad():
        for inputs, labels in loader:
            inputs = inputs.to(device)
            outputs = model(inputs)
            _originalInputs = scaler.inverse_transform(inputs.squeeze().cpu().numpy())
            originalInputs.append(_originalInputs)
            predictions.append(outputs.cpu().numpy())
            actuals.append(labels.cpu().numpy())

    predictions = np.concatenate(predictions)
    actuals = np.concatenate(actuals)

    pred_ys = []
    actual_ys = []
    zs = []

    pred_ys = [fix_points(x, scaler) for x in predictions.flatten()]
    actual_ys =  [fix_points(y, scaler) for y in actuals.flatten()]

    for date in originalInputs:
        date = date[:, 0]
        for i, row in enumerate(date.flatten()):
            zs.append(row)


    # Assuming 'zs' is a tensor, first convert it to a NumPy array or list
    zs = zs.cpu().numpy() if isinstance(zs, torch.Tensor) else zs

    # Now you can apply set and sort operations
    zs = sorted(list(set(zs)))
    zs = [fix_date(date) for date in zs][-30:] if len(zs) > 30 else [fix_date(date) for date in zs][15:]

    # Create a figure
    fig = go.Figure()

    # pred_ys = [np.nan] * 15 + pred_ys
    # actual_ys = [np.nan] * 15 + actual_ys
    # Add Predictions plot
    fig.add_trace(go.Scatter(x=zs, y=pred_ys, mode='lines+markers', name='Predictions'))

    # Add Actual plot
    fig.add_trace(go.Scatter(x=zs, y=actual_ys, mode='lines+markers', name='Actual'))

    # Customize the layout (title, labels, etc.)
    fig.update_layout(
        title='Stephen Curry Points Prediction',
        xaxis_title='Date',
        yaxis_title='Points',
        xaxis=dict(tickmode='array', tickvals=zs, ticktext=zs),
        template='ggplot2',  # Optional: Change plot theme
        width=1500,  # Width of the plot
        height=400  # Height of the plot
    )
    if not no_log:
        wandb.log({"predictions vs actuals": fig})

    # Show the plot
    fig.show()

def predict_only(config: Config | WandbConfig, scaler, device):
    model, _, _ = get_model(config)
    with wandb.init(project="curry-scraper", config=config.to_dict()):
        _, _, validation_loader = get_loaders(config, scaler)
        predict(model, validation_loader, scaler, device)

def predict_next_game(model, scaler, device):
    last_15_df = load_data().tail(15)
    last_15_games_features = last_15_df.copy()
    last_15_games_features = scaler.fit_transform(last_15_games_features)
    last_15_games_features = last_15_games_features.reshape(1, 15, 59)  # Reshape for the model
    last_15_df['Date'] = last_15_df['Date'].apply(fix_date)
    model.to(device)  # Ensure model is on the correct device
    model.eval()  # Set model to evaluation mode

    # Assuming last_15_games_features is a NumPy array with the features of the last 15 games
    inputs = torch.tensor(last_15_games_features).float().to(device)  # Ensure the correct type and device

    # Get the prediction from the model
    with torch.no_grad():
        prediction = model(inputs)  # Output should be for the next game

    # Return the predicted value (e.g., predicted points)
    return fix_points(prediction.cpu().numpy(), scaler)