import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from scipy.stats import pearsonr


def main():
    data = pd.read_csv('engineer.csv')
    colX=['Atmosphere', 'lecture quality', 'paper supporting', 'payment',
       'personality']
    colY=[ 'N corresponding author', 'citation index', 'alumni',
       'alumni(dr)', 'alumni(ms)', 'avg semesters(dr)']

    dataX=data[colX]
    dataY=data[colY]

    #print(dataX)
    #corr = dataX.corrwith(dataY)
    corr= pd.concat([dataX, dataY], axis=1).corr()
    corr = corr[colY].ix[colX]
    print(corr)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(corr, cmap='coolwarm', vmin=-1, vmax=1)
    fig.colorbar(cax)
    xticks=np.arange(6)
    yticks=np.arange(5)

    ax.set_xticks(xticks)
    plt.xticks(rotation=90)
    ax.set_yticks(yticks)
    ax.set_xticklabels(colY)
    ax.set_yticklabels(colX)
    print(corr.columns)
    plt.show()
    corr.to_csv("correlation.csv", mode='w')
if __name__ =='__main__':
    main()