from __future__ import annotations

import numpy as np
from typing import Any, Literal

from .array_base import ArrayBase


CoordinateType = Literal["polar", "spherical", "cylindrical", "complex"]


class Vector(ArrayBase):
    def __init__(
        self,
        data: Any,
        type: CoordinateType | None = None,
        degrees: bool = True,
    ) -> None:
        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 == "cartesian" or type is None:
            pass  # No transformation needed for Cartesian coordinates
            
        elif type is not None:
            raise ValueError(f"Unknown coordinate type: {type}")

    # ---- iteration ----

    def __iter__(self):
        """Yield Python-native int or float for each element (no trailing dot)."""
        for x in self.data:
            val = float(x)
            yield int(val) if val.is_integer() else val

    # ---- vector methods ----

    def magnitude(self) -> float:
        """Return the magnitude (length) of the vector."""
        return float(np.linalg.norm(self.data))

    def normalize(self) -> Vector:
        """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)
