import json
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor
from xgboost import XGBRegressor
from catboost import CatBoostRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error as MSE, r2_score
import math

# with open('../../data/output/training_data/annual_mean_training_dataset_08-11-2020.json', 'r') as f:
#     data = json.load(f)

all_zones_df = pd.read_csv("../data/scsb_all_zones.csv")
zone_25_df = pd.read_csv("../data/scsb_zone_25.csv")
zone_26_df = pd.read_csv("../data/scsb_zone_26.csv")
zone_27_df = pd.read_csv("../data/scsb_zone_27.csv")

month_dependant_variables = ['jan_dist','feb_dist','mar_dist','apr_dist','may_dist','jun_dist','jul_dist','aug_dist','sep_dist','oct_dist','nov_dist','dec_dist']
month_labels = [x[0:3] for x in month_dependant_variables]

data = zone_26_df

xgb_results = []
rfr_results = []
dtr_results = []

# calculate monthly estimations for 3 models
for dependant_month in month_dependant_variables:
    features_df = data[['median_elevation', 'glacial_coverage', 'annual_precipitation', 'potential_evapo_transpiration', dependant_month]]
    X = features_df.drop([dependant_month], axis=1)
    y = features_df.get(dependant_month)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.30, random_state=42)
    
    xgb = XGBRegressor(random_state=42)
    xgb.fit(X_train, y_train)
    xgb_results.append(xgb.predict(X))

    rfr = RandomForestRegressor(random_state=42)
    rfr.fit(X_train, y_train)
    rfr_results.append(rfr.predict(X))

    dtr = DecisionTreeRegressor(random_state=42)
    dtr.fit(X_train, y_train)
    dtr_results.append(dtr.predict(X))

# compare the outputs of scsb against the 3 models
for row_target_index in range(20):
    xgb_row = []
    rfr_row = []
    dtr_row = []

    for month in range(12):
        xgb_row.append(xgb_results[month][row_target_index])
        rfr_row.append(rfr_results[month][row_target_index])
        dtr_row.append(dtr_results[month][row_target_index])

    plt.plot(data[month_dependant_variables].iloc[row_target_index], '-', label='scsb', color='blue', alpha=0.5)
    plt.plot(xgb_row, '-', label='xgboost', color='red', alpha=0.5)
    plt.plot(rfr_row, '-', label='randomforest', color='green', alpha=0.5)
    plt.plot(dtr_row, '-', label='decisiontree', color='purple', alpha=0.5)
    plt.legend(loc='best')
    plt.xticks(month_dependant_variables, month_labels)
    plt.xlabel('Month')
    plt.ylabel('Monthly Distribution')
    name = data['name'].iloc[row_target_index]
    plt.title(name)
    plt.savefig('../plots/{}.png'.format(name))
    plt.show()

[((559, 600), 'pandas.read_csv', 'pd.read_csv', (['"""../data/scsb_all_zones.csv"""'], {}), "('../data/scsb_all_zones.csv')\n", (570, 600), True, 'import pandas as pd\n'), ((614, 653), 'pandas.read_csv', 'pd.read_csv', (['"""../data/scsb_zone_25.csv"""'], {}), "('../data/scsb_zone_25.csv')\n", (625, 653), True, 'import pandas as pd\n'), ((667, 706), 'pandas.read_csv', 'pd.read_csv', (['"""../data/scsb_zone_26.csv"""'], {}), "('../data/scsb_zone_26.csv')\n", (678, 706), True, 'import pandas as pd\n'), ((720, 759), 'pandas.read_csv', 'pd.read_csv', (['"""../data/scsb_zone_27.csv"""'], {}), "('../data/scsb_zone_27.csv')\n", (731, 759), True, 'import pandas as pd\n'), ((1420, 1474), 'sklearn.model_selection.train_test_split', 'train_test_split', (['X', 'y'], {'test_size': '(0.3)', 'random_state': '(42)'}), '(X, y, test_size=0.3, random_state=42)\n', (1436, 1474), False, 'from sklearn.model_selection import train_test_split\n'), ((1491, 1520), 'xgboost.XGBRegressor', 'XGBRegressor', ([], {'random_state': '(42)'}), '(random_state=42)\n', (1503, 1520), False, 'from xgboost import XGBRegressor\n'), ((1601, 1639), 'sklearn.ensemble.RandomForestRegressor', 'RandomForestRegressor', ([], {'random_state': '(42)'}), '(random_state=42)\n', (1622, 1639), False, 'from sklearn.ensemble import RandomForestRegressor\n'), ((1720, 1758), 'sklearn.tree.DecisionTreeRegressor', 'DecisionTreeRegressor', ([], {'random_state': '(42)'}), '(random_state=42)\n', (1741, 1758), False, 'from sklearn.tree import DecisionTreeRegressor\n'), ((2183, 2296), 'matplotlib.pyplot.plot', 'plt.plot', (['data[month_dependant_variables].iloc[row_target_index]', '"""-"""'], {'label': '"""scsb"""', 'color': '"""blue"""', 'alpha': '(0.5)'}), "(data[month_dependant_variables].iloc[row_target_index], '-', label\n    ='scsb', color='blue', alpha=0.5)\n", (2191, 2296), True, 'import matplotlib.pyplot as plt\n'), ((2296, 2359), 'matplotlib.pyplot.plot', 'plt.plot', (['xgb_row', '"""-"""'], {'label': '"""xgboost"""', 'color': '"""red"""', 'alpha': '(0.5)'}), "(xgb_row, '-', label='xgboost', color='red', alpha=0.5)\n", (2304, 2359), True, 'import matplotlib.pyplot as plt\n'), ((2364, 2434), 'matplotlib.pyplot.plot', 'plt.plot', (['rfr_row', '"""-"""'], {'label': '"""randomforest"""', 'color': '"""green"""', 'alpha': '(0.5)'}), "(rfr_row, '-', label='randomforest', color='green', alpha=0.5)\n", (2372, 2434), True, 'import matplotlib.pyplot as plt\n'), ((2439, 2510), 'matplotlib.pyplot.plot', 'plt.plot', (['dtr_row', '"""-"""'], {'label': '"""decisiontree"""', 'color': '"""purple"""', 'alpha': '(0.5)'}), "(dtr_row, '-', label='decisiontree', color='purple', alpha=0.5)\n", (2447, 2510), True, 'import matplotlib.pyplot as plt\n'), ((2515, 2537), 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': '"""best"""'}), "(loc='best')\n", (2525, 2537), True, 'import matplotlib.pyplot as plt\n'), ((2542, 2593), 'matplotlib.pyplot.xticks', 'plt.xticks', (['month_dependant_variables', 'month_labels'], {}), '(month_dependant_variables, month_labels)\n', (2552, 2593), True, 'import matplotlib.pyplot as plt\n'), ((2598, 2617), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Month"""'], {}), "('Month')\n", (2608, 2617), True, 'import matplotlib.pyplot as plt\n'), ((2622, 2656), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Monthly Distribution"""'], {}), "('Monthly Distribution')\n", (2632, 2656), True, 'import matplotlib.pyplot as plt\n'), ((2708, 2723), 'matplotlib.pyplot.title', 'plt.title', (['name'], {}), '(name)\n', (2717, 2723), True, 'import matplotlib.pyplot as plt\n'), ((2776, 2786), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (2784, 2786), True, 'import matplotlib.pyplot as plt\n')]