import jax.numpy as jnp
from jax import Array

from ..._problem import AbstractConstrainedMinimisation


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

    The problem is to fit (orthogonally) a planar curve to a set of points
    in the plane. This set of points is generated by perturbing a
    first set lying exactly on the predefined curve.
    The curve is referred to as a cardioid in the original paper,
    but is in fact a circle.

    This problem is a modification of ORTHREGD.SIF.
    The start point was moved to keep any data points from
    converging to the center of the circle, because such
    points make the constraint Jacobian singular.

    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,
               modified by T. Plantagena, May 1994.

    classification QOR2-AY-V-V

    Number of data points (number of variables = 2 NPTS + 3)
    Default NPTS = 2000 (original value)
    """

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

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

    @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.1
        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 (modified to avoid singular Jacobian)
        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
        xd, yd = self._generate_data_points()
        y = y.at[3 : 3 + self.npts].set(xd)
        y = y.at[3 + self.npts :].set(yd)

        return y

    def objective(self, y: Array, args) -> Array:
        """Compute the objective function."""
        # Extract variables
        x = y[3 : 3 + self.npts]
        y_points = y[3 + self.npts :]

        # 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 : 3 + self.npts]
        y_points = y[3 + self.npts :]

        # Equality constraints
        # ((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."""
        # Should be near zero if fitting is perfect
        return jnp.array(0.0)
