import fire
import ray
import pandas as pd
import tensorflow as tf
from ray import tune
from tensorflow import keras
from utils import *
from model import Model

B_params = {
    'batch_size': tune.grid_search([256]),
    'layers': tune.grid_search([[512, 512]]),
    'lr': tune.grid_search([3e-4]),
    'optimizer': tune.grid_search([tf.keras.optimizers.Adam]),
}

df = pd.read_pickle('data.p')

class TuneB(tune.Trainable):
    def _setup(self, config):
        logger.debug('Ray Tune model configuration %s', config)

        self.model = Model('b', epochs=1)

        optimizer = config['optimizer']
        optimizer = config['optimizer'](lr=config['lr'])

        self.model.prepare_for_use(df=df, batch_size=config['batch_size'], layers=config['layers'], optimizer=optimizer)

    def _train(self):
        logs = self.model.train(self.config)

        metrics = {
            'mean_accuracy': logs.history['acc'][0],
            'loss': logs.history['loss'][0],
            'val_accuracy': logs.history['val_acc'][0],
            'val_loss': logs.history['val_loss'][0],
        }

        return metrics

    def _save(self, checkpoint_dir):
        return self.model.save(checkpoint_dir)

    def _restore(self, path):
        return self.model.restore(path)

A_params = {
    'batch_size': tune.grid_search([128]),
    'layers': tune.grid_search([[64, 64]]),
    'lr': tune.grid_search([1e-4]),
    'optimizer': tune.grid_search([tf.keras.optimizers.Adam]),
}

class TuneA(tune.Trainable):
    def _setup(self, config):
        logger.debug('Ray Tune model configuration %s', config)

        self.model = Model('a', epochs=1)

        optimizer = config['optimizer']
        optimizer = config['optimizer'](lr=config['lr'])

        self.model.prepare_for_use(
                df=df,
                batch_size=config['batch_size'],
                layers=config['layers'],
                optimizer=optimizer,
                dataset_fn=dataframe_to_dataset_temp_precip,
                loss='mse'
        )

    def _train(self):
        logs = self.model.train(self.config)

        metrics = {
            'mean_accuracy': logs.history['acc'][0],
            'loss': logs.history['loss'][0],
            'val_accuracy': logs.history['val_acc'][0],
            'val_loss': logs.history['val_loss'][0],
        }

        return metrics

    def _save(self, checkpoint_dir):
        return self.model.save(checkpoint_dir)

    def _restore(self, path):
        return self.model.restore(path)

def start_tuning(model, cpu=1, gpu=2, checkpoint_freq=1, checkpoint_at_end=True, resume=False, restore=None, stop=500):
    ray.init()

    if model == 'a':
        t = TuneA
        params = A_params
    else:
        t = TuneB
        params = B_params

    tune.run(t,
            config=params,
            resources_per_trial={
                "cpu": cpu,
                "gpu": gpu
            },
            resume=resume,
            checkpoint_at_end=checkpoint_at_end,
            checkpoint_freq=checkpoint_freq,
            restore=restore,
            stop={
              'training_iteration': stop
            })

if __name__ == "__main__":
    fire.Fire(start_tuning)