__author__ = 'chunk'

import numpy as np
import math
import matplotlib.pyplot as plt
import seaborn as sns

from chap3.cholesky import calc_upper, calc_lower, cholesky

plt.ticklabel_format(style='sci', axis='both')


def fiiting(t, f, phi):
    m, n = len(t), len(phi)
    assert m == len(f)

    A = np.array([phi[j](t[i]) for i in range(m) for j in range(n)]).reshape(m, -1)
    G = np.dot(A.T, A)
    b = np.dot(A.T, f)
    L = cholesky(G)
    y = calc_lower(L, b)
    x = calc_upper(L.T, y)
    return x


def test0():
    t0 = [1, 2, 3, 4, 5]
    f0 = [4, 4.5, 6, 8, 8.5]
    phi = [lambda x: 1, lambda x: x]
    coef = fiiting(t0, f0, phi)
    print coef

    fit = lambda x: coef[0] + coef[1] * x * x
    t1 = np.linspace(1, 8, num=1000).tolist()
    f1 = [fit(i) for i in t1]

    plt.scatter(t0, f0)
    plt.plot(t1, f1)
    plt.show()


def test():
    # print  np.linspace(1, 8, num=15).tolist()



    t = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0]
    f = [33.40, 79.50, 122.65, 159.05, 189.15, 214.15, 238.65, 252.2, 267.55, 280.50, 296.65,
         301.65, 310.40,
         318.15, 325.15]
    phi = [lambda x: 1, lambda x: x, lambda x: x * x]
    coef = fiiting(t, f, phi)
    print coef
    fit = lambda x: coef[0] + coef[1] * x + coef[2] * x * x
    t1 = np.linspace(1, 8, num=1000).tolist()
    f1 = [fit(i) for i in t1]


    ff = [np.log(i) for i in f]
    print ff
    phi2 = [lambda x: 1, lambda x: x]
    coef2 = fiiting(t, ff, phi2)
    print coef2
    fit2 = lambda x: np.exp(coef2[0]) * np.exp(coef2[1] *  x)
    t2 = np.linspace(1, 8, num=1000).tolist()
    f2 = [fit2(i) for i in t2]

    tt = [1.0/i for i in t]
    ff = [np.log(i) for i in f]
    print ff
    phi2 = [lambda x: 1, lambda x: x]
    coef2 = fiiting(tt, ff, phi2)
    print coef2
    fit2 = lambda x: np.exp(coef2[0]) * np.exp(coef2[1] * 1.0 / x)
    t3 = np.linspace(1, 8, num=1000).tolist()
    f3 = [fit2(i) for i in t3]

    plt.scatter(t, f)
    plt.plot(t1, f1, label='polynomial')
    plt.plot(t2, f2, label='exponential', linestyle='--')
    plt.plot(t3, f3, label='exponential 1/t', linestyle='-')
    plt.xlabel("t")
    plt.ylabel("y")
    plt.legend(loc=2)
    plt.show()


if __name__ == '__main__':
    test()