# canonical.py

from __future__ import annotations
import numpy as np
from typing import Tuple


# ---------------------------------------------------------------------
# Canonicalization configuration
# ---------------------------------------------------------------------

# Numerical tolerance for zero detection
EPSILON: float = 1e-12


# ---------------------------------------------------------------------
# Canonicalization helpers
# ---------------------------------------------------------------------

def _normalize(values: np.ndarray) -> np.ndarray:
    """
    Normalize a complex amplitude vector.
    """
    norm = np.linalg.norm(values)
    if norm == 0:
        raise ValueError("Cannot canonicalize zero-norm state")
    return values / norm


def _remove_global_phase(values: np.ndarray) -> np.ndarray:
    """
    Remove global phase by forcing the first non-zero amplitude
    to be real and non-negative.
    """
    for v in values:
        if abs(v) > EPSILON:
            phase = np.angle(v)
            values = values * np.exp(-1j * phase)
            if values.real[0] < 0:
                values *= -1
            break
    return values


# ---------------------------------------------------------------------
# Public canonicalization API
# ---------------------------------------------------------------------

def canonicalize_sparse(
    indices: np.ndarray,
    values: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Canonicalize a sparse amplitude representation.

    Guarantees:
    - Deterministic normalization
    - Global phase removed
    - Output arrays are copies (caller mutation-safe)
    - Index ordering preserved (caller responsibility)

    Parameters
    ----------
    indices:
        Integer basis indices (shape: [k])
    values:
        Complex amplitudes (shape: [k])

    Returns
    -------
    (indices, values):
        Canonicalized sparse representation
    """
    if indices.ndim != 1 or values.ndim != 1:
        raise ValueError("indices and values must be 1-D arrays")

    if len(indices) != len(values):
        raise ValueError("indices and values must have the same length")

    # Copy defensively
    ci = np.array(indices, dtype=np.int64, copy=True)
    cv = np.array(values, dtype=np.complex128, copy=True)

    # Normalize
    cv = _normalize(cv)

    # Canonical global phase
    cv = _remove_global_phase(cv)

    return ci, cv


# ---------------------------------------------------------------------
# Optional utilities (explicit, not implicit)
# ---------------------------------------------------------------------

def canonicalize_dense(
    amplitudes: np.ndarray,
) -> np.ndarray:
    """
    Canonicalize a dense amplitude vector.

    Provided for completeness and testing;
    sparse canonicalization is preferred for infrastructure.
    """
    if amplitudes.ndim != 1:
        raise ValueError("amplitudes must be a 1-D array")

    values = np.array(amplitudes, dtype=np.complex128, copy=True)
    values = _normalize(values)
    values = _remove_global_phase(values)
    return values