#!/usr/bin/env python
"""
Potentials

Functions of the different potentials (and their derivatives for the evolution)

@ Author: Moussouni, Yaël (MSc student) & Bhat, Junaid Ramzan (MSc student)
@ Institution:  Université de Strasbourg, CNRS, Observatoire astronomique
                de Strasbourg, UMR 7550, F-67000 Strasbourg, France
@ Date: 2024-11-29
"""

import numpy as np

MAX_VAL = 1e3

def kepler_potential(W_grid: np.ndarray, 
                     position_only: bool = False) -> np.ndarray:
    """Computes the Kepler potential: V(R) = -G*m1*m2/R 
    (assuming G = 1, m1 = 1, m2 = 1)
    assuming the point mass at (x = 0, y = 0).
    @params:
        - W: Phase-space vector
        - position_only: True if W is np.array([X, Y])
    @returns:
        - computed potential
    """
    if position_only: 
        X = W_grid[0]
        Y = W_grid[1]
    else:
        X = W_grid[0,0]
        Y = W_grid[0,1]
    # If X or Y is not an array (or a list), but rather a scalar, then we
    # create a list of one element so that it can work either way
    if np.ndim(X) == 0: X = np.array([X])
    if np.ndim(Y) == 0: Y = np.array([Y])
    R = np.sqrt(X**2 + Y**2)
    return -1/R

def hh_potential(W_grid: np.ndarray, 
                 position_only=False) -> np.ndarray:
    """Computes the Hénon-Heiles potential.
    @params:
        - W: Phase-space vector
        - position_only: True if W is np.array([X, Y])
    @returns:
        - POT: Potential
    """
    if position_only:
        X = W_grid[0]
        Y = W_grid[1]
    else:
        X = W_grid[0, 0]
        Y = W_grid[0, 1]
        
    # If X or Y is not an array (or a list), but rather a scalar, then we
    # create a list of one element so that it can work either way
    if np.ndim(X) == 0: X = np.array([X])
    if np.ndim(Y) == 0: Y = np.array([Y])

    POT = (X**2 + Y**2 + 2*X**2*Y - 2*Y**3/3)/2
    return POT        

def hh_evolution(t: np.ndarray, W: np.ndarray):
    """Computes the evolution from the HH potential derivative
    @params
        - t: Time (not used)
        - W: Phase space vector
    &returns 
        - dot W: Time derivative of the phase space vector
    """
    X = W[0 ,0]
    Y = W[0, 1]
    U = W[1, 0]
    V = W[1, 1]
    DX = U 
    DY = V
    DU = -(2*X*Y + X)
    DV = -(X**2 - Y**2 + Y)

    return np.array([[DX, DY], [DU, DV]])
