import numpy as np
import itertools
# A small data set
#T_in = np.array([20,24,35,40])+273.15
#x_in = np.array([47,55,70,78,82])/100.0
#rho_in = np.array([[1047,1033,1020,1000,990],[997,983,970,950,940],[947,933,920,900,895],[987,883,870,850,845]])
#
# large data set
T_in   = np.array([   -45 ,    -40 ,    -35 ,    -30 ,    -25 ,    -20 ,    -15 ,    -10])+273.15 # Kelvin
x_in   = np.array([     5 ,     10 ,     15 ,     20 ,     25 ,     30 ,     35 ])/100.0 # mass fraction
        
rho_in = np.array([
          [1064.0,    1054.6,    1045.3,    1036.3,    1027.4,    1018.6,    1010.0],
          [1061.3,    1052.1,    1043.1,    1034.3,    1025.6,    1017.0,    1008.6],
          [1057.6,    1048.8,    1040.1,    1031.5,    1023.1,    1014.8,    1006.7],
          [1053.1,    1044.6,    1036.2,    1028.0,    1019.9,    1012.0,    1004.1],
          [1047.5,    1039.4,    1031.5,    1023.7,    1016.0,    1008.4,    1000.9],
          [1040.7,    1033.2,    1025.7,    1018.4,    1011.2,    1004.0,     997.0],
          [1032.3,    1025.3,    1018.5,    1011.7,    1005.1,     998.5,     992.0],
          [1021.5,    1015.3,    1009.2,    1003.1,     997.1,     991.2,     985.4]]) # kg/m3

def getCoeffs1d(x,z,order):
    if (len(x)<order+1): 
        raise ValueError("You have only {0} elements and try to fit {1} coefficients, please reduce the order.".format(len(x),order+1))
    A = np.vander(x,order+1)[:,::-1]
    #Anew = np.dot(A.T,A)
    #znew = np.dot(A.T,z)
    #coeffs = np.linalg.solve(Anew, znew)
    coeffs = np.linalg.lstsq(A, z)[0]
    return coeffs

xorder = 2
x = T_in[0:4]#T_in#T_in[0:4]#
z = rho_in[0:4,0]#rho_in#rho_in[0:4,0]#
coeffs = getCoeffs1d(x,z,xorder)
print coeffs
zf = np.polynomial.polynomial.polyval(x,coeffs)
print z[0]
print zf[0]
print z[-1]
print zf[-1]

def getCoeffs2dmatrix(x_in,y_in,z_in,x_order,y_order):
    
    x_order += 1
    y_order += 1
    
    #To avoid overfitting, we only use the upper left triangle of the coefficient matrix
    x_exp = range(x_order)
    y_exp = range(y_order)
    limit = max(x_order,y_order)
    
    xy_exp = []
    
    # Construct the upper left triangle of coefficients    
    for i in x_exp:
        for j in y_exp:
            if(i+j<limit): xy_exp.append((i,j))
                
    x_num = len(x_in)
    y_num = len(y_in)
                
    cols = len(xy_exp)
    eqns = x_num * y_num
    #if (eqns<cols):
    #    raise ValueError("You have only {0} equations and try to fit {1} coefficients, please reduce the order.".format(eqns,cols))   
    if (x_num<x_order):
        raise ValueError("You have only {0} x-entries and try to fit {1} x-coefficients, please reduce the x_order.".format(x_num,x_order))
    if (y_num<y_order):
        raise ValueError("You have only {0} y-entries and try to fit {1} y-coefficients, please reduce the y_order.".format(y_num,y_order))
    
    #Create functional matrix
    A = np.zeros((x_num,y_num,x_order,y_order))
    for i in range(x_num):
        for j in range(y_num):
            for (xk,yk) in xy_exp:
                A[i][j][xk][yk] = x[i]**xk * y[j]**yk
                
    raise NotImplementedError("No 4-dimensional solver implemented")
    
    
def getCoeffs2d(x_in,y_in,z_in,x_order,y_order):
    
    x_order += 1
    y_order += 1
    
    #To avoid overfitting, we only use the upper left triangle of the coefficient matrix
    x_exp = range(x_order)
    y_exp = range(y_order)
    limit = max(x_order,y_order)
    
    xy_exp = []
    
    # Construct the upper left triangle of coefficients    
    for i in x_exp:
        for j in y_exp:
            if(i+j<limit): xy_exp.append((i,j))
    
    # Construct input pairs   
    xx, yy = np.meshgrid(x_in,y_in,indexing='ij')
    xx = np.array(xx.flat)
    yy = np.array(yy.flat)
    zz = np.array(z_in.flat)
    
    x_num = len(x_in)
    y_num = len(y_in)
                
    cols = len(xy_exp)
    eqns = x_num * y_num
    #if (eqns<cols):
    #    raise ValueError("You have only {0} equations and try to fit {1} coefficients, please reduce the order.".format(eqns,cols))   
    if (x_num<x_order):
        raise ValueError("You have only {0} x-entries and try to fit {1} x-coefficients, please reduce the x_order.".format(x_num,x_order))
    if (y_num<y_order):
        raise ValueError("You have only {0} y-entries and try to fit {1} y-coefficients, please reduce the y_order.".format(y_num,y_order))
    
    # Build the functional matrix
    A = np.zeros((eqns,cols))
    for i in range(eqns): # row loop
        for j, (xj,yj) in enumerate(xy_exp): # makes columns
            A[i][j] = xx[i]**xj * yy[i]**yj
             
    coeffs = np.linalg.lstsq(A, zz)[0]
    
    #Rearrange coefficients to a matrix shape
    C = np.zeros((x_order,y_order))
    for i, (xi,yi) in enumerate(xy_exp): # makes columns
        C[xi][yi] = coeffs[i]
        
    return C

xorder = 3
yorder = 3

x = T_in
y = x_in
z = rho_in

coeffs = getCoeffs2d(x,y,z,xorder,yorder)
print coeffs

print z[0][0]
print np.polynomial.polynomial.polyval2d(x[0],y[0],coeffs)

print z[-1][-1]
print np.polynomial.polynomial.polyval2d(x[-1],y[-1],coeffs)

%pylab inline
import matplotlib
matplotlib.rcParams['savefig.dpi'] = 2 * matplotlib.rcParams['savefig.dpi']
from mpl_toolkits.mplot3d import Axes3D

# Construct input pairs   
xx, yy = np.meshgrid(x,y,indexing='ij')
xx = np.array(xx.flat)
yy = np.array(yy.flat)
zz = np.array(z.flat)

zf = np.polynomial.polynomial.polyval2d(xx,yy,coeffs)

fig  = plt.figure()
ax   = plt.axes(projection='3d')
ax.scatter(xx, yy, zz, color='blue', label='Original data')
ax.scatter(xx, yy, zf, color='red',  label='Fitted data')
#ax.plot(xx[0], yy[0], zs=zz[0], label='Fitted line', zdir='z')
ax.legend()