# src/optimizer/models.py

import torch
import torch.nn as nn
import geoopt

from .utils import to_device  # Optional: for easy device placement in scripts


class SphereRosenbrockModel(nn.Module):
    """
    A simple batched model representing points on the 3-sphere (S^3) parameterized as unit quaternions.
    
    Used primarily for testing optimizers on the stereographically compactified 3D Rosenbrock function,
    which creates a challenging landscape with narrow valleys and pole singularities.
    
    Attributes:
        q: ManifoldParameter on the Sphere manifold (shape: [num_instances, 4])
    """
    def __init__(self, num_instances: int = 32, device: torch.device | None = None):
        super().__init__()
        
        # Initialize near the north pole (challenging starting region)
        init = torch.randn(num_instances, 4, dtype=torch.float64)
        init[..., 0] = 0.95 + 0.05 * torch.randn(num_instances)  # w component biased high
        init = init / init.norm(dim=-1, keepdim=True)  # Project to unit sphere
        
        self.manifold = geoopt.manifolds.Sphere()
        self.q = geoopt.ManifoldParameter(init, manifold=self.manifold)
        
        # Optional: move to device immediately if specified
        if device is not None:
            self.to(device)

    def forward(self) -> torch.Tensor:
        """
        Forward pass: simply return the quaternion parameters on the sphere.
        
        Returns:
            q: Tensor of shape [num_instances, 4]
        """
        return self.q


# Future-proof placeholders for additional benchmark models
# ------------------------------------------------------------------
# class StiefelOrthogonalModel(nn.Module):
#     """Example: Model with parameters on the Stiefel manifold (orthogonal frames)."""
#     ...
#
# class PoincareBallModel(nn.Module):
#     """Example: Hyperbolic embedding model on the Poincaré ball."""
#     ...
# ------------------------------------------------------------------