from pypge.benchmarks import explicit

import numpy as np

# visualization libraries
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# plot the visuals in ipython
%matplotlib inline


"""
Explanation: Explicit 1D Benchmarks
This file demonstrates how to generate, plot, and output data for 1d benchmarks
Choose from:

Koza_01
Koza_02
Koza_03
Lipson_01
Lipson_02
Lipson_03
Nguyen_01
Nguyen_02
Nguyen_03
Nguyen_04
Nguyen_05
Nguyen_06
Nguyen_07
Nguyen_08

Imports
End of explanation
"""


# choose your problem here
prob = explicit.Koza_01()

# you can also specify the following params as keyword arguments
#
# params = {
#     'name': "Koza_01",
#     'xs_str': ["x"],
#     'eqn_str': "x**4 + x**3 + x**2 + x",
#     'xs_params': [ (-4.0,4.0) ],
#     'npts': 200,
#     'noise': 0.1
# }

# or make your own with the following
#
# explicit.Explicit_1D(params):


"""
Explanation: Generate the data with noise
End of explanation
"""


print prob['name'], prob['eqn']
print prob['xpts'].shape

fig = plt.figure()
fig.set_size_inches(16, 12)

plt.plot(prob['xpts'][0], prob['ypure'], 'r.')
plt.legend(loc='center left', bbox_to_anchor=(0.67, 0.12))

plt.title(prob['name'] + " Clean", fontsize=36)
plt.savefig("img/benchmarks/" + prob['name'].lower() + "_clean.png", dpi=200)
# plt.show()

### You can only do one of 'savefig()' or 'show()'

fig = plt.figure()
fig.set_size_inches(16, 12)

plt.plot(prob['xpts'][0], prob['ypts'], 'b.')
plt.legend(loc='center left', bbox_to_anchor=(0.67, 0.12))

plt.title(prob['name'] + " Noisy", fontsize=36)
plt.savefig("img/benchmarks/" + prob['name'].lower() + "_noisy.png", dpi=200)
# plt.show()


"""
Explanation: Plot inline and save image
End of explanation
"""


data = np.array([prob['xpts'][0], prob['ypts']]).T
print data.shape

cols = [['x', 'out']]
out_data = cols + data.tolist()
import json
json_out = json.dumps( out_data, indent=4)
# print json_out

f_json = open("data/benchmarks/" + prob['name'].lower() + ".json", 'w')
f_json.write(json_out)
f_json.close()

f_csv = open("data/benchmarks/" + prob['name'].lower() + ".csv", 'w')
for row in out_data:
    line = ", ".join([str(col) for col in row]) + "\n"
    f_csv.write(line)
f_csv.close()

"""
Explanation: Output json and csv data
End of explanation
"""


