import jax.numpy as jnp
from jax import Array

from ..._problem import AbstractConstrainedMinimisation


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

    The problem is to fit (orthogonally) a torus to a
    set of points in 3D space. This set of points is generated by
    perturbing a first set lying exactly on a predefined torus
    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.
               minor correction by Ph. Shott, Jan 1995.

    classification QOR2-AY-V-V

    Square root of the number of data points (number of variables = 3 * NPTS**2 + 5)
    Default NPTS = 40 (so 1600 data points total)
    """

    npts: int = 40  # Square root of number of data points (default from SIF)
    y0_iD: int = 0
    provided_y0s: frozenset = frozenset({0})

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

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

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

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

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

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

        # Generate angles using meshgrid
        i_vals = jnp.arange(self.npts, dtype=jnp.float64)
        j_vals = jnp.arange(self.npts, dtype=jnp.float64)
        ii, jj = jnp.meshgrid(i_vals, j_vals, indexing="ij")

        # Flatten to match the expected order
        ii_flat = ii.ravel()
        jj_flat = jj.ravel()

        # Compute angles
        theta1 = ii_flat * incr
        theta2 = jj_flat * incr

        # Trig functions
        st1 = jnp.sin(theta1)
        ct1 = jnp.cos(theta1)
        st2 = jnp.sin(theta2)
        ct2 = jnp.cos(theta2)

        # Intermediate values
        p5ct1 = tp5 * ct1
        p4p5ct1 = tp4 + p5ct1
        r3 = tp5 * st1

        r1 = p4p5ct1 * ct2
        r2 = p4p5ct1 * st2

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

        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)

        # Torus parameters P1-P5
        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.5)  # P5

        # Point projections initialized to data points (interleaved X,Y,Z)
        xd, yd, zd = self._generate_data_points()

        # Variables are ordered as X(1,1), Y(1,1), Z(1,1), X(1,2), Y(1,2), Z(1,2), ...
        y = y.at[5::3].set(xd)  # Set all X values at indices 5, 8, 11, ...
        y = y.at[6::3].set(yd)  # Set all Y values at indices 6, 9, 12, ...
        y = y.at[7::3].set(zd)  # Set all Z values at indices 7, 10, 13, ...

        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 (interleaved X,Y,Z)
        x_proj = y[5::3]  # Gets indices 5, 8, 11, ... (all X values)
        y_proj = y[6::3]  # Gets indices 6, 9, 12, ... (all Y values)
        z_proj = y[7::3]  # Gets indices 7, 10, 13, ... (all Z values)

        # Sum of squared distances to data points
        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 torus parameters
        p1, p2, p3, p4, p5 = y[0], y[1], y[2], y[3], y[4]

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

        # Torus constraints from the SIF file:
        # A(I,J): ((X(I,J)-P1)^2 + (Y(I,J)-P2)^2)^0.5 * P4^-1 + (Z(I,J)-P3)^2
        # - P5^2 = 0
        # This can be rewritten as: (sqrt((X-P1)^2 + (Y-P2)^2) - P4)^2 +
        # (Z-P3)^2 - P5^2 = 0

        dx = x_proj - p1
        dy = y_proj - p2
        dz = z_proj - p3

        # Distance in XY plane
        r_xy = jnp.sqrt(dx**2 + dy**2)

        # Torus constraint
        eq_constraints = (r_xy - p4) ** 2 + dz**2 - p5**2

        ineq_constraints = None

        return eq_constraints, ineq_constraints

    @property
    def bounds(self) -> tuple[Array, Array] | None:
        """P4 and P5 have lower bounds 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)

        # P4 (index 3) and P5 (index 4) have lower bounds of 0.001
        lower = lower.at[3].set(0.001)
        lower = lower.at[4].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."""
        # Different values for different NPTS sizes given in SIF comments
        return jnp.array(0.0)
