import json
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import joblib
from typing import Optional, Dict, Any
from huggingface_hub import hf_hub_download

class LSTMClassifier(nn.Module):
    def __init__(self, input_size: int, hidden_size: int = 64, num_layers: int = 1, dropout: float = 0.0):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
            bidirectional=False
        )
        self.head = nn.Linear(hidden_size, 2)

    def forward(self, x):
        _, (h_n, _) = self.lstm(x)
        last_h = h_n[-1]
        return self.head(last_h)

def load_model_and_scaler(repo_id: str, revision: Optional[str] = None, device: Optional[str] = None):
    cfg_path = hf_hub_download(repo_id, "config.json", revision=revision)
    scaler_path = hf_hub_download(repo_id, "scaler.joblib", revision=revision)

    with open(cfg_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

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

    model = LSTMClassifier(
        input_size=int(cfg["input_size"]),
        hidden_size=int(cfg["hidden_size"]),
        num_layers=int(cfg["num_layers"]),
        dropout=float(cfg["dropout"]),
    ).to(device)

    weights_name = cfg.get("weights_file", "model.safetensors")
    weights_path = hf_hub_download(repo_id, weights_name, revision=revision)

    if weights_name.endswith(".safetensors"):
        from safetensors.torch import load_file
        state = load_file(weights_path)
        model.load_state_dict({k: v for k, v in state.items()}, strict=True)
    else:
        state = torch.load(weights_path, map_location="cpu")
        model.load_state_dict(state, strict=True)

    model.eval()
    scaler = joblib.load(scaler_path)
    return model, scaler, cfg

def predict_df(df: pd.DataFrame, model: nn.Module, scaler, cfg: Dict[str, Any]) -> np.ndarray:
    from numpy.lib.stride_tricks import sliding_window_view

    feature_cols = cfg["feature_cols"]
    W = int(cfg["window_size"])
    stride = int(cfg.get("stride", 1))

    X = df[feature_cols].to_numpy(np.float32)
    if len(X) < W:
        return np.empty((0,), dtype=np.int64)

    Xw = sliding_window_view(X, window_shape=(W, X.shape[1])).squeeze(1)
    Xw = Xw[::stride]

    F = Xw.shape[2]
    Xw_scaled = scaler.transform(Xw.reshape(-1, F)).reshape(Xw.shape).astype(np.float32)

    device = next(model.parameters()).device
    with torch.no_grad():
        xb = torch.tensor(Xw_scaled, device=device)
        logits = model(xb)
        y_pred = torch.argmax(logits, dim=1).detach().cpu().numpy()
    return y_pred