import csv
import pandas as pd
import numpy as np

from scipy.interpolate import Rbf

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm

filename = 'Training_wells.csv'
data = pd.read_csv(filename, sep=',',header=None, names=['Well','X','Y','NTG'])
data = data.apply(pd.to_numeric, errors='coerce')
data = data[1:]

filename = 'Empty_part.csv'
data_test = pd.read_csv(filename, sep=',',header=None, names=['Well','X','Y'])
data_test = data_test.apply(pd.to_numeric, errors='coerce')
data_test = data_test[1:]

# 2-d tests - setup scattered data
x = list(data.X)
y = list(data.Y)
z = list(data.NTG)

if True:
    avg = sum(z) / len(z)

    x.append(250)
    y.append(900)
    z.append(0.45)

    x.append(250)
    y.append(932)
    z.append(avg)

    x.append(250)
    y.append(920)
    z.append(0.35)


    x.append(250)
    y.append(923)
    z.append(0.35)

    x.append(235)
    y.append(932)
    z.append(avg)

    x.append(212)
    y.append(906)
    z.append(0.20)

    x.append(240)
    y.append(905)
    z.append(0.35)


ti = np.linspace(200, 250, 101)
XI, YI = np.meshgrid(np.linspace(200, 250, 101), np.linspace(900, 950, 101))

# use RBF
rbf = Rbf(x, y, z, epsilon=2)
ZI = rbf(XI, YI)

# plot the result
fig, axes = plt.subplots(nrows = 1, ncols = 1)
fig.set_size_inches(13.5, 7.5)
plt.subplot(1, 1, 1)
plt.pcolor(XI, YI, ZI, cmap=cm.jet, shading='auto')
plt.scatter(x, y, 100, np.asarray(z), cmap=cm.jet, edgecolors = 'black')
plt.xlim(200, 250)
plt.ylim(900, 932)
plt.colorbar()

plt.grid(True)
plt.show()