from pandas import read_csv
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras import Sequential
import tensorflow as tf
from tensorflow.keras.layers import Dense
import datetime
import os
import keras
import json
import glob


print("Tensforflow version: ", tf.__version__)
print("Availabe devices: ", devices)
if len(devices) == 0:
    print("No devices for mac found")
    exit(1)

    

directory = os.environ['TF_FOLDER']

checkpoint_dir = os.path.join(directory, "ckpt")
if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir)

path = 'https://raw.githubusercontent.com/jbrownlee/Datasets/master/ionosphere.csv'
df = read_csv(path, header=None)

X, y = df.values[:, :-1], df.values[:, -1]

X = X.astype('float32')

y = LabelEncoder().fit_transform(y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33)
print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)

n_features = X_train.shape[1]


def get_compiled_model():
       model = tf.keras.Sequential([
            tf.keras.layers.Flatten(input_shape=(n_features,)),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10)
        ])
       model.summary()
       model.compile(optimizer=tf.keras.optimizers.legacy.Adam(learning_rate=0.0001), loss='binary_crossentropy', metrics=['accuracy'])
       return model

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

with strategy.scope():
    model = get_compiled_model()

log_dir = os.path.join(directory, "logs/fit/") + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)


class CustomModelCheckpoint(keras.callbacks.ModelCheckpoint):
    def __init__(self, filepath, max_to_keep=100, **kwargs):
        super().__init__(filepath, **kwargs)
        self.filepath = filepath
        self.max_to_keep = max_to_keep

    def on_epoch_end(self, epoch, logs=None):
        super().on_epoch_end(epoch, logs)
        files = sorted(glob.glob(self.filepath.format(epoch='*')))
        if len(files) > self.max_to_keep:
            for f in files[:-self.max_to_keep]:
                os.remove(f)


callbacks = [
        CustomModelCheckpoint(
            filepath=checkpoint_dir + "/ckpt-{epoch}", save_freq="epoch", max_to_keep=10, save_weights_only=True
        ),
       keras.callbacks.TensorBoard('tensorboard_logs')
    ]

latest = tf.train.latest_checkpoint(checkpoint_dir)
if latest:
    print("Loading model checkpoint {} ...\n".format(latest))
    model.load_weights(latest)

model.fit(X_train, y_train, epochs=100, batch_size=32, verbose=1,callbacks=callbacks)

loss, acc = model.evaluate(X_test, y_test, verbose=0)
print(f'Test Accuracy: {acc:.3f}')

2024-01-28 14:46:40.395499: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:117] Plugin optimizer for device_type GPU is enabled.
Epoch 1/100
2024-01-28 14:46:40.778281: W tensorflow/core/framework/op_kernel.cc:1803] INTERNAL: Failed to build OpKernel for Add : No registered 'Add' OpKernel for 'GPU' devices compatible with node {{node Add}}
	 (OpKernel was found, but attributes didn't match) Requested Attributes: T=DT_INT64, _device="/job:worker/replica:0/task:0/device:GPU:0"
	.  Registered:  device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, DT_INT8, DT_COMPLEX64, DT_INT64, DT_BFLOAT16, DT_COMPLEX128, DT_HALF]
  device='GPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_HALF]
  device='GPU'; T in [DT_BFLOAT16]
  device='DEFAULT'; T in [DT_INT32]
  device='CPU'; T in [DT_FLOAT]
  device='CPU'; T in [DT_HALF]
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_INT32]
  device='CPU'; T in [DT_INT64]
  device='CPU'; T in [DT_BFLOAT16]
  device='CPU'; T in [DT_INT8]
  device='CPU'; T in [DT_INT16]
  device='CPU'; T in [DT_COMPLEX64]
  device='CPU'; T in [DT_UINT8]
  device='CPU'; T in [DT_COMPLEX128]
  device='CPU'; T in [DT_STRING]

Traceback (most recent call last):
  File "/Users/inet11/git/tensorflow/model.py", line 129, in <module>
    model.fit(X_train, y_train, epochs=100, batch_size=32, verbose=1,callbacks=callbacks)
  File "/Users/inet11/git/tensorflow/email/tf/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/inet11/git/tensorflow/email/tf/lib/python3.11/site-packages/tensorflow/python/eager/execute.py", line 53, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.InternalError: Graph execution error:

Detected at node Add defined at (most recent call last):
<stack traces unavailable>
2 root error(s) found.
  (0) INTERNAL:  Failed to build OpKernel for Add : No registered 'Add' OpKernel for 'GPU' devices compatible with node {{node Add}}
	 (OpKernel was found, but attributes didn't match) Requested Attributes: T=DT_INT64, _device="/job:worker/replica:0/task:0/device:GPU:0"
	.  Registered:  device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, DT_INT8, DT_COMPLEX64, DT_INT64, DT_BFLOAT16, DT_COMPLEX128, DT_HALF]
  device='GPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_HALF]
  device='GPU'; T in [DT_BFLOAT16]
  device='DEFAULT'; T in [DT_INT32]
  device='CPU'; T in [DT_FLOAT]
  device='CPU'; T in [DT_HALF]
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_INT32]
  device='CPU'; T in [DT_INT64]
  device='CPU'; T in [DT_BFLOAT16]
  device='CPU'; T in [DT_INT8]
  device='CPU'; T in [DT_INT16]
  device='CPU'; T in [DT_COMPLEX64]
  device='CPU'; T in [DT_UINT8]
  device='CPU'; T in [DT_COMPLEX128]
  device='CPU'; T in [DT_STRING]

	 [[CollectiveReduceV2]]
  (1) CANCELLED:  Function was cancelled before it was started
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_1260]