import jax.numpy as jnp

from ..._problem import AbstractConstrainedMinimisation


class ORTHRGDM(AbstractConstrainedMinimisation):
    """An orthogonal regression problem with modified start.

    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.

    The start point is modified here to keep any data points from
    converging to the center.

    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, Sep 1993.

    Classification: QOR2-AY-V-V
    """

    y0_iD: int = 0
    provided_y0s: frozenset = frozenset({0})

    # Number of data points
    NPTS: int = 5000  # Default from SIF file (line 43)

    # True circle parameter (centered at the origin)
    TZ3: float = 1.7

    # Perturbation parameters
    PSEED: float = 237.1531
    PSIZE: float = 0.2

    @property
    def n(self):
        """Number of variables: 2 * NPTS + 3."""
        return 2 * self.NPTS + 3

    @property
    def m(self):
        """Number of constraints."""
        return self.NPTS

    def _generate_data_points(self):
        """Generate perturbed data points on a circle."""
        PI = jnp.pi
        INCR = 2.0 * PI / self.NPTS
        Z3SQ = self.TZ3 * self.TZ3
        FACT_BASE = 1.0 + Z3SQ

        # Vectorized computation
        i = jnp.arange(self.NPTS, dtype=jnp.float64)
        theta = i * INCR
        st = jnp.sin(theta)
        ct = jnp.cos(theta)

        # Cardioid factor
        fact = FACT_BASE + 0.25 * ct

        # Base circle points
        r1 = fact * ct
        r2 = fact * st

        # Perturbation
        xseed = theta * self.PSEED
        sseed = jnp.cos(xseed)
        pert = 1.0 + self.PSIZE * sseed

        # Perturbed data points
        xd = r1 * pert
        yd = r2 * pert

        return xd, yd

    @property
    def y0(self):
        """Initial guess."""
        xd, yd = self._generate_data_points()

        # Circle parameters
        z_params = jnp.array([1.0, 0.0, 1.0])  # Z1, Z2, Z3

        # Interleave x and y coordinates as pycutest does
        xy_interleaved = jnp.zeros(2 * self.NPTS)
        xy_interleaved = xy_interleaved.at[::2].set(xd)
        xy_interleaved = xy_interleaved.at[1::2].set(yd)

        # Concatenate: z_params, then interleaved x,y coordinates
        y0 = jnp.concatenate([z_params, xy_interleaved])

        return y0

    @property
    def args(self):
        """Generate data points."""
        return self._generate_data_points()

    def objective(self, y, args):
        """Compute the least-squares objective function."""
        xd, yd = args

        # Extract projection coordinates from interleaved layout
        xy_interleaved = y[3:]
        x_proj = xy_interleaved[::2]  # Even indices
        y_proj = xy_interleaved[1::2]  # Odd indices

        # Vectorized computation of squared distances
        ox = (x_proj - xd) ** 2
        oy = (y_proj - yd) ** 2

        # Sum all squared distances
        obj = jnp.sum(ox + oy)

        return obj

    def constraint(self, y):
        """Compute the constraint values."""
        z1 = y[0]
        z2 = y[1]
        z3 = y[2]

        # Extract projection coordinates from interleaved layout
        xy_interleaved = y[3:]
        x_proj = xy_interleaved[::2]  # Even indices
        y_proj = xy_interleaved[1::2]  # Odd indices

        # Vectorized computation of EA(i) and EB(i)
        dx = x_proj - z1
        dy = y_proj - z2
        t = dx * dx + dy * dy  # This is T in the SIF file

        # EA(i) = T^2
        ea = t * t

        # EB(i) = T * (1 + z3^2)^2
        z3_sq = z3 * z3
        t1 = 1.0 + z3_sq
        eb = t * t1 * t1

        # Constraint: EA(i) - EB(i) = 0
        constraints = ea - eb

        # Return (equality_constraints, inequality_constraints)
        # All constraints are equality constraints
        return constraints, None

    @property
    def bounds(self):
        """Returns the bounds on the variable y."""
        # No bounds specified in the SIF file
        return None

    @property
    def expected_result(self):
        """Expected optimal solution (not provided in SIF)."""
        return None

    @property
    def expected_objective_value(self):
        """Expected optimal objective value."""
        if self.NPTS == 10:
            return jnp.array(3.412121065)
        elif self.NPTS == 50:
            return jnp.array(15.59042181)
        elif self.NPTS == 250:
            return jnp.array(76.10435792)
        elif self.NPTS == 500:
            return jnp.array(151.2351183)
        elif self.NPTS == 1000:
            return jnp.array(304.899)
        else:
            return None
