from __future__ import annotations
import numpy as np
from pyvelora.core.array_base import ArrayBase

class Vector(ArrayBase):
    def __init__(self, data, type=None, degrees=True):
        super().__init__(data)
        
        if type == "polar":
            if self.ndim != 1 or self.size != 2:
                raise ValueError("Polar coordinates must be a 1D vector of size 2 (r, theta)")
            r, theta = self.data
            if degrees:
                theta = np.deg2rad(theta)
            self.data = np.array([r * np.cos(theta), r * np.sin(theta)])
        
        elif type == "spherical":
            if self.ndim != 1 or self.size != 3:
                raise ValueError("Spherical coordinates must be a 1D vector of size 3 (r, theta, phi)")
            r, theta, phi = self.data
            if degrees:
                theta = np.deg2rad(theta)
                phi = np.deg2rad(phi)
            self.data = np.array([
                r * np.sin(theta) * np.cos(phi),
                r * np.sin(theta) * np.sin(phi),
                r * np.cos(theta)
            ])
        
        elif type == "cylindrical":
            if self.ndim != 1 or self.size != 3:
                raise ValueError("Cylindrical coordinates must be a 1D vector of size 3 (rho, phi, z)")
            rho, phi, z = self.data
            if degrees:
                phi = np.deg2rad(phi)
            self.data = np.array([
                rho * np.cos(phi),
                rho * np.sin(phi),
                z
            ])
        
        elif type is not None:
            raise ValueError(f"Unknown coordinate type: {type}")
            
    def magnitude(self):
        """Return the magnitude (length) of the vector."""
        return np.linalg.norm(self.data)
    
    def normalize(self):
        """Return a normalized (unit length) version of the vector."""
        mag = self.magnitude()
        if mag == 0:
            raise ValueError("Cannot normalize a zero vector")
        return Vector(self.data / mag)
    
if __name__ == "__main__":
    v = Vector([5, 45], type="polar", degrees=True)
    print(v.data)
    print(v.magnitude())
    print(v.normalize().data)
    print(type(v))
    
