from __future__ import print_function
import json
import numpy as np

from networkx.readwrite import json_graph
from argparse import ArgumentParser

from sklearn.linear_model import SGDRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_percentage_error, r2_score, explained_variance_score
from sklearn.dummy import DummyRegressor

''' To evaluate the embeddings, we run a logistic regression.
Run this script after running unsupervised training.
Baseline of using features-only can be run by setting data_dir as 'feat'
Example:
  python eval_scripts/lkpg-osm_eval.py ../data/lkpg-osm unsup-lkpg-osm/n2v_big_0.000010 test
'''

def run_dummy_regressor(train_embeds, train_targets, test_embeds, test_targets, scaler=None):
    dummy = DummyRegressor(strategy="median")
    dummy.fit(train_embeds, train_targets)  

    if scaler == None:
        rmse= mean_squared_error(test_targets, dummy.predict(test_embeds), squared=False)
    else:
        rmse = mean_squared_error(test_targets, dummy.predict(test_embeds), squared=False)
        rmse = scaler.inverse_transform(np.array(rmse).reshape(1,-1))[0][0]
    print('Dummy regressor RMSE:', rmse)
    print('Dummy regressor MAPE:', mean_absolute_percentage_error(test_targets, dummy.predict(test_embeds)))
    print('Dummy regressor R2:', r2_score(test_targets, dummy.predict(test_embeds)))

def run_regression(train_embeds, train_targets, test_embeds, test_targets, scaler=None):
    

    print('Ravelling train targets...')
    train_targets_ravelled = train_targets.ravel()

    print('Fitting squared loss Regressor...')
    print(np.shape(train_embeds), np.shape(train_targets_ravelled))
    regressor = SGDRegressor(loss="squared_loss", max_iter=1000, tol=1e-3)
    regressor.fit(train_embeds, train_targets_ravelled)

    print('Predicting outputs...')
    tr_res = []
    if scaler == None:
        tr_res.append(mean_squared_error(train_targets, regressor.predict(train_embeds), squared=False))
    else:
        rmse = mean_squared_error(train_targets, regressor.predict(train_embeds), squared=False)
        tr_res.append(scaler.inverse_transform(np.array(rmse).reshape(1,-1))[0][0])
    tr_res.append(mean_absolute_percentage_error(train_targets, regressor.predict(train_embeds)))
    tr_res.append(r2_score(train_targets, regressor.predict(train_embeds)))
    tr_res.append(explained_variance_score(train_targets, regressor.predict(train_embeds)))

    ts_res = []
    if scaler == None:
        ts_res.append(mean_squared_error(test_targets, regressor.predict(test_embeds), squared=False))
    else:
        rmse = mean_squared_error(test_targets, regressor.predict(test_embeds), squared=False)
        ts_res.append(scaler.inverse_transform(np.array(rmse).reshape(1,-1))[0][0])
    ts_res.append(mean_absolute_percentage_error(test_targets, regressor.predict(test_embeds)))
    ts_res.append(r2_score(test_targets, regressor.predict(test_embeds)))
    ts_res.append(explained_variance_score(test_targets, regressor.predict(test_embeds)))

    return tr_res, ts_res


def calculate_results(tr_res, ts_res, setting):
    prefixes = ['train', setting]
    results = [tr_res, ts_res]
    dicts = [{}, {}]
    cols = ['RMSE','MAPE', 'R2', 'expl_var']#, 'accuracy', 'precision', 'recall']
    for prefix, result, dict in zip(prefixes, results, dicts):
        for i, col in enumerate(cols):
            dict[prefix + '_' + col + '_avg'] = np.mean(result[:, i])
            #TODO: Add stds back in
            dict[prefix + '_' + col + '_std'] = np.std(result[:, i])
    return dicts


def evaluate(train_prefix, embed_dir, *, setting ='val',n_iter=100, scaling=False):
    print("Loading data...")
    G = json_graph.node_link_graph(json.load(open(train_prefix + "-G.json")))
    targets = json.load(open(train_prefix + "-target_map.json"))
    # targets = {int(i): l for i, l in targets.iteritems()}
    targets = {int(i): float(l) for i, l in targets.items()}
    
    train_ids = [n for n in G.nodes() if not G.node[n]['val'] and not G.node[n]['test']]
    test_ids = [n for n in G.nodes() if G.node[n][setting]]
    train_targets = np.array([targets[i] for i in train_ids])
    if train_targets.ndim == 1:
        train_targets = np.expand_dims(train_targets, 1)
    test_targets = np.array([targets[i] for i in test_ids])
    if test_targets.ndim == 1:
        test_targets = np.expand_dims(test_targets, 1)

    if scaling:
        print('Scaling targets...')
        from sklearn.preprocessing import StandardScaler, MinMaxScaler
        target_scaler = MinMaxScaler()
        target_scaler.fit(train_targets)
        train_targets = target_scaler.transform(train_targets)
        test_targets = target_scaler.transform(test_targets)
    else:
        target_scaler=None
    
    print('train targets:', np.shape(train_targets))
    print('test targets:', np.shape(test_targets))

    if embed_dir == "feat":
        print("Using only features..")
        feats = np.load(train_prefix + "-feats.npy")
        #TODO: Test using histograms only
        # feats = feats[:,:19]
        # print('----- FEATURES----')
        # for val in feats[0]:
        #     print(val)
        # print('^^^^^ FEATURES ^^^^^')
        # ## Logistic gets thrown off by big counts, so log transform num comments and score
        # feats[:, 0] = np.log(feats[:, 0] + 1.0)
        # feats[:, 1] = np.log(feats[:, 1] - min(np.min(feats[:, 1]), -1))
        feat_id_map = json.load(open( train_prefix + "-id_map.json"))
        feat_id_map = {int(id): val for id, val in feat_id_map.items()}

        train_feats = feats[[feat_id_map[id] for id in train_ids]]
        test_feats = feats[[feat_id_map[id] for id in test_ids]]
        print('train features:',np.shape(train_feats))
        print('test features:',np.shape(test_feats))

        print('Scaling features...')
        from sklearn.preprocessing import StandardScaler
        scaler = StandardScaler()
        # scaler.fit(train_feats)
        # train_feats = scaler.transform(train_feats)
        # test_feats = scaler.transform(test_feats)
        # Exclude histogram features from scaling!
        scaler.fit(train_feats[:,:19])
        train_feats[:,:19] = scaler.transform(train_feats[:,:19])
        test_feats[:,:19] = scaler.transform(test_feats[:,:19])
        
        print("Running regression...")
        run_dummy_regressor(train_feats, train_targets, test_feats, test_targets, target_scaler)

        all_tr_res = []
        all_ts_res = []
        for i in range(n_iter):
            tr_res, ts_res = run_regression(train_feats, train_targets, test_feats, test_targets, target_scaler)
            all_tr_res.append(tr_res)
            all_ts_res.append(ts_res)
            # print('(%d/%d) %s_f1_avg %.4f' % (i, n_iter, setting, np.mean(np.array(all_ts_res)[:,0])), end='\r')
    else:
        embeds = np.load(embed_dir + "/val.npy")
        id_map = {}
        with open(embed_dir + "/val.txt") as fp:
            for i, line in enumerate(fp):
                id_map[int(line.strip())] = i
        train_embeds = embeds[[id_map[id] for id in train_ids]]
        test_embeds = embeds[[id_map[id] for id in test_ids]]

        print("Running regression...")

        all_tr_res = []
        all_ts_res = []
        for i in range(n_iter):
            tr_res, ts_res = run_regression(train_embeds, train_targets, test_embeds, test_targets, target_scaler)
            all_tr_res.append(tr_res)
            all_ts_res.append(ts_res)
            # print('(%d/%d) %s_f1_avg %.4f' % (i, n_iter, setting, np.mean(np.array(all_ts_res)[:,0])), end='\r')

    tr_res_dict, ts_res_dict = calculate_results(np.array(all_tr_res), np.array(all_ts_res), setting)
    return tr_res_dict, ts_res_dict

if __name__ == '__main__':
    #TODO: remove the fix seed
    np.random.seed(1)
    parser = ArgumentParser("Run evaluation on osm regression data.")
    parser.add_argument("train_prefix", help="Path to directory containing the dataset.")
    parser.add_argument("embed_dir", help="Path to directory containing the learned node embeddings. Set to 'feat' for raw features.")
    parser.add_argument("setting", help="Either val or test.")
    parser.add_argument("n_iter", help="Number of re-runs (for avg and std).")
    parser.add_argument("scaling", help="Wheter to scale target values.")
    args = parser.parse_args()
    train_prefix = args.train_prefix
    embed_dir = args.embed_dir
    setting = args.setting
    n_iter = int(args.n_iter)
    scaling = True if args.scaling == "true" else False
    print('train prefix:', train_prefix)
    print('embed dir:', embed_dir)
    print('setting:', setting)
    print('n iter:', n_iter, type(n_iter))
    print('scaling:', scaling)
    print('\n\n')

    tr_res_dict, ts_res_dict = evaluate(train_prefix, embed_dir, setting=setting, n_iter=n_iter, scaling=scaling)
    print('Train stats')
    print(tr_res_dict)
    print(setting + ' stats')
    print(ts_res_dict)

