import jax.numpy as jnp
from jax import Array

from ..._problem import AbstractConstrainedMinimisation


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

    The problem is to fit (orthogonally) a circle to a set of points
    in the plane. This set of points is generated by perturbing a
    first set lying exactly on a predefined circle centered at the
    origin. Each point mapped to the center of the circle has a constraint
    gradient of zero, making the Jacobian of the constraints singular.

    Certain data points (which serve to define the initial estimate
    of the unknowns) are modified so that they will converge to
    center of the circle, allowing us to study ill-conditioned
    Jacobians.
    With 10 singular points, the final objective is SOLTN = 118.037

    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, Mar 1991 and T. Plantenga, May 1993.

    classification QOR2-AY-V-V

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

    npts: int = 2500  # Number of data points (default from SIF)
    tdp_lo: int = 5  # First modified data point (1-indexed in SIF, 0-indexed here)
    tdp_hi: int = 14  # Last modified data point (1-indexed in SIF, 0-indexed here)
    y0_iD: int = 0
    provided_y0s: frozenset = frozenset({0})

    def __init__(self, npts: int = 2500, tdp_lo: int = 5, tdp_hi: int = 14):
        self.npts = npts
        self.tdp_lo = tdp_lo
        self.tdp_hi = tdp_hi

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

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

    def _generate_data_points(self):
        """Generate the perturbed data points."""
        tz3 = 1.7
        pseed = 237.1531
        psize = 0.2
        pi = jnp.pi

        # Increment for angle
        icr0 = 1.0 / self.npts
        incr = icr0 * 2.0 * pi

        # Computed parameters
        z3sq = tz3 * tz3
        one_plus_z3sq = 1.0 + z3sq

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

        # Radius factor
        fact = one_plus_z3sq + ct
        r1 = fact * ct
        r2 = fact * st

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

        # Z parameters
        y = y.at[0].set(1.0)  # z1
        y = y.at[1].set(0.0)  # z2
        y = y.at[2].set(1.0)  # z3

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

        # Modify certain starting points as specified in SIF
        # From TDP_lo to TDP_hi (1-indexed in SIF, so convert to 0-indexed)
        for i in range(self.tdp_lo - 1, min(self.tdp_hi, self.npts)):
            y = y.at[3 + 2 * i].set(1.8)  # X(I) = 1.8
            y = y.at[4 + 2 * i].set(1.0)  # Y(I) = 1.0

        return y

    def objective(self, y: Array, args) -> Array:
        """Compute the objective function."""
        # Extract variables (interleaved)
        x = y[3::2]  # Gets indices 3, 5, 7, ... (all X values)
        y_points = y[4::2]  # Gets indices 4, 6, 8, ... (all Y values)

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

        # Objective: sum of squared distances to data points
        obj = jnp.sum((x - xd) ** 2 + (y_points - yd) ** 2)

        return obj

    def constraint(self, y: Array):
        """Compute the equality and inequality constraints."""
        # Extract variables
        z1 = y[0]
        z2 = y[1]
        z3 = y[2]
        x = y[3::2]  # Gets indices 3, 5, 7, ... (all X values)
        y_points = y[4::2]  # Gets indices 4, 6, 8, ... (all Y values)

        # Equality constraints from SIF:
        # E(I): ((x[i]-z1)^2+(y[i]-z2)^2)^2 - ((x[i]-z1)^2+(y[i]-z2)^2)*(1+z3^2)^2 = 0

        dx = x - z1
        dy = y_points - z2
        dist_sq = dx**2 + dy**2
        z3sq = z3 * z3
        one_plus_z3sq = 1.0 + z3sq

        eq_constraints = dist_sq * dist_sq - dist_sq * one_plus_z3sq * one_plus_z3sq

        # No inequality constraints
        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)
