from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from plot_helpers import finish_figure, matlab_figure, require_data_file, set_axes_font, set_color_order


def fig10_6() -> None:
    #  this needs data file crime_data_2009.txt

    #  finds best 7D-planar fit, where variable from column IC is planar
    #  function of the other 7 variables. After this, it compares test data
    #  for column IC (using its 7 variables values) with the
    #  plane prediction

    # Columns: 1=Pop; 2=Murder; 3=Rape; 4=Robbery; 5=Assault; 6=Burglary;
    #          7=Larceny; 8=Vehicle

    # pick variable IC=2,8
    IC = 2

    data_path = require_data_file("crime_data_2009.txt", __file__)
    crime_data_2009 = np.loadtxt(data_path)
    AA = crime_data_2009[:, :]
    m, n = AA.shape
    print(m, n)
    means = np.mean(AA, axis=0)
    print(means)

    # interchange IC column with first column
    AA[:, [0, IC - 1]] = AA[:, [IC - 1, 0]]

    #  A=traing data   B = testing data
    A = AA[0 : m - 1 : 2, :].copy()
    B = AA[1:m:2, :].copy()
    m, n = A.shape
    print(m, n)
    means = np.mean(A, axis=0)
    print(means)

    # normalization of A
    D = np.zeros((m, n))
    for j in range(n):
        D[:, j] = A[:, j] - means[j] * np.ones(m)
    DD = np.zeros((m, n))
    DM = np.zeros(n)
    for j in range(n):
        # DM(j)=max(abs(D(:,j)))
        DM[j] = np.linalg.norm(D[:, j], 2) / np.sqrt(m)
        DD[:, j] = D[:, j] / DM[j]

    #  DD=U*S*V'
    U, singular_values, Vh = np.linalg.svd(DD, full_matrices=False)
    V = Vh.T
    S = np.diag(singular_values)
    ss = S.shape
    print(ss)

    id = np.arange(1, ss[1] + 1)
    sigmas = np.array([S[i, i] for i in range(ss[1])])
    # print(f' {i:d}   {S[i,i]:19.14e}')
    print(sigmas)
    print(V[:, 0])
    del id

    Q = np.column_stack([V[1:8, i] for i in range(7)])

    #  normalize testing data (using same normalization as used for A)
    M, N = B.shape
    print(M, N)
    DB = np.zeros((M, N))
    for j in range(N):
        DB[:, j] = B[:, j] - means[j] * np.ones(M)
    DDB = np.zeros((M, N))
    for j in range(N):
        DDB[:, j] = DB[:, j] / DM[j]

    print(DD[m - 1, :])
    print(DDB[m - 1, :])

    PP = np.zeros(M)
    AP = np.zeros(M)
    for i in range(M):
        alpha = np.linalg.solve(Q, DDB[i, 1:N])
        PP[i] = np.dot(alpha, V[0, 0:7])
        AP[i] = DDB[i, 0]

    set_color_order()
    fig = matlab_figure([1, 1046, 367, 299])
    ax = fig.add_subplot(111)
    ax.plot(AP, PP, "ob", markersize=8, linewidth=2)

    if IC == 8:
        ax.text(1.5, -0.5, "Vehicle", fontsize=14, fontweight="bold")
        LX = -1; RX = 3; LY = -1; RY = 3
    elif IC == 2:
        ax.text(1.5, -7, "Murder", fontsize=14, fontweight="bold")
        LX = -1; RX = 3; LY = -15; RY = 54
    else:
        raise ValueError("IC must be 2 or 8 to match the original plotting limits.")

    ax.axis([LX, RX, LY, RY])
    ax.plot([LX, RX], [LX, RX], "r", linewidth=1.6)

    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel("Training")
    ax.set_ylabel("Testing")
    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, 14, "bold")

    # if IC==8:
    #     fig.savefig('/Users/mark/Desktop/v.eps')
    # elif IC==2:
    #     fig.savefig('/Users/mark/Desktop/m.eps')
    finish_figure(fig)
    plt.show()


if __name__ == "__main__":
    fig10_6()
