from __future__ import annotations

# third party
import numpy as np
import torch
import torch.nn as nn
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler

# built-in
import logging
from typing import Type

# *----------------------------------------------------*
#                        GLOBALS
# *----------------------------------------------------*

_logger = logging.getLogger(__name__)

# *----------------------------------------------------*
#                        CLASSES
# *----------------------------------------------------*

class SGNS_Torch:
    """PyTorch implementation of Skip-Gram with Negative Sampling."""

    def __init__(self,
                 V: int,
                 D: int,
                 in_weights: torch.Tensor | np.ndarray | None = None,
                 out_weights: torch.Tensor | np.ndarray | None = None,
                 *,
                 seed: int | None = None,
                 optim: Type[Optimizer] = torch.optim.SGD,
                 optim_kwargs: dict | None = None,
                 lr_sched: Type[LRScheduler] | None = None,
                 lr_sched_kwargs: dict | None = None,
                 device: str | None = None):
        """Initialize SGNS (negative sampling) in PyTorch.

        Shapes:
            - Embedding tables:
                in_weights:  (V, D) or None — row i is the “input” vector for token i.
                out_weights: (V, D) or None — row i is the “output” vector for token i.

        Args:
            V: Vocabulary size (number of nodes/tokens).
            D: Embedding dimensionality.
            in_weights: Optional starting input-embedding matrix of shape (V, D).
            out_weights: Optional starting output-embedding matrix of shape (V, D).
            seed: Optional RNG seed for PyTorch (controls init, sampling, and shuffles).
            optim: Optimizer class to instantiate. Defaults to plain SGD.
            optim_kwargs: Keyword arguments for the optimizer. Defaults to {"lr": 0.1}.
            lr_sched: Optional learning-rate scheduler class.
            lr_sched_kwargs: Keyword arguments for the scheduler (required if lr_sched is provided).
            device: Target device string (e.g. "cuda"). Defaults to CUDA if available, else CPU.
        """
        optim_kwargs = optim_kwargs or {"lr": 0.1}
        if lr_sched is not None and lr_sched_kwargs is None:
            raise ValueError("lr_sched_kwargs required when lr_sched is provided")

        self.V, self.D = int(V), int(D)

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

        # Seed torch
        self.seed = None if seed is None else int(seed)
        if self.seed is not None:
            torch.manual_seed(self.seed)
            if self.device.type == "cuda":
                torch.cuda.manual_seed_all(self.seed)

        # two embeddings as in/out matrices
        self.in_emb  = nn.Embedding(self.V, self.D, device=self.device)
        self.out_emb = nn.Embedding(self.V, self.D, device=self.device)

        # init / warm-start
        with torch.no_grad():
            if in_weights is not None:
                w = torch.as_tensor(in_weights, dtype=torch.float32, device=self.device)
                if w.shape != (self.V, self.D):
                    raise ValueError(f"in_weights must be (V,D); got {tuple(w.shape)}")
                self.in_emb.weight.copy_(w)
            else:
                nn.init.uniform_(self.in_emb.weight, -0.5 / self.D, 0.5 / self.D)

            if out_weights is not None:
                w = torch.as_tensor(out_weights, dtype=torch.float32, device=self.device)
                if w.shape != (self.V, self.D):
                    raise ValueError(f"out_weights must be (V,D); got {tuple(w.shape)}")
                self.out_emb.weight.copy_(w)
            else:
                nn.init.zeros_(self.out_emb.weight)

        self.to(self.device)
        _logger.info("SGNS_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, self.seed)

        params = list(self.in_emb.parameters()) + list(self.out_emb.parameters())
        # optimizer / scheduler
        self.opt = optim(params=params, **optim_kwargs)
        self.lr_sched = lr_sched(self.opt, **lr_sched_kwargs) if lr_sched is not None else None

    def predict(self,
                center: torch.Tensor,
                pos: torch.Tensor,
                neg: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute positive/negative logits for SGNS.

        Inputs:
            center: int tensor of shape (B,), values in [0, V)
            pos:    int tensor of shape (B,), values in [0, V)
            neg:    int tensor of shape (B, K), values in [0, V)

        Returns:
            pos_logits: (B,)
            neg_logits: (B, K)
        """
        center = center.to(self.device, dtype=torch.long)
        pos    = pos.to(self.device, dtype=torch.long)
        neg    = neg.to(self.device, dtype=torch.long)

        c  = self.in_emb(center)  # (B, D)
        pe = self.out_emb(pos)    # (B, D)
        ne = self.out_emb(neg)    # (B, K, D)

        pos_logits = (c * pe).sum(dim=-1)               # (B,)
        neg_logits = (c.unsqueeze(1) * ne).sum(dim=-1)  # (B, K)
        return pos_logits, neg_logits

    __call__ = predict

    def fit(self,
            centers: np.ndarray,
            contexts: np.ndarray,
            num_epochs: int,
            batch_size: int,
            num_negative_samples: int,
            noise_dist: np.ndarray,
            shuffle_data: bool,
            lr_step_per_batch: bool):
        """Train SGNS on the provided center/context pairs."""
        if noise_dist.ndim != 1 or noise_dist.size != self.V:
            raise ValueError(f"noise_dist must be 1-D with length {self.V}; got {noise_dist.shape}")
        _logger.info(
            "SGNS_Torch fit: epochs=%d batch=%d negatives=%d shuffle=%s",
            num_epochs, batch_size, num_negative_samples, shuffle_data
        )
        bce = nn.BCEWithLogitsLoss(reduction="mean")

        N = centers.shape[0]
        idx = np.arange(N)

        noise_probs = torch.as_tensor(noise_dist, dtype=torch.float32, device=self.device)
        # normalize if slightly off; enforce nonnegativity + finite sum
        if (noise_probs < 0).any():
            raise ValueError("noise_dist has negative entries")
        s = noise_probs.sum()
        if not torch.isfinite(s) or float(s.item()) <= 0.0:
            raise ValueError("noise_dist must have positive finite sum")
        if abs(float(s.item()) - 1.0) > 1e-6:
            noise_probs = noise_probs / s

        for epoch in range(1, int(num_epochs) + 1):
            epoch_loss = 0.0
            batches = 0

            if shuffle_data:
                if self.seed is None:
                    np.random.shuffle(idx)
                else:
                    np.random.default_rng(self.seed + epoch).shuffle(idx)

            for s_ in range(0, N, int(batch_size)):
                take = idx[s_:s_+int(batch_size)]
                if take.size == 0:
                    continue
                K = int(num_negative_samples)
                B = len(take)

                cen = torch.as_tensor(centers[take],  dtype=torch.long, device=self.device)  # (B,)
                pos = torch.as_tensor(contexts[take], dtype=torch.long, device=self.device)  # (B,)
                neg = torch.multinomial(noise_probs, num_samples=B * K, replacement=True).view(B, K)  # (B,K)

                pos_logits, neg_logits = self(cen, pos, neg)

                y_pos = torch.ones_like(pos_logits)
                y_neg = torch.zeros_like(neg_logits)
                loss_pos = bce(pos_logits, y_pos)
                loss_neg = bce(neg_logits, y_neg)

                loss = loss_pos + K * loss_neg

                self.opt.zero_grad(set_to_none=True)
                loss.backward()
                self.opt.step()

                if lr_step_per_batch and self.lr_sched is not None:
                    self.lr_sched.step()

                epoch_loss += float(loss.detach().cpu().item())
                batches += 1
                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss.item())

            if not lr_step_per_batch and self.lr_sched is not None:
                self.lr_sched.step()

            mean_loss = epoch_loss / max(batches, 1)
            _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)

    @property
    def in_embeddings(self) -> np.ndarray:
        W = self.in_emb.weight.detach().cpu().numpy()  # (V, D)
        _logger.debug("In emb shape: %s", W.shape)
        return W

    @property
    def out_embeddings(self) -> np.ndarray:
        W = self.out_emb.weight.detach().cpu().numpy()  # (V, D)
        _logger.debug("Out emb shape: %s", W.shape)
        return W

    @property
    def avg_embeddings(self) -> np.ndarray:
        return 0.5 * (self.in_embeddings + self.out_embeddings)

    # tiny helper for device move
    def to(self, device):
        self.in_emb.to(device)
        self.out_emb.to(device)
        return self


class SG_Torch:
    """PyTorch implementation of Skip-Gram (full softmax, **no biases**).

    This variant uses **no bias terms**: both projections are pure linear maps.

    Computation:
        x = one_hot(center, V)          # (B, V)
        y = x @ W_in                    # (B, D), with W_in ∈ R^{VxD}
        logits = y @ W_out              # (B, V), with W_out ∈ R^{DxV}
        loss = CrossEntropyLoss(logits, context)

    Embeddings:
        - Input embeddings  = rows of W_in        → shape (V, D)
        - Output embeddings = rows of W_outᵀ      → shape (V, D)
    """

    def __init__(self,
                 V: int,
                 D: int,
                 in_weights: torch.Tensor | np.ndarray | None = None,
                 out_weights: torch.Tensor | np.ndarray | None = None,
                 *,
                 seed: int | None = None,
                 optim: Type[Optimizer] = torch.optim.SGD,
                 optim_kwargs: dict | None = None,
                 lr_sched: Type[LRScheduler] | None = None,
                 lr_sched_kwargs: dict | None = None,
                 device: str | None = None):
        """Initialize the plain Skip-Gram (full softmax, **no biases**) model in PyTorch.

        Shapes:
            - Linear maps (no bias):
                W_in:  (V, D) — rows are input embeddings for tokens.
                W_out: (D, V) — maps D→V; rows of W_outᵀ are output embeddings.

            - Warm-starts:
                in_weights:  (V, D) or None — copied into W_in if provided.
                out_weights: (D, V) or None — copied into W_out if provided.

        Args:
            V: Vocabulary size (number of nodes/tokens).
            D: Embedding dimensionality.
            in_weights: Optional starting matrix for W_in with shape (V, D).
            out_weights: Optional starting matrix for W_out with shape (D, V).
            seed: Optional RNG seed for reproducibility.
            optim: Optimizer class to instantiate. Defaults to :class:`torch.optim.SGD`.
            optim_kwargs: Keyword args for the optimizer. Defaults to ``{"lr": 0.1}``.
            lr_sched: Optional learning-rate scheduler class.
            lr_sched_kwargs: Keyword args for the scheduler (required if ``lr_sched`` is provided).
            device: Target device string (e.g., ``"cuda"``). Defaults to CUDA if available, else CPU.

        Notes:
            The encoder/decoder are **bias-free** linear layers acting on one-hot centers:
            • ``in_emb = nn.Linear(V, D, bias=False)``
            • ``out_emb = nn.Linear(D, V, bias=False)``
            Forward pass produces vocabulary-sized logits and is trained with CrossEntropyLoss.
        """
        optim_kwargs = optim_kwargs or {"lr": 0.1}
        if lr_sched is not None and lr_sched_kwargs is None:
            raise ValueError("lr_sched_kwargs required when lr_sched is provided")

        self.V, self.D = int(V), int(D)

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

        # Seed torch (no global NumPy seeding)
        self.seed = None if seed is None else int(seed)
        if self.seed is not None:
            torch.manual_seed(self.seed)
            if self.device.type == "cuda":
                torch.cuda.manual_seed_all(self.seed)

        self.in_emb  = nn.Linear(self.V, self.D, bias=False, device=self.device)
        self.out_emb = nn.Linear(self.D, self.V, bias=False, device=self.device)

        # warm-starts (note Linear weights are (out_features, in_features))
        with torch.no_grad():
            if in_weights is not None:
                w_in = torch.as_tensor(in_weights, dtype=torch.float32, device=self.device)
                if w_in.shape != (self.V, self.D):
                    raise ValueError(f"in_weights must be (V,D); got {tuple(w_in.shape)}")
                self.in_emb.weight.copy_(w_in.T)  # (D,V)
            # else: use default PyTorch init

            if out_weights is not None:
                w_out = torch.as_tensor(out_weights, dtype=torch.float32, device=self.device)
                if w_out.shape != (self.D, self.V):
                    raise ValueError(f"out_weights must be (D,V); got {tuple(w_out.shape)}")
                self.out_emb.weight.copy_(w_out)  # (V,D) weight is (V,D) because (out=in V, in=D)
            # else: default init

        self.to(self.device)
        _logger.info("SG_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, self.seed)

        params = list(self.in_emb.parameters()) + list(self.out_emb.parameters())
        # optimizer / scheduler
        self.opt = optim(params=params, **optim_kwargs)
        self.lr_sched = lr_sched(self.opt, **lr_sched_kwargs) if lr_sched is not None else None

    def predict(self, center: torch.Tensor) -> torch.Tensor:
        center = center.to(self.device, dtype=torch.long)
        c = nn.functional.one_hot(center, num_classes=self.V).to(dtype=torch.float32, device=self.device)
        y = self.in_emb(c)
        z = self.out_emb(y)
        return z

    __call__ = predict

    def fit(self,
            centers: np.ndarray,
            contexts: np.ndarray,
            num_epochs: int,
            batch_size: int,
            shuffle_data: bool,
            lr_step_per_batch: bool,
            **_ignore):
        cce = nn.CrossEntropyLoss(reduction="mean")

        N = centers.shape[0]
        idx = np.arange(N)

        for epoch in range(1, int(num_epochs) + 1):
            epoch_loss = 0.0
            batches = 0

            if shuffle_data:
                if self.seed is None:
                    np.random.shuffle(idx)
                else:
                    np.random.default_rng(self.seed + epoch).shuffle(idx)

            for s in range(0, N, int(batch_size)):
                take = idx[s:s+int(batch_size)]
                if take.size == 0:
                    continue

                cen = torch.as_tensor(centers[take], dtype=torch.long, device=self.device)
                ctx = torch.as_tensor(contexts[take], dtype=torch.long, device=self.device)

                logits = self(cen)
                loss = cce(logits, ctx)

                self.opt.zero_grad(set_to_none=True)
                loss.backward()
                self.opt.step()

                if lr_step_per_batch and self.lr_sched is not None:
                    self.lr_sched.step()

                epoch_loss += float(loss.detach().cpu().item())
                batches += 1
                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss.item())

            if not lr_step_per_batch and self.lr_sched is not None:
                self.lr_sched.step()

            mean_loss = epoch_loss / max(batches, 1)
            _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)

    @property
    def in_embeddings(self) -> np.ndarray:
        W = self.in_emb.weight.detach().T.cpu().numpy()  # (V, D)
        _logger.debug("In emb shape: %s", W.shape)
        return W

    @property
    def out_embeddings(self) -> np.ndarray:
        W = self.out_emb.weight.detach().cpu().numpy()   # (V, D)
        _logger.debug("Out emb shape: %s", W.shape)
        return W

    @property
    def avg_embeddings(self) -> np.ndarray:
        return 0.5 * (self.in_embeddings + self.out_embeddings)

    # tiny helper for device move
    def to(self, device):
        self.in_emb.to(device)
        self.out_emb.to(device)
        return self


__all__ = ["SGNS_Torch", "SG_Torch"]

if __name__ == "__main__":
    pass
