import jax.numpy as jnp
from jax import Array

from ..._problem import AbstractConstrainedMinimisation


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

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

    Source: adapted from:
    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 QQR2-AN-V-V

    Number of data points (number of variables = 2 NPTS + 5)
    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: 5 ellipse parameters + 2 * NPTS projections."""
        return 5 + 2 * self.npts

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

    def _generate_data_points(self):
        """Generate the perturbed ellipse data points."""
        # True ellipse parameters
        v11, v12, v22 = 2.0, 1.0, 2.0

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

        # Computed parameters
        incr = 2.0 * pi / self.npts
        c3 = jnp.cos(v22)
        s3 = jnp.sin(v22)

        # 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)

        # Ellipse points before rotation
        u1 = v11 * ct
        u2 = v12 * st

        # Apply rotation by angle v22
        r1 = u1 * c3 - u2 * s3
        r2 = u1 * s3 + u2 * c3

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

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

        return xd, yd

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

        # Ellipse parameters (H11, H12, H22, G1, G2)
        y = y.at[0].set(1.0)  # H11
        y = y.at[1].set(0.0)  # H12
        y = y.at[2].set(1.0)  # H22
        y = y.at[3].set(1.0)  # G1
        y = y.at[4].set(1.0)  # G2

        # Point projections initialized to data points (interleaved X,Y)
        xd, yd = self._generate_data_points()
        y = y.at[5::2].set(xd)  # Set all X values at indices 5, 7, 9, ...
        y = y.at[6::2].set(yd)  # Set all Y values at indices 6, 8, 10, ...

        return y

    def objective(self, y: Array, args) -> Array:
        """Compute the objective function."""
        # Extract projected points (interleaved X,Y)
        x_proj = y[5::2]  # Gets indices 5, 7, 9, ... (all X values)
        y_proj = y[6::2]  # Gets indices 6, 8, 10, ... (all Y values)

        # Get data points
        xd, yd = self._generate_data_points()

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

        return obj

    def constraint(self, y: Array):
        """Compute the equality constraints."""
        # Extract ellipse parameters
        h11, h12, h22, g1, g2 = y[0], y[1], y[2], y[3], y[4]

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

        # Ellipse constraints: H11*x^2 + 2*H12*x*y + H22*y^2 - 2*G1*x - 2*G2*y - 1 = 0
        eq_constraints = (
            h11 * x_proj**2
            + 2.0 * h12 * x_proj * y_proj
            + h22 * y_proj**2
            - 2.0 * g1 * x_proj
            - 2.0 * g2 * y_proj
            - 1.0
        )

        ineq_constraints = None

        return eq_constraints, ineq_constraints

    @property
    def bounds(self) -> tuple[Array, Array] | None:
        """All variables are free."""
        return None

    @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)
