import json
import torch
from torch import nn


class Model(nn.Module):
    """LSTM + MLP regressor that predicts the next step's numeric features.

    Each timestep is a vector of ``numeric_features`` normalized values
    followed by an integer ticker id, which is embedded and concatenated
    to the numeric features before the LSTM.
    """

    def __init__(
        self,
        num_tickers,
        numeric_features=5,
        embedding_dim=16,
        output_size=5,
    ):
        super().__init__()

        self.numeric_features = numeric_features

        self.ticker_embedding = nn.Embedding(
            num_embeddings=num_tickers + 1,
            embedding_dim=embedding_dim,
        )

        self.lstm = nn.LSTM(
            input_size=numeric_features + embedding_dim,
            hidden_size=128,
            num_layers=3,
            batch_first=True,
        )

        self.mlp = nn.Sequential(
            nn.Linear(128, 128),
            nn.GELU(),
            nn.Linear(128, 256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, 128),
            nn.GELU(),
            nn.Linear(128, output_size),
        )

    def forward(self, x):
        """Predict the next step.

        Args:
            x: Tensor of shape (batch, sequence, numeric_features + 1) where
                the last channel holds the ticker id.

        Returns:
            Tensor of shape (batch, output_size).
        """
        numeric = x[..., : self.numeric_features]
        ticker_ids = x[..., self.numeric_features].long()
        ticker_vectors = self.ticker_embedding(ticker_ids)
        x = torch.cat([numeric, ticker_vectors], dim=-1)
        output, _ = self.lstm(x)
        return self.mlp(output[:, -1, :])


class Normalizer(object):
    """Applies the per-ticker mean/std normalization stored in a config file."""

    def __init__(self, config_path: str | None = None):
        """Load the normalization config.

        Args:
            config_path: Path to the JSON config. Defaults to `./config.json`.
        """
        if config_path is None:
            config_path = "./config.json"

        with open(config_path) as f:
            self.config = json.load(f)

        self.ticker_ids = self.config["ticker_ids"]
        self.mean = torch.tensor(self.config["mean"], dtype=torch.float)
        self.std = torch.tensor(self.config["std"], dtype=torch.float)

    def __call__(self, x, ticker_name) -> torch.Tensor:
        """Normalize ``x`` with the statistics of ``ticker_name``."""
        return self.normalize(x, ticker_name)

    def normalize(self, x, ticker_name: str) -> torch.Tensor:
        """Return ``(x - mean) / std`` for the given ticker.

        Args:
            x: Array-like of shape (..., num_features). Volume is expected to
                already be log1p-transformed, as in training.
            ticker_name: Ticker symbol present in the config's ``ticker_ids``.

        Raises:
            KeyError: If the ticker is not in the config.
        """
        ticker_id = self.ticker_ids[ticker_name]
        x = torch.as_tensor(x, dtype=torch.float)
        return (x - self.mean[ticker_id]) / self.std[ticker_id]


def load_model(
    weights_path: str = "stock_predictor/model.pth",
    config_path: str | None = None,
    device: str | torch.device | None = None,
) -> Model:
    """Build a ``Model`` from the config and load trained weights into it.

    Args:
        weights_path: Path to the saved ``state_dict``.
        config_path: Path to the JSON config. Defaults to ``./config.json``.
        device: Device to load onto. Defaults to CUDA if available, else CPU.

    Returns:
        The model on ``device`` in evaluation mode.
    """
    if config_path is None:
        config_path = "./config.json"

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

    with open(config_path) as f:
        config = json.load(f)

    model = Model(
        num_tickers=len(config["ticker_ids"]),
        numeric_features=len(config["mean"][0]),
        output_size=len(config["mean"][0]),
    )
    state_dict = torch.load(weights_path, map_location=device, weights_only=True)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model