from __future__ import annotations

import os

import matplotlib.pyplot as plt
import numpy as np

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes


thetad = np.array([])
xd = np.array([])
yd = np.array([])


def F(x):
    # F=dot(f,f)
    a = x[0]
    b = x[1]
    phi = x[2]
    n = len(xd)
    f = np.zeros(n)
    for i in range(n):
        xf = a * np.cos(phi) * np.cos(thetad[i]) - b * np.sin(phi) * np.sin(thetad[i]) - xd[i]
        yf = b * np.cos(phi) * np.sin(thetad[i]) + a * np.sin(phi) * np.cos(thetad[i]) - yd[i]
        f[i] = xf**2 + yf**2
    return f


def Jac(x):
    # jacobian
    a = x[0]
    b = x[1]
    phi = x[2]
    n = len(xd)
    g = np.zeros((n, 3))
    for i in range(n):
        xf = a * np.cos(phi) * np.cos(thetad[i]) - b * np.sin(phi) * np.sin(thetad[i]) - xd[i]
        yf = b * np.cos(phi) * np.sin(thetad[i]) + a * np.sin(phi) * np.cos(thetad[i]) - yd[i]
        g[i, 0] = 2 * xf * np.cos(phi) * np.cos(thetad[i]) + 2 * yf * np.sin(phi) * np.cos(thetad[i])
        g[i, 1] = -2 * xf * np.sin(phi) * np.sin(thetad[i]) + 2 * yf * np.cos(phi) * np.sin(thetad[i])

        g31 = 2 * xf * (-a * np.sin(phi) * np.cos(thetad[i]) - b * np.cos(phi) * np.sin(thetad[i]))
        g31 = g31 + 2 * yf * (-b * np.sin(phi) * np.sin(thetad[i]) + a * np.cos(phi) * np.cos(thetad[i]))
        g[i, 2] = g31
    return g


def fig9_18():
    global thetad, xd, yd

    # generate some data points
    n = 50
    thetad = 2 * np.pi * np.random.rand(n)
    xd = np.zeros(n)
    yd = np.zeros(n)
    a_data = np.zeros(n)
    b_data = np.zeros(n)
    phi_data = np.zeros(n)
    for i in range(n):
        a_data[i] = 2 + 0.4 * (np.random.rand() - 0.5)
        b_data[i] = 1 + 0.4 * (np.random.rand() - 0.5)
        phi_data[i] = np.pi * (1 + 0.3 * (np.random.rand() - 0.5)) / 3
        xd[i] = a_data[i] * np.cos(phi_data[i]) * np.cos(thetad[i]) - b_data[i] * np.sin(phi_data[i]) * np.sin(thetad[i])
        yd[i] = b_data[i] * np.cos(phi_data[i]) * np.sin(thetad[i]) + a_data[i] * np.sin(phi_data[i]) * np.cos(thetad[i])

    # circle_data=[cxd;yd; a ;b ;phi];
    # save('/Users/mark/Desktop/cicrle_data.txt','circle_data','-ascii');

    # circle_data=importdata('cicrle_data.txt');
    # xd=circle_data(1,:);
    # yd=circle_data(2,:);
    # a=circle_data(3,1); b=circle_data(4,1); phi=circle_data(5,1);

    x = np.array([1.0, 0.5, 0.0])
    tol = 1e-6
    xpoints = [x[0]]
    ypoints = [x[1]]
    f = F(x)
    ff = np.linalg.norm(f, 2) ** 2
    D = np.zeros((3, 3))
    mu = 10
    alpha = 3
    beta = 2
    counter = 1
    err = 10 * tol
    max_iter = int(os.getenv("FIG9_18_MAX_ITER", "10000"))
    while err > tol and counter < max_iter:  # & counter<200
        counter += 1
        xx = x.copy()
        ff0 = ff

        J = Jac(x)
        JJ = J.T @ J
        gk = 2 * J.T @ f
        D[0, 0] = 2 * JJ[0, 0]
        D[1, 1] = 2 * JJ[1, 1]
        D[2, 2] = 2 * JJ[2, 2]
        S = 2 * JJ + mu * D
        dx = -np.linalg.solve(S, gk)
        x = xx + dx
        f = F(x)
        ff = np.linalg.norm(f, 2) ** 2

        if ff < ff0:
            mu = mu / beta
            xpoints.append(x[0])
            ypoints.append(x[1])
            with np.errstate(divide="ignore", invalid="ignore"):
                err = np.linalg.norm(-np.linalg.solve(S, gk) / x, np.inf)
        else:
            mu = alpha * mu
            x = xx
            ff = ff0

    print("counter =", counter)

    fig = matlab_figure(None, 421, 395)
    ax = fig.add_subplot(111)
    set_color_order(ax)

    ax.plot(xd, yd, "bo", markersize=8, linewidth=2)
    ax.set_aspect("equal", adjustable="box")
    nt = 100
    t = np.linspace(0, 2 * np.pi, nt)
    a = x[0]
    b = x[1]
    phi = x[2]
    print("a =", a)
    print("b =", b)
    print("phi =", phi)
    xc = np.zeros(nt)
    yc = np.zeros(nt)
    for i in range(nt):
        xc[i] = a * np.cos(phi) * np.cos(t[i]) - b * np.sin(phi) * np.sin(t[i])
        yc[i] = b * np.cos(phi) * np.sin(t[i]) + a * np.sin(phi) * np.cos(t[i])
    ax.plot(xc, yc, "r", linewidth=2)
    ax.plot([0, 0], [-2, 2], "k", linewidth=0.5)
    ax.plot([-2, 2], [0, 0], "k", linewidth=0.5)
    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")
    style_axes(ax, 14)
    fig.tight_layout()

    # exportgraphics(gcf,'/Users/mark/Desktop/ellipse2.eps')
    finish_figure(fig, "ellipse2.png")
    plt.show()


if __name__ == "__main__":
    fig9_18()
