"""Concrete 2D section ansatz profiles.

Section ansatzes define local director components in the canonical section
frame ``(u, mid, v)``. They are primarily used directly on 2D charts or as
building blocks for 3D embeddings such as straight and loop fingers.
"""

from __future__ import annotations

import numpy as np

from ._utils import pack_vector
from .base import Ansatz2D
from .frame import CoordinateSampleND, CoordinateSystemND

__all__ = [
    "Ansatz2D",
    "AxisSection",
    "BlochSection",
    "SpiralSection",
    "CF1",
    "CF2",
    "CF3",
    "CF4",
    "Twist",
]


def _sample_xz(sample: CoordinateSampleND) -> tuple[np.ndarray, np.ndarray]:
    """Return the two coordinate components used by section formulas."""

    return sample.coords[..., 0], sample.coords[..., 1]


def _first_last(array: np.ndarray):
    """Return the value at the first leading index and last trailing index."""

    index = [0] * array.ndim
    index[-1] = -1
    return array[tuple(index)]


def _last_first(array: np.ndarray):
    """Return the value at the last leading index and first trailing index."""

    index = [0] * array.ndim
    index[0] = -1
    return array[tuple(index)]


def _first_value(array: np.ndarray):
    """Return the first scalar entry of an array-like section sample."""

    return array[(0,) * array.ndim]


class AxisSection(Ansatz2D):
    """Axis-aligned section profile with radial amplitude decay.

    The supplied three-vector is interpreted as a fixed local director
    direction in the section frame ``(u, mid, v)``. Its magnitude is modulated
    by a simple inverse-radius envelope.
    """

    def __init__(self, axis, coordinates: CoordinateSystemND | None = None):
        super().__init__(coordinates=coordinates)
        axis = np.asarray(axis, dtype=np.float64)
        if axis.shape != (3,):
            raise ValueError(f"`axis` must have shape (3,), got {axis.shape}.")
        self.axis = axis

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        d2 = x**2 + z**2 + 1e-15
        f = d2 ** (-1 / 2) * 3
        return f[..., None] * self.axis


class BlochSection(Ansatz2D):
    """Bloch-like section texture in local section-frame components."""

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        d2 = x**2 + z**2 + 1e-15
        f = d2 ** (-1 / 1) * 5
        return pack_vector(x * f, np.ones_like(x) * f, z * f)


class SpiralSection(Ansatz2D):
    """Localized spiral section profile in local coordinates.

    The profile oscillates along the chosen in-section axis and decays away
    from the section origin.
    """

    def __init__(
        self,
        axis=(1, 0),
        pf: float = 1.1,
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        axis = np.asarray(axis, dtype=np.float64)
        if axis.shape != (2,):
            raise ValueError(f"`axis` must have shape (2,), got {axis.shape}.")
        norm = np.sqrt(np.sum(axis**2))
        if norm == 0:
            raise ValueError("`axis` must be non-zero.")
        self.axis = axis / norm
        self.pf = pf

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        axis = self.axis
        L = _first_last(z) * 2
        p = L * self.pf
        a = axis[0] * x + axis[1] * z
        d2 = x**2 + z**2
        phi = a / p * np.pi
        c, s = np.cos(phi), np.sin(phi)
        f = np.exp(-d2 / L**2 * 8) * 4
        return pack_vector(f * (axis[1] * c), f * s, f * (-axis[0] * c))


class CF1(Ansatz2D):
    """CF1 section profile in local section-frame components.

    This is the default closed-finger profile used in the notebook examples
    and loop-based toron constructions.
    """

    def __init__(
        self,
        pf: float = 0.75,
        wf: float = 0.2,
        A: float = -0.9 * np.pi,
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        self.pf = pf
        self.wf = wf
        self.A = A

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        L = _first_last(z) * 2
        p = L * self.pf
        w = L * self.wf
        alpha = self.A
        a = (z * np.sin(alpha) - x * np.cos(alpha)) / np.sqrt(2)
        phi = 2 * np.pi * a / p
        theta = self.A / np.cosh(x / w) * np.cos(np.pi * z / L)
        ctheta, stheta = np.cos(theta), np.sin(theta)
        return pack_vector(stheta * np.cos(phi), stheta * np.sin(phi), ctheta)


class CF2(Ansatz2D):
    """CF2 section profile in local section-frame components."""

    def __init__(
        self,
        pf: float = 0.75,
        wf: float = 0.2,
        A: float = -1.0 * np.pi,
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        self.pf = pf
        self.wf = wf
        self.A = A

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        L = _first_last(z) * 2
        p = L * self.pf
        w = L * self.wf
        phi = np.arctan2(x, z)
        theta = self.A / np.cosh(x / w) * np.cos(np.pi * z / L)
        ctheta, stheta = np.cos(theta), np.sin(theta)
        return pack_vector(stheta * np.cos(phi), stheta * np.sin(phi), ctheta)


class CF3(Ansatz2D):
    """CF3 section profile in local section-frame components."""

    def __init__(self, wf: float = 0.1, coordinates: CoordinateSystemND | None = None):
        super().__init__(coordinates=coordinates)
        self.wf = wf

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        L = _first_last(z) * 2
        w = L * self.wf
        phi = -np.pi / 2 * np.sin(np.pi / 2 * z / L)
        theta = np.pi / 2 * np.tanh(x / w / np.cos(np.pi * z / L))
        ctheta, stheta = np.cos(theta), np.sin(theta)
        return pack_vector(ctheta * np.sin(phi), ctheta * np.cos(phi), stheta)


class CF4(Ansatz2D):
    """CF4 section profile in local section-frame components."""

    def __init__(
        self,
        pf: float = 0.7,
        wf: float = 0.2,
        coordinates: CoordinateSystemND | None = None,
    ):
        super().__init__(coordinates=coordinates)
        self.pf = pf
        self.wf = wf

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, z = _sample_xz(sample)
        L = _first_last(z) * 2
        p = L * self.pf
        w = L * self.wf
        a = x / np.sqrt(2)
        phi = np.pi * (0.5 + 2 * a / p)
        theta = 0.5 * np.pi * np.tanh(
            (1 - 2 / np.cosh(x / w) * np.sin(np.pi * (0.5 - z / L) / 2))
            / np.sin(np.pi * (0.5 + z / L) / 2) ** (1 / 2)
        )
        ctheta, stheta = np.cos(theta), np.sin(theta)
        return pack_vector(ctheta * np.sin(phi), ctheta * np.cos(phi), stheta)


class Twist(Ansatz2D):
    """Monotone twist profile in local section-frame components."""

    def _call_sample(self, sample: CoordinateSampleND) -> np.ndarray:
        x, _z = _sample_xz(sample)
        xx = x - _first_value(x)
        xx = xx * np.pi / 2 / _last_first(xx)
        phi = xx * 0
        theta = xx
        ctheta, stheta = np.cos(theta), np.sin(theta)
        return pack_vector(ctheta * np.sin(phi), ctheta * np.cos(phi), stheta)
