import numpy as np
import matplotlib 
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib import colors

def plot_loss(metrics,expName):
    plt.figure() 
    plt.plot(metrics['train'],color='blue',label='train')
    plt.plot(metrics['val'],color='red',label='val')
    plt.xlabel('iterations')
    plt.ylabel('L2 Loss')
    #plt.yscale('log')
    plt.legend(loc="best")
    plt.tight_layout() 
    plt.savefig('loss_%s.pdf' %expName)

def plot_error(y_pred,y_from_ds,expName):
    plt.figure() 
    plt.hist([y_pred[:,0],y_from_ds[:,0]],label=['estimate','true'],bins=int(np.sqrt(len(y_pred[:,0]))),color=['red','blue'],alpha=0.5)
    #plt.hist(y_from_ds[:,0], label='true',  bins=int(np.sqrt(len(y_from_ds[:,0]))),color='blue',alpha=0.5)
    plt.xlabel('x coord')
    plt.ylabel('frequency')
    plt.xlim([0.4,0.55])
    plt.legend() 
    plt.tight_layout()
    plt.savefig('x_coord_kde_%s.pdf' %expName)

    plt.figure() 
    plt.hist([y_pred[:,1],y_from_ds[:,1]],label=['estimate','true'],bins=int(np.sqrt(len(y_pred[:,1]))),color=['red','blue'],alpha=0.5)
    #plt.hist(y_from_ds[:,1], label='true',  bins=int(np.sqrt(len(y_from_ds[:,1]))),color='blue',alpha=0.5)
    plt.xlabel('y coord')
    plt.ylabel('frequency')
    plt.xlim([0.4,0.55])
    plt.legend()
    plt.tight_layout()
    plt.savefig('y_coord_kde_%s.pdf' %expName)
    
    plt.figure() 
    plt.hist(y_from_ds[:,0].squeeze()-y_pred[:,0].squeeze(),bins=int(np.sqrt(len(y_from_ds[:,0]))),label='error')
    plt.xlabel('X Coordinate Absolute Error')
    plt.ylabel('frequency')
    plt.xlim([-0.005,0.005])
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('x_coord_mse_%s.pdf' %expName)

    plt.figure() 
    plt.hist(y_from_ds[:,1].squeeze()-y_pred[:,1].squeeze(),bins=int(np.sqrt(len(y_from_ds[:,1]))),label='error')
    plt.xlabel('Y Coordinate Absolute Error')
    plt.ylabel('frequency')
    plt.xlim([-0.005,0.005])
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('y_coord_mse_%s.pdf' %expName)

    plt.figure() 
    plt.hist(np.sqrt(np.power(y_from_ds[:,0].squeeze()-y_pred[:,0].squeeze(),2)),bins=int(np.sqrt(len(y_from_ds[:,0]))),label='error')
    plt.xlabel('X Coordinate RMSE')
    plt.ylabel('frequency')
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('x_coord_rmse_%s.pdf' %expName)

    plt.figure() 
    plt.hist(np.sqrt(np.power(y_from_ds[:,1].squeeze()-y_pred[:,1].squeeze(),2)),bins=int(np.sqrt(len(y_from_ds[:,1]))),label='error')
    plt.xlabel('Y Coordinate RMSE')
    plt.ylabel('frequency')
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('y_coord_rmse_%s.pdf' %expName)

    plt.figure() 
    plt.scatter(y_from_ds[:,0],np.power(y_from_ds[:,0].squeeze()-y_pred[:,0].squeeze(),2))
    plt.xlabel('X Coordinate')
    plt.ylabel('Absolute Error')
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('x_coord_positional_mse_%s.pdf' %expName)

    plt.figure() 
    plt.scatter(y_from_ds[:,1],np.power(y_from_ds[:,1].squeeze()-y_pred[:,1].squeeze(),2))
    plt.xlabel('Y Coordinate')
    plt.ylabel('Absolute Error')
    #plt.yscale('log')
    plt.tight_layout()
    plt.savefig('y_coord_positional_mse_%s.pdf' %expName)

    plt.figure() 
    plt.hist2d(y_from_ds[:,0].squeeze()-y_pred[:,0].squeeze(),y_from_ds[:,1].squeeze()-y_pred[:,1].squeeze(),bins=int(np.sqrt(len(y_from_ds[:,0]))),norm=colors.LogNorm())
    plt.xlabel('X Coordinate Absolute Error')
    plt.ylabel('Y Coordinate Absolute Error')
    plt.colorbar()
    plt.tight_layout()
    plt.savefig('error_map_mse_%s.pdf' %expName)
    
    data_range_x = np.arange(0.42,0.52,50)
    data_range_y = np.arange(0.42,0.52,50)
    x,y = np.meshgrid(data_range_x,data_range_y)
    error = np.zeros(shape=(50,50))
    i,j=0,0
    for x_val in x:
        print(np.where(y_from_ds[:,0]==x_val))
    '''
        for y_val in y:
            error[i][j] = y_from_ds[np.where(y_from_ds[:,0]==x_val)[0],0] - y_pred[np.where(y_from_ds[:,0]==x_val)[0],0] + \
            y_from_ds[np.where(y_from_ds[:,1]==y_val)[0],1] - y_pred[np.where(y_from_ds[:,1]==y_val)[0],1]
            j +=1 
        i+=1 
    plt.figure() 
    plt.contourf(x,y,error)
    plt.tight_layout()
    plt.savefig('test_%s.pdf' %expName)
    '''
    data_range_x = np.linspace(min(y_from_ds[:,0]),max(y_from_ds[:,0]),len(y_from_ds[:,0]))
    data_range_y = np.linspace(min(y_from_ds[:,1]),max(y_from_ds[:,1]),len(y_from_ds[:,1]))
    x,y = np.meshgrid(data_range_x,data_range_y)
    error = np.zeros(shape=(len(y_from_ds[:,0]),len(y_from_ds[:,1])))
    print(np.shape(error))
    print(np.shape(data_range_x))
    print(np.shape(data_range_y))
    i,j=0,0
    for i in range(len(x)):
        for j in range(len(y)):
            error[i][j] = y_from_ds[i,0] - y_pred[i,0] + \
            y_from_ds[j,1] - y_pred[j,1]
    plt.figure() 
    plt.contourf(x,y,error)
    plt.tight_layout()
    plt.savefig('test_%s.pdf' %expName)


    plt.figure() 
    plt.scatter(y_from_ds[:,0],y_from_ds[:,0].squeeze()-y_pred[:,0].squeeze(),bins=int(np.sqrt(len(y_from_ds[:,0]))),norm=colors.NoNorm())
    plt.xlabel('X Coordinate Absolute Error')
    plt.ylabel('Y Coordinate Absolute Error')
    plt.colorbar()
    plt.tight_layout()
    plt.savefig('error_map_mse_no_norm_%s.pdf' %expName)