from Schenk_Brandenberger_S7_Aufg1 import Schenk_Brandenberger_S6_Aufg2
import numpy as np
import matplotlib.pyplot as plt



if __name__ == '__main__':
    # Aufgabe 3a
    x_plt_label_offset = 1997
    x = np.array([0, 2, 9, 13])
    y = np.array([150, 104, 172, 152])

    A = np.array([[x[0]**3, x[0]**2, x[0]**1, x[0]**0],
         [x[1]**3, x[1]**2, x[1]**1, x[1]**0],
         [x[2]**3, x[2]**2, x[2]**1, x[2]**0],
         [x[3]**3, x[3]**2, x[3]**1, x[3]**0]])
    b = np.array([[y[0]],
         [y[1]],
         [y[2]],
         [y[3]]])


    p1 = Schenk_Brandenberger_S6_Aufg2(A, b)[0]
    x_plt1_min = np.min(x)
    x_plt1_max = np.max(x)
    x_plt1_steps = (float(x_plt1_max) - float(x_plt1_min)) / 100.0
    x_plt = np.arange(x_plt1_min, x_plt1_max + x_plt1_steps, x_plt1_steps)
    plt.plot(x_plt + x_plt_label_offset, np.polyval(p1, x_plt), label="Aufgabe 3a")


    # Aufgabe 3b
    print("Schätzwert Aufgabe 3b für das Jahr 2003:", np.polyval(p1, 2003 - x_plt_label_offset))
    print("Schätzwert Aufgabe 3b für das Jahr 2004:", np.polyval(p1, 2004 - x_plt_label_offset))

    print()
    # Aufgabe 3c
    p2 = np.polyfit(x, y, 3)
    x_plt2_min = np.min(x)
    x_plt2_max = np.max(x)
    x_plt2_steps = (float(x_plt2_max) - float(x_plt2_min)) / 100.0
    x_plt = np.arange(x_plt2_min, x_plt2_max + x_plt2_steps, x_plt2_steps)
    plt.plot(x_plt + x_plt_label_offset, np.polyval(p2, x_plt), label="Aufgabe 3c")
    print("Schätzwert Aufgabe 3c für das Jahr 2003:", np.polyval(p2, 2003 - x_plt_label_offset))
    print("Schätzwert Aufgabe 3c für das Jahr 2004:", np.polyval(p2, 2004 - x_plt_label_offset))

    # Plots anzeigen
    plt.title("Aufgabe 3")
    plt.grid()
    plt.legend()
    plt.show()