import jax.numpy as jnp
from jax import Array

from ..._problem import AbstractConstrainedMinimisation


class ORTHREGE(AbstractConstrainedMinimisation):
    """
    An orthogonal regression problem.

    The problem is to fit (orthogonally) an elliptic helix to a
    set of points in 3D space. This set of points is generated by
    perturbing a first set lying exactly on a predefined helix
    centered at the origin.

    Source:
    M. Gulliksson,
    "Algorithms for nonlinear Least-squares with Applications to
    Orthogonal Regression",
    UMINF-178.90, University of Umea, Sweden, 1990.

    SIF input: Ph. Toint, June 1990.

    classification QOR2-AY-V-V

    Number of data points (number of variables = 3 NPTS + 6)
    Default NPTS = 2500
    """

    npts: int = 2500  # Number of data points (default from SIF)
    y0_iD: int = 0
    provided_y0s: frozenset = frozenset({0})

    def __init__(self, npts: int = 2500):
        self.npts = npts

    @property
    def n(self) -> int:
        """Total number of variables: 6 helix parameters + 3 * NPTS projections."""
        return 6 + 3 * self.npts

    @property
    def m(self) -> int:
        """Number of constraints: 2 * NPTS equality constraints."""
        return 2 * self.npts

    def _generate_data_points(self):
        """Generate the perturbed helix data points."""
        # True helix parameters
        tp4, tp5, tp6 = 1.7, 0.8, 2.0

        # Perturbation parameters
        pseed = 237.1531
        psize = 0.2
        pi = jnp.pi

        # Computed parameters
        incr = 2.0 * pi / self.npts

        # Generate angles
        i_vals = jnp.arange(self.npts, dtype=jnp.float64)
        theta = i_vals * incr

        # Compute sin and cos
        st = jnp.sin(theta)
        ct = jnp.cos(theta)

        # Helix points
        r1 = tp4 * ct
        r2 = tp5 * st
        r3 = tp6 * theta

        # Add perturbation
        xseed = theta * pseed
        sseed = jnp.cos(xseed)
        pert = 1.0 + psize * sseed

        # Data points
        xd = r1 * pert
        yd = r2 * pert
        zd = r3 * pert

        return xd, yd, zd

    def starting_point(self) -> Array:
        """Return the starting point for the problem."""
        y = jnp.zeros(self.n, dtype=jnp.float64)

        # Helix parameters P1-P6
        y = y.at[0].set(1.0)  # P1
        y = y.at[1].set(0.0)  # P2
        y = y.at[2].set(1.0)  # P3
        y = y.at[3].set(1.0)  # P4
        y = y.at[4].set(0.0)  # P5
        y = y.at[5].set(0.25)  # P6

        # Point projections initialized to data points (vectorized)
        xd, yd, zd = self._generate_data_points()
        y = y.at[6::3].set(xd)  # Set all X values
        y = y.at[7::3].set(yd)  # Set all Y values
        y = y.at[8::3].set(zd)  # Set all Z values

        return y

    def objective(self, y: Array, args) -> Array:
        """Compute the objective function."""
        # Get data points
        xd, yd, zd = self._generate_data_points()

        # Extract projected points (vectorized)
        x_proj = y[6::3]  # Gets indices 6, 9, 12, ... (all X values)
        y_proj = y[7::3]  # Gets indices 7, 10, 13, ... (all Y values)
        z_proj = y[8::3]  # Gets indices 8, 11, 14, ... (all Z values)

        # Sum of squared distances to data points (vectorized)
        obj = jnp.sum((x_proj - xd) ** 2 + (y_proj - yd) ** 2 + (z_proj - zd) ** 2)

        return obj

    def constraint(self, y: Array):
        """Compute the equality constraints."""
        # Extract helix parameters
        p1, p2, p3, p4, p5, p6 = y[0], y[1], y[2], y[3], y[4], y[5]

        # Extract projected points (vectorized)
        x = y[6::3]  # Gets indices 6, 9, 12, ... (all X values)
        y_coord = y[7::3]  # Gets indices 7, 10, 13, ... (all Y values)
        z = y[8::3]  # Gets indices 8, 11, 14, ... (all Z values)

        # Helix constraints from the SIF file:
        # A(I): X(I) - P1 - P4*cos((Z(I)-P3)/P6) = 0
        # B(I): Y(I) - P2 - P5*cos((Z(I)-P3)/P6) = 0  (Note: both use cos, not sin!)

        t = (z - p3) / p6

        constraint_a = x - p1 - p4 * jnp.cos(t)
        constraint_b = y_coord - p2 - p5 * jnp.cos(t)

        # Interleave constraints: A(1), B(1), A(2), B(2), ...
        eq_constraints = jnp.zeros(2 * self.npts, dtype=y.dtype)
        eq_constraints = eq_constraints.at[::2].set(
            constraint_a
        )  # A constraints at even indices
        eq_constraints = eq_constraints.at[1::2].set(
            constraint_b
        )  # B constraints at odd indices

        ineq_constraints = None

        return eq_constraints, ineq_constraints

    @property
    def bounds(self) -> tuple[Array, Array] | None:
        """P6 has lower bound of 0.001, others are free."""
        lower = jnp.full(self.n, -jnp.inf, dtype=jnp.float64)
        upper = jnp.full(self.n, jnp.inf, dtype=jnp.float64)

        # P6 (index 5) has lower bound of 0.001
        lower = lower.at[5].set(0.001)

        return lower, upper

    @property
    def y0(self) -> Array:
        """Initial guess for the optimization problem."""
        return self.starting_point()

    @property
    def args(self):
        """Additional arguments for the objective and constraint functions."""
        return None

    @property
    def expected_result(self) -> Array:
        """Expected result of the optimization problem."""
        # Not explicitly given in the SIF file
        return jnp.zeros(self.n, dtype=jnp.float64)

    @property
    def expected_objective_value(self) -> Array:
        """Expected value of the objective at the solution."""
        # Should be near zero if fitting is perfect
        return jnp.array(0.0)
