from dataclasses import dataclass
import numpy as np
import pandas as pd
import torch
from torch import nn
from huggingface_hub import hf_hub_download
from sklearn.model_selection import train_test_split
from torch.utils.data import TensorDataset, DataLoader
from wandb.sdk.wandb_config import Config as WandbConfig
from sklearn.preprocessing import MinMaxScaler

@dataclass
class Config:
    criterion: str = "MSE"
    epochs: int = 30
    input_dim: int = 59
    hidden_dim: int = 128
    num_layers: int = 2
    output_dim: int = 1
    batch_size: int = 1
    optimizer: str = "adam"
    learning_rate: float = 0.0005
    dataset: str = "jawwaadsabree/CurryData"
    architecture: str = "LSTM"
    # Number of previous games to consider
    window_size: int = 15

    def to_dict(self) -> dict[str, str | int | float]:
        return {
            "criterion": self.criterion,
            "epochs": self.epochs,
            "input_dim": self.input_dim,
            "hidden_dim": self.hidden_dim,
            "num_layers": self.num_layers,
            "output_dim": self.output_dim,
            "batch_size": self.batch_size,
            "optimizer": self.optimizer,
            "learning_rate": self.learning_rate,
            "dataset": self.dataset,
            "architecture": self.architecture,
            "window_size": self.window_size
        }

def get_loaders(config: Config | WandbConfig, scaler):
    data = load_data()
    stats = normalize(data, scaler)
    X, y = create_sliding_windows(stats, config)
    X_train, X_val, X_test, y_train, y_val, y_test = split_data(X, y)

    ### Convert the data into PyTorch tensors
    # Train
    X_train_tensor = torch.Tensor(X_train)
    y_train_tensor = torch.Tensor(y_train)
    # Val
    X_val_tensor = torch.Tensor(X_val)
    y_val_tensor = torch.Tensor(y_val)
    # Test
    X_test_tensor = torch.Tensor(X_test)
    y_test_tensor = torch.Tensor(y_test)

    ### Create TensorDataset for each dataset (train, validation, test)
    # Train
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    # Val
    val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
    # Test
    test_dataset = TensorDataset(X_test_tensor, y_test_tensor)

    ### Create DataLoader for each dataset
    batch_size = config.batch_size
    # Train
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
    # Val
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    # Test
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    return train_loader, test_loader, val_loader
    
def  split_data(X, y, debug=False):
    # Split the data into train and test (80% train, 20% test)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
    # Further split the training data into train and validation sets (80% train, 20% validation)
    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, shuffle=False)
    print(f"Training set size: {X_train.shape}") if debug else None
    print(f"Test set size: {X_test.shape}") if debug else None
    print(f"Validation set size: {X_val.shape}") if debug else None
    return X_train, X_val, X_test, y_train, y_val, y_test

def normalize(stats_df: pd.DataFrame, scaler: MinMaxScaler, debug=False):
    # Convert the DataFrame to float
    stats_df_float = stats_df.astype(float)
    
    # Fit and transform the data using the scaler
    scaled_values = scaler.fit_transform(stats_df_float)
    
    # Create a new DataFrame from the scaled values to ensure the correct float dtype
    stats_normalized = pd.DataFrame(scaled_values, columns=stats_df.columns, index=stats_df.index)
    
    if debug:
        print(stats_normalized.to_string())
    
    return stats_normalized.to_numpy()

def create_sliding_windows(data, config: Config | WandbConfig):
    window_size = config.window_size
    X, y = [], []
    for i in range(len(data) - window_size):
        X.append(data[i:i + window_size])   # Input sequence
        y.append(data[i + window_size][9])  # 9 is the index of the target column (points)
    return np.array(X), np.array(y)

def fix_date(date: int) -> str:
    return pd.to_datetime(date).strftime('%Y-%m-%d')

def load_data(real_dates=False):
  from datasets import load_dataset, DatasetDict
  
  # Load the dataset from Hugging Face
  dataset = load_dataset("jawwaadsabree/CurryData")
  if not isinstance(dataset, DatasetDict):
    raise ValueError("Dataset is not a DatasetDict")
  
  # # Convert to pandas dataframes
  stats_df = dataset['train'].to_pandas()
  if not isinstance(stats_df, pd.DataFrame):
    raise ValueError("Stats DataFrame is not a pandas DataFrame")
  # # order by date
  stats_df = stats_df.sort_values(by='Date')
  # # reset index
  stats_df = stats_df.reset_index(drop=True)

  if real_dates:
    stats_df['Date'] = stats_df['Date'].apply(fix_date)

  return stats_df

# Set device for PyTorch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        # Bidirectional LSTM
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, output_size)  # Adjust for bidirectional LSTM
        self.tanh = nn.Tanh()
        self.dropout = nn.Dropout(0.2)  # Add dropout with a 20% probability of zeroing
        self.layer_norm = nn.LayerNorm(hidden_size * 2)  # LayerNorm after LSTM output

    def forward(self, x):
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)

        out, _ = self.lstm(x, (h0, c0))
        out = self.layer_norm(out)  # Apply layer normalization
        out = self.dropout(out)  # Apply dropout
        out = self.fc(out[:, -1, :])  # Get the last time-step's output
        out = self.tanh(out)
        return out
    
def get_model(config: Config | WandbConfig) -> tuple[LSTMModel, nn.L1Loss | nn.MSELoss, torch.optim.Adam | torch.optim.SGD]:
    model = LSTMModel(input_size=config.input_dim, hidden_size=config.hidden_dim, num_layers=config.num_layers, output_size=config.output_dim) # 20 is the number of features on the data

    # Define loss function and optimizer
    criterion = nn.MSELoss() if config.criterion == "MSE" else nn.L1Loss()
    if config.optimizer == "adam":
        optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    elif config.optimizer == "sgd":
        optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)
    
    return model, criterion, optimizer

def load_model():
    """Load the model from Hugging Face"""
    model_path = hf_hub_download(repo_id="jawwaadsabree/CurryModel", filename="model.pth", revision="v0.0.3")
    model = LSTMModel(input_size=59, hidden_size=128, num_layers=2, output_size=1)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()
    return model