"""Literal Python conversion of fig10_18.m."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from subaxis import subaxis


BASE_DIR = Path(__file__).resolve().parent


def _set_bold_tick_font(ax: plt.Axes, fontsize: int) -> None:
    ax.tick_params(labelsize=fontsize)
    ax.xaxis.label.set_size(fontsize)
    ax.yaxis.label.set_size(fontsize)
    ax.xaxis.label.set_fontweight("bold")
    ax.yaxis.label.set_fontweight("bold")
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontweight("bold")


def fig10_18() -> None:
    # y(t+dt)=Ax(t)

    # measure dt in days
    dt = 7

    # https://usafacts.org/visualizations/coronavirus-covid-19-spread-map/

    NY = np.loadtxt(BASE_DIR / "covid_data.txt")
    NY = NY.T
    N, M = NY.shape
    print("N, M =", N, M)

    # XD = new cases; YD = new cases at t+dt
    XD = NY[1 : N - dt, :] - NY[0 : N - dt - 1, :]
    YD = NY[dt + 1 : N, :] - NY[dt : N - 1, :]
    n, m = XD.shape
    print("n, m =", n, m)

    # data for table 10.5
    # NY[0:4,0:3]
    # XD[0:3,0:3]
    # YD[n-1,0:3]
    # XD[0:3,m-2:m]
    # YD[n-1,m-2:m]

    # use normalized data
    meansX = np.mean(XD, axis=0)
    meansY = np.mean(YD, axis=0)
    X = XD - meansX
    Y = YD - meansY
    XM = np.zeros(m)
    YM = np.zeros(m)
    for j in range(m):
        XM[j] = np.linalg.norm(X[:, j], np.inf)
        YM[j] = np.linalg.norm(Y[:, j], np.inf)
        X[:, j] = X[:, j] / XM[j]
        Y[:, j] = Y[:, j] / YM[j]

    G = X.T @ X
    GI = np.linalg.inv(G)
    P = Y.T @ X @ GI
    SSx = np.diag(1 / XM)
    Sy = np.diag(YM)
    A = Sy @ P @ SSx
    b = meansY.T - A @ meansX.T
    # np.linalg.norm(b, np.inf)

    # check on accuracy of y=Ax+b assumption
    # Albany = 1
    # Erie = 15
    # Rensselaer = 42
    # ulster = 56
    # saratoga=46, king=24, westchester=60

    r1 = 1
    r2 = 15
    r3 = 56

    fig1 = plt.figure(1, figsize=(7.2, 7.0))
    fig1.clf()
    # MATLAB figure position was [4 925 560 420]; figsize is scaled to fit the plot window.
    co = np.array(
        [
            [0, 0, 1],
            [0, 0.5, 0],
            [1, 0, 0],
            [0, 0.75, 0.75],
            [0.75, 0, 0.75],
            [0.75, 0.75, 0],
            [0.25, 0.25, 0.25],
        ]
    )

    ax1 = subaxis(3, 1, 1, 1, "MT", 0.01, "MB", 0.06, "MR", -0.001, "ML", 0.06, "P", 0.04)
    ax1.set_prop_cycle(color=co.tolist())
    r = r1
    r_index = r - 1
    yp = np.zeros(n)
    for i in range(n):
        yp[i] = np.dot(A[r_index, :], XD[i, :]) + b[r_index]
    max_yp = np.max(yp)
    print("max_yp =", max_yp)
    t = np.arange(dt + 1, n + dt + 1)
    ax1.plot(t, yp, "r", linewidth=1.5)
    ax1.plot(t, YD[:, r_index], "--b", linewidth=1.5)
    # ax1.set_xlabel('x-axis')
    ax1.set_ylabel("Albany", fontweight="bold")
    ax1.set_xlim([1, 275])
    ax1.set_ylim([0, 300])
    ax1.set_xticks([1, 91, 181, 274])
    ax1.set_xticklabels(["April 1", "June 1", "Oct 1", "Jan 1"])
    ax1.set_yticks([0, 150, 300])
    ax1.grid(True)
    ax1.set_frame_on(True)
    _set_bold_tick_font(ax1, 12)  # MATLAB FontSize 14 reduced by 2

    # error1=np.linalg.norm(yp-YD[:,r_index]+XD[:,r_index],np.inf)

    ax2 = subaxis(3, 1, 1, 2)
    ax2.set_prop_cycle(color=co.tolist())
    r = r2
    r_index = r - 1
    yp = np.zeros(n)
    for i in range(n):
        yp[i] = np.dot(A[r_index, :], XD[i, :]) + b[r_index]
    max_yp = np.max(yp)
    print("max_yp =", max_yp)
    ax2.plot(t, yp, "r", linewidth=1.5)
    ax2.plot(t, YD[:, r_index], "--b", linewidth=1.5)
    # ax2.set_xlabel('x-axis')
    ax2.set_ylabel("Erie", fontweight="bold")
    ax2.set_xlim([1, 275])
    ax2.set_ylim([0, 800])
    ax2.set_xticks([1, 91, 181, 274])
    ax2.set_xticklabels(["April 1", "June 1", "Oct 1", "Jan 1"])
    ax2.set_yticks([0, 400, 800])
    ax2.grid(True)
    ax2.set_frame_on(True)
    legend = ax2.legend([" Predicted", " Actual"], loc="upper center", fontsize=12)
    for text in legend.get_texts():
        text.set_fontweight("bold")
    _set_bold_tick_font(ax2, 12)  # MATLAB FontSize 14 reduced by 2

    # error2=np.linalg.norm(yp-YD[:,r_index]+XD[:,r_index],np.inf)

    ax3 = subaxis(3, 1, 1, 3)
    ax3.set_prop_cycle(color=co.tolist())
    r = r3
    r_index = r - 1
    yp = np.zeros(n)
    for i in range(n):
        yp[i] = np.dot(A[r_index, :], XD[i, :]) + b[r_index]
    ax3.plot(t, yp, "r", linewidth=1.5)
    ax3.plot(t, YD[:, r_index], "--b", linewidth=1.5)
    ax3.set_xlabel("t-axis", fontweight="bold")
    ax3.set_ylabel("Ulster", fontweight="bold")
    ax3.set_xlim([1, 275])
    ax3.set_ylim([0, 120])
    ax3.set_yticks([0, 60, 120])
    ax3.set_xticks([1, 91, 181, 274])
    ax3.set_xticklabels(["April 1", "June 1", "Oct 1", "Jan 1"])
    ax3.grid(True)
    ax3.set_frame_on(True)
    _set_bold_tick_font(ax3, 12)  # MATLAB FontSize 14 reduced by 2

    # error3=np.linalg.norm(yp-YD[:,r_index]+XD[:,r_index],np.inf)

    fig1.savefig(BASE_DIR / "covid.eps", format="eps", bbox_inches="tight")

    fig2 = plt.figure(2, figsize=(7.2, 3.2))
    fig2.clf()
    sings = np.linalg.svd(P, compute_uv=False)
    ratios = sings[0:4] / sings[0]
    print("sigma(i)/sigma(1) = " + "".join(f"{value:6.2f}, " for value in ratios))
    print()

    # MATLAB figure position was [4 647 592 204]; figsize is scaled to fit the plot window.
    ax4 = fig2.add_subplot(111)
    ax4.plot(np.arange(1, len(sings) + 1), sings / sings[0], "--o", linewidth=1.5, markersize=8)
    # ax4.axis([0, len(error), 0, 0.1])
    ax4.set_xlabel("k-axis", fontweight="bold")
    ax4.set_ylabel("sigma/sigma(1)", fontweight="bold")
    ax4.grid(True)
    ax4.set_frame_on(True)
    _set_bold_tick_font(ax4, 12)  # MATLAB FontSize 14 reduced by 2
    fig2.tight_layout()

    plt.show()


if __name__ == "__main__":
    fig10_18()
