from sklearn.linear_model import RidgeCV

from scipy.stats.stats import pearsonr

import os

from matplotlib import pyplot as plt

def run_ridge_main(X_tr, y_tr, X_vl, y_val, output, main):

    os.chdir(output)

    if not os.path.exists("ridge"):

            os.makedirs("ridge")

            dir = output + "ridge/"

    os.chdir(os.path.join(output, 'ridge/'))

    if not os.path.exists("figures"):

        os.makedirs("figures")

        dir = output + "ridge/figures/"

    os.chdir(os.path.join(output, 'ridge/figures/'))

    for i in range(0, y_tr.shape[1]):

        ridge = RidgeCV(alphas=[1e-3, 1e-2, 1e-1, 1]).fit(X_tr, y_tr[:,i])

        ridge.score(X_tr, y_tr[:, i])

        y_hat = ridge.predict(X_vl)

        # correlation btw predicted and observed

        corr = pearsonr(y_hat, y_val[:,i])

        fig = plt.figure()

        # plot observed vs. predicted targets

        plt.title('Ridge: Observed vs Predicted Y_trait_' + str(i) + 'cor:' + str(corr[0]))

        plt.ylabel('Predicted')

        plt.xlabel('Observed')

        plt.scatter(y_val[:,i], y_hat, marker='o')

        fig.savefig("Ridge_Out" + str(i) + '.png', dpi=300)

        plt.close(fig)