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
from subaxis import subaxis


def _plot_fit(ax, DDB, V, IX, IY, xlabel, ylabel, fs, printf_format):
    # calculate best line
    slope = V[IY, 0] / V[IX, 0]
    xL = np.min(DDB[:, IX])
    yL = slope * xL
    xR = np.max(DDB[:, IX])
    yR = slope * xR
    ax.plot([xL, xR], [yL, yR], "r", linewidth=1.6)
    ax.plot(DDB[:, IX], DDB[:, IY], "ob", markersize=8, linewidth=2)
    LX = np.min(DDB[:, IX]); RX = np.max(DDB[:, IX]); LY = np.min(DDB[:, IY]); RY = np.max(DDB[:, IY])
    # LX=-0.2; RX=0.5; LY=-0.2; RY=0.6;
    ax.axis([LX, RX, LY, RY])
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, fs, "bold")

    # format shortG
    # mean squared prediction error
    MSPE = np.linalg.norm(slope * DDB[:, IX] - DDB[:, IY], 2) ** 2 / len(DDB[:, IY])
    # mean absolute deviation
    MAD = np.linalg.norm(slope * DDB[:, IX] - DDB[:, IY], 1) / len(DDB[:, IY])
    # coefficient of determination R^2
    R2 = (
        np.dot(slope * DDB[:, IX], DDB[:, IY])
        / (np.linalg.norm(DDB[:, IX], 2) * np.linalg.norm(DDB[:, IY], 2))
    ) ** 2
    print(printf_format.format(MSPE, MAD, R2))
    return np.array([MSPE, MAD, R2])


def fig10_5() -> None:
    #  this needs data file crime_data_2009.txt
    #  it also needs subaxis.m and parseArgs.m

    #  finds best line fit, then plots test data for variable for column IY
    #  as function of variable from column IX
    #  Columns: 1=Pop; 2=Murder; 3=Rape; 4=Robbery; 5=Assault; 6=Burglary;
    #          7=Larceny; 8=Motor vehicle theft

    # pick prediction variable
    IX = 2 - 1

    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)

    #  split using every other row A=training 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)

    # randomize rows then split D into training set A and testing set B
    # rows=np.random.permutation(m)
    # fac=2/3
    # mm=round(fac*m)
    # A=AA[rows[0:mm],:]
    # B=AA[rows[mm:m],:]
    # 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(n)
        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

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

    set_color_order()
    fig = matlab_figure([4, 715, 700, 630])

    # fontsize
    fs = 12

    #  IY choices used in different plots
    ax = subaxis(2, 2, 2, 2, "MT", 0.008, "MB", 0.06, "MR", 0.003, "ML", 0.042, "P", 0.03)
    assert ax is not None
    IY = 1 - 1
    errors_pop = _plot_fit(ax, DDB, V, IX, IY, "Murder", "Population", fs, "\n Population  &  {:3.2e}  &  {:3.2e}  &  {:10.3f}")

    ax = subaxis(2, 2, 1, 2)
    assert ax is not None
    IY = 4 - 1
    errors_ass = _plot_fit(ax, DDB, V, IX, IY, "Murder", "Robbery", fs, "\n Robbery  &  {:3.2e}  &  {:3.2e}  &  {:10.3f}")

    ax = subaxis(2, 2, 2, 1)
    assert ax is not None
    IY = 5 - 1
    errors_ass = _plot_fit(ax, DDB, V, IX, IY, "Murder", "Assault", fs, "\n Assault  &  {:3.2e}  &  {:3.2e}  &  {:10.3f}")

    ax = subaxis(2, 2, 1, 1)
    assert ax is not None
    IY = 8 - 1
    # calculate best line
    errors_ass = _plot_fit(ax, DDB, V, IX, IY, "Murder", "Vehicle", fs, "\n Vehicle  &  {:3.1e}  &  {:3.1e}  &  {:7.3f}\n")
    del errors_pop, errors_ass

    # fig.savefig('/Users/mark/Desktop/murder.eps')
    finish_figure(fig)
    plt.show()


if __name__ == "__main__":
    fig10_5()
