import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
import itertools as it

from sklearn.metrics import r2_score

from math import sin

def f(x, y): return (1 + sin(x**2 + 5*y)) / 2

def train_generic(model, marker):
    model.compile(optimizer = "adam", loss = "mse")

    X = np.linspace(0, 10, 300)
    Y = np.linspace(0, 10, 300)

    ins = np.array(list(it.product(X, Y)))
    ous = np.array(list(f(*i) for i in ins))

    result = model.fit(ins, ous, epochs = 200, batch_size = 2048)

    model.save_weights(f"save-{marker}.weights.h5")

def verify_generic(model, marker):
    model.compile(optimizer = "adam", loss = "mse")
    model.load_weights(f"save-{marker}.weights.h5")

    X = np.linspace(0, 10, 50)
    Y = np.linspace(0, 10, 50)

    ins = np.array(list(it.product(X, Y)))
    ous = np.array(list(f(*i) for i in ins))

    preds = model.predict(ins)

    print(f"Model {marker} has {r2_score(ous, preds)} r2 score")

    preds_flat = [i[0] for i in preds]

    px = np.array([X] * len(X))
    py = np.array([[x] * (len(X)) for x in X])
    pz1 = np.array([ous[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])
    pz2 = np.array([preds_flat[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])
    
    #print([ous[i*len(X):(i+1)*len(X)] for i, _ in enumerate(X)])

    p = plt.figure().add_subplot(projection='3d')
    p.plot_surface(px, py, pz1, edgecolor = "lime", alpha = 0.1)
    p.plot_surface(px, py, pz2, edgecolor = "blue", alpha = 0.3)
    plt.show()