"""Concrete 3D ansatz constructions built on coordinate-system semantics.

This module contains direct 3D director fields and section-to-3D lifting
helpers. All fields are expressed in local coordinate components and pushed
into ambient space by the bound coordinate system.
"""

from __future__ import annotations

from typing import Callable

import numpy as np

from ._utils import normalize_vector, pack_vector
from .base import Ansatz3D, _pushforward_tangent_components
from .frame import CoordinateSampleND, CoordinateSystemND
from .section import Ansatz2D

__all__ = [
    "Ansatz3D",
    "ConstantField3D",
    "CholestericSpiral3D",
    "CoordinateSpiral3D",
    "Hopfion",
    "SectionEmbedding3D",
    "StraightFinger",
    "LoopFinger",
    "ArtLoopFinger",
]


def _sample_xyz(sample: CoordinateSampleND) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return the first three coordinate components of a 3D sample."""

    coords = np.asarray(sample.coords)
    return coords[..., 0], coords[..., 1], coords[..., 2]


def _as_component_vector(value, name: str) -> np.ndarray:
    """Validate a fixed 3-component vector parameter."""

    value = np.asarray(value, dtype=np.float64)
    if value.shape != (3,):
        raise ValueError(f"`{name}` must have shape (3,), got {value.shape}.")
    return value


def _identity_lift_builder(coordinates: CoordinateSystemND) -> CoordinateSystemND:
    """Return the ambient chart unchanged for straight section lifts."""

    return coordinates


def _cylindrical_lift_builder(coordinates: CoordinateSystemND) -> CoordinateSystemND:
    """Derive the default cylindrical lift chart used by loop embeddings."""

    return coordinates.cylindrize(
        radius=0,
        angle=1,
        angle_min=0.0,
        angle_max=2.0 * np.pi,
        tol=1e-15,
        orientation=1.0,
    )


def _loop_section_input_map(back: bool, radial_shift: float = 0.0) -> Callable[[CoordinateSampleND], tuple[np.ndarray, np.ndarray]]:
    """Build the section-input selector for loop and art-loop embeddings."""

    def section_input_map(sample: CoordinateSampleND) -> tuple[np.ndarray, np.ndarray]:
        rho = np.asarray(sample.coords[..., 0])
        z = np.asarray(sample.coords[..., 2])
        radius = rho - radial_shift
        if back:
            radius = -radius
        return radius, z

    return section_input_map


class SectionEmbedding3D(Ansatz3D):
    """Generic lifting of a 2D section ansatz into a 3D coordinate chart.

    Public ``coordinates`` describe the ambient chart seen by callers and by
    inline masking helpers. The actual tangent-space lift may use a derived
    internal chart built from that ambient chart via ``lift_builder``.

    The lifting pipeline is:

    1. evaluate the ambient chart bound to this ansatz
    2. build internal ``lift_coordinates`` from it
    3. evaluate ``lift_coordinates`` on the same Euclidean points
    4. materialize the section raw field on two chosen lift coordinates
    5. optionally normalize that raw section field
    6. apply ``component_scale`` and add ``background_components``
    7. reinterpret the resulting three entries as components in the lift basis
       and push them through ``lift_sample.basis``

    This keeps masking semantics tied to the public ambient chart while allowing
    the section lift itself to happen in a different internal coordinate system.
    """

    def __init__(
        self,
        section: Ansatz2D,
        *,
        coordinates: CoordinateSystemND | None = None,
        lift_builder: Callable[[CoordinateSystemND], CoordinateSystemND] | None = None,
        section_axes: tuple[int, int] | None = None,
        section_input_map: Callable[[CoordinateSampleND], tuple[np.ndarray, np.ndarray]] | None = None,
        background_components: tuple[float, float, float] | np.ndarray = (0.0, 0.0, 0.0),
        normalize_section: bool = False,
        component_scale: tuple[float, float, float] | np.ndarray = (1.0, 1.0, 1.0),
    ):
        super().__init__(coordinates=coordinates)
        if (section_axes is None) == (section_input_map is None):
            raise TypeError("Pass exactly one of `section_axes` or `section_input_map`.")
        self.section = section
        self.lift_builder = lift_builder or _identity_lift_builder
        self.section_axes = None if section_axes is None else tuple(int(axis) for axis in section_axes)
        self.section_input_map = section_input_map
        self.background_components = _as_component_vector(background_components, "background_components")
        self.normalize_section = bool(normalize_section)
        self.component_scale = _as_component_vector(component_scale, "component_scale")
        if self.section_axes is not None:
            if len(self.section_axes) != 2:
                raise ValueError("`section_axes` must contain exactly two indices.")
            if any(axis not in (0, 1, 2) for axis in self.section_axes):
                raise ValueError(f"`section_axes` must contain only 0, 1, 2, got {self.section_axes}.")

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        raise RuntimeError("SectionEmbedding3D.call() evaluates its lift coordinates directly.")

    def _lift_coordinates(self) -> CoordinateSystemND:
        lift_coordinates = self.lift_builder(self.coordinates)
        if not isinstance(lift_coordinates, CoordinateSystemND):
            raise TypeError(
                "`lift_builder` must return a CoordinateSystemND, "
                f"got {type(lift_coordinates).__name__}."
            )
        if lift_coordinates.ndim != 3:
            raise ValueError(f"`lift_builder` must return a 3D coordinate system, got ndim={lift_coordinates.ndim}.")
        return lift_coordinates

    def _section_inputs(self, lift_sample: CoordinateSampleND) -> tuple[np.ndarray, np.ndarray]:
        if self.section_input_map is not None:
            inputs = self.section_input_map(lift_sample)
            if not isinstance(inputs, (tuple, list)) or len(inputs) != 2:
                raise ValueError("`section_input_map` must return exactly two arrays.")
            return np.asarray(inputs[0]), np.asarray(inputs[1])
        assert self.section_axes is not None
        return (
            np.asarray(lift_sample.coords[..., self.section_axes[0]]),
            np.asarray(lift_sample.coords[..., self.section_axes[1]]),
        )

    def _section_components(self, lift_sample: CoordinateSampleND) -> np.ndarray:
        u, v = self._section_inputs(lift_sample)
        components = np.asarray(self.section.call(u, v))
        expected_shape = lift_sample.coords.shape[:-1] + (3,)
        try:
            components = np.broadcast_to(components, expected_shape)
        except ValueError as exc:
            raise ValueError(
                f"Section output with shape {components.shape} cannot be broadcast to {expected_shape}."
            ) from exc
        if self.normalize_section:
            components = normalize_vector(components)
        dtype = np.result_type(lift_sample.coords, components, np.float64)
        scale = self.component_scale.astype(dtype, copy=False)
        background = self.background_components.astype(dtype, copy=False)
        components = np.asarray(components, dtype=dtype) * scale
        components = components + np.broadcast_to(background, expected_shape)
        return components

    def call(self, points) -> np.ndarray:
        """Evaluate the embedded 3D field on Euclidean sample points."""

        points = np.asarray(points)
        if points.shape[-1] != self.ndim:
            raise ValueError(f"Expected points with trailing dimension {self.ndim}, got {points.shape}.")
        lift_sample = self._lift_coordinates()(points)
        components = self._section_components(lift_sample)
        return _pushforward_tangent_components(components, lift_sample.basis)

    def in_coordinates(self, coordinates: CoordinateSystemND) -> "SectionEmbedding3D":
        """Rebind the ambient chart while preserving the lift configuration."""

        return SectionEmbedding3D(
            self.section,
            coordinates=coordinates,
            lift_builder=self.lift_builder,
            section_axes=self.section_axes,
            section_input_map=self.section_input_map,
            background_components=self.background_components,
            normalize_section=self.normalize_section,
            component_scale=self.component_scale,
        )


class ConstantField3D(Ansatz3D):
    """Constant director components in the local coordinate basis.

    The stored three-vector is interpreted as fixed local components, not as a
    Cartesian ambient vector. Ambient deformation therefore follows the bound
    coordinate system through raw pushforward.
    """

    def __init__(
        self,
        vector: tuple[float, float, float] | np.ndarray = (0.0, 0.0, 1.0),
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        vector = np.asarray(vector, dtype=np.float64)
        if vector.shape != (3,):
            raise ValueError(f"`vector` must have shape (3,), got {vector.shape}.")
        self.vector = vector

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        components = self.vector.astype(np.result_type(sample.coords, np.float64), copy=False)
        components = np.broadcast_to(components, sample.coords.shape[:-1] + (3,))
        return components


class CholestericSpiral3D(Ansatz3D):
    """Cholesteric helical director field in a bound curvilinear coordinate system.

    The spiral phase is read from ``sample.coords[..., phase_axis]``. The raw
    field is defined by local components in the two selected coordinate
    directions and pushed into ambient space through the tangent basis of the
    bound coordinate system. This makes the helix covariant under coordinate
    deformation of the ambient chart.
    """

    def __init__(
        self,
        phase_axis: int = 0,
        plane_axes: tuple[int, int] = (1, 2),
        period: float = 1.0,
        shift: float = 0.0,
        amplitude: float = 1.0,
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        if period == 0:
            raise ValueError("`period` must be non-zero.")
        phase_axis = int(phase_axis)
        plane_axes = tuple(int(axis) for axis in plane_axes)
        if len(plane_axes) != 2:
            raise ValueError("`plane_axes` must contain exactly two basis indices.")
        if phase_axis not in (0, 1, 2):
            raise ValueError(f"`phase_axis` must be one of 0, 1, 2, got {phase_axis}.")
        if any(axis not in (0, 1, 2) for axis in plane_axes):
            raise ValueError(f"`plane_axes` must contain only 0, 1, 2, got {plane_axes}.")
        if plane_axes[0] == plane_axes[1]:
            raise ValueError("`plane_axes` must contain distinct indices.")
        self.phase_axis = phase_axis
        self.plane_axes = plane_axes
        self.period = float(period)
        self.shift = float(shift)
        self.amplitude = float(amplitude)

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        phase_coordinate = np.asarray(
            sample.coords[..., self.phase_axis],
            dtype=np.result_type(sample.coords, np.float64),
        )
        phase = ((phase_coordinate + self.shift) * 2.0 * np.pi / self.period)[..., None]
        components = np.zeros(sample.coords.shape[:-1] + (3,), dtype=np.result_type(sample.coords, np.float64))
        components[..., self.plane_axes[0]] = self.amplitude * np.cos(phase[..., 0])
        components[..., self.plane_axes[1]] = self.amplitude * np.sin(phase[..., 0])
        return components


# Backward-compatible alias. Keep the old name exported for existing notebooks
# and user code, but prefer ``CholestericSpiral3D`` in new code.
CoordinateSpiral3D = CholestericSpiral3D


class Hopfion(Ansatz3D):
    """Hopfion ansatz expressed in local coordinate-basis components.

    The usual Hopfion angular formulas are evaluated on the associated
    curvilinear coordinates. The resulting three local director components are
    then pushed forward through the ambient tangent basis.
    """

    def __init__(self, size: float = 1.0, eps: float = 1e-8, coordinates: CoordinateSystemND | None = None):
        super().__init__(coordinates=coordinates)
        self.size = float(size)
        self.eps = float(eps)

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, y, z = _sample_xyz(sample)
        rho = np.sqrt(x**2 + y**2 + z**2)
        theta = np.arccos(z / (rho + self.eps))
        phi = np.arctan2(y, x)
        t = rho**2 / self.size + self.eps
        f = np.pi * (1 - np.power(1 + 4.22 / t**2, -0.5))
        stheta = np.arccos(1 - 2 * (np.sin(f) * np.sin(theta)) ** 2)
        sphi = -phi + np.arctan2(1 / (np.tan(f) + self.eps), np.cos(theta))
        components = pack_vector(
            np.sin(stheta) * np.cos(sphi),
            np.sin(stheta) * np.sin(sphi),
            np.cos(stheta),
        )
        return components


class StraightFinger(SectionEmbedding3D):
    """Convenience specialization of :class:`SectionEmbedding3D` for straight lifts.

    The public coordinate system is also the lift chart, so inline masking and
    section lifting both use the same ambient coordinates.
    """

    def __init__(self, section: Ansatz2D, uni: float = 1.0, coordinates: CoordinateSystemND | None = None):
        self.uni = float(uni)
        super().__init__(
            section,
            coordinates=coordinates,
            lift_builder=_identity_lift_builder,
            section_axes=(0, 2),
            background_components=(0.0, 0.0, self.uni),
            normalize_section=False,
            component_scale=(1.0, 1.0, 1.0),
        )

    def in_coordinates(self, coordinates: CoordinateSystemND) -> "StraightFinger":
        """Rebind the straight-finger construction to another ambient chart."""

        return StraightFinger(self.section, uni=self.uni, coordinates=coordinates)


class LoopFinger(SectionEmbedding3D):
    """Convenience specialization of :class:`SectionEmbedding3D` for loop lifts.

    Public masking semantics stay tied to the ambient chart bound to this
    ansatz. The internal cylindrical lift chart is derived from that ambient
    chart for section embedding only.
    """

    def __init__(
        self,
        section: Ansatz2D,
        uni: float = 1.0,
        back: bool = False,
        coordinates: CoordinateSystemND | None = None,
    ):
        self.uni = float(uni)
        self.back = bool(back)
        super().__init__(
            section,
            coordinates=coordinates,
            lift_builder=_cylindrical_lift_builder,
            section_input_map=_loop_section_input_map(back=self.back, radial_shift=0.0),
            background_components=(0.0, 0.0, self.uni),
            normalize_section=False,
            component_scale=(-1.0, -1.0, 1.0) if self.back else (1.0, 1.0, 1.0),
        )

    def in_coordinates(self, coordinates: CoordinateSystemND) -> "LoopFinger":
        """Rebind the loop-finger construction to another ambient chart."""

        return LoopFinger(self.section, uni=self.uni, back=self.back, coordinates=coordinates)


class ArtLoopFinger(SectionEmbedding3D):
    """Convenience specialization of :class:`SectionEmbedding3D` for art-loop lifts.

    As with :class:`LoopFinger`, ambient coordinates remain public while the
    cylindrical lift chart is used only internally. The sampled section is
    normalized before being lifted into 3D.
    """

    def __init__(
        self,
        section: Ansatz2D,
        rad: float,
        back: bool = False,
        coordinates: CoordinateSystemND | None = None,
    ):
        self.rad = float(rad)
        self.back = bool(back)
        super().__init__(
            section,
            coordinates=coordinates,
            lift_builder=_cylindrical_lift_builder,
            section_input_map=_loop_section_input_map(back=self.back, radial_shift=self.rad),
            background_components=(0.0, 0.0, 0.0),
            normalize_section=True,
            component_scale=(-1.0, -1.0, 1.0) if self.back else (1.0, 1.0, 1.0),
        )

    def in_coordinates(self, coordinates: CoordinateSystemND) -> "ArtLoopFinger":
        """Rebind the art-loop construction to another ambient chart."""

        return ArtLoopFinger(self.section, rad=self.rad, back=self.back, coordinates=coordinates)
