"""
Harmonic Oscillator Calculator for EZGA.

This module provides a simple harmonic oscillator potential (V = 0.5 * k * r^2)
implemented using NumPy for compatibility with the functional registry.
"""

from __future__ import annotations
import numpy as np
from typing import Union, Sequence, Optional, Tuple, Dict, Any, List
from . import register_calculator

def _harmonic_oscillator_kernel(positions: np.ndarray, k_spring: float) -> Tuple[float, np.ndarray]:
    """
    Computes energy and forces for a harmonic oscillator centered at the origin.
    
    V = 0.5 * k * sum(r_i^2)
    F = -k * r_i
    """
    if positions.size == 0:
        return 0.0, np.zeros((0, 3))
    
    # Energy: 0.5 * k * sum(x^2 + y^2 + z^2)
    r_sq = np.sum(positions**2, axis=1)
    energy = 0.5 * k_spring * np.sum(r_sq)
    
    # Forces: -dV/dr = -k * r
    forces = -k_spring * positions
    
    return float(energy), forces

def harmonic_oscillator(
    k_spring: float = 1.0,
    **factory_kwargs
) -> Callable[..., Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float, Dict[str, Any]]]:
    """
    Factory that returns a run function for the Harmonic Oscillator potential.
    """
    
    def run(
        symbols: Union[np.ndarray, Sequence[str]],
        positions: np.ndarray,
        cell: Optional[np.ndarray] = None,
        *,
        fixed: Optional[np.ndarray] = None,
        sampling_temperature: float = 0.0,
        **kwargs
    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float, Dict[str, Any]]:
        
        positions_arr = np.asarray(positions, dtype=float)
        symbols_arr = np.asarray(symbols, dtype=object)
        
        # Compute physics
        energy, forces = _harmonic_oscillator_kernel(positions_arr, k_spring)
        
        # Metadata construction
        # We include forces in metadata because the simulator expects 
        # a tuple (pos, sym, cell, energy, metadata)
        corrections = {
            "energy": energy,
            "forces": forces,
            "k_spring": k_spring,
            "calculator_type": "harmonic_oscillator"
        }
        
        # The simulator expects the FINAL state. 
        # In a simple calculator without relaxation, it's just the input state.
        return (
            positions_arr,
            symbols_arr,
            cell,
            energy,
            corrections
        )

    return run

# Self-register if imported
register_calculator("harmonic_oscillator", harmonic_oscillator())
register_calculator("harmonic_oscillator_k10", harmonic_oscillator(k_spring=10.0))
