from datetime import datetime
import json
import logging
from os import path
import pytz

import numpy as np
import rethinkdb as r
import yaml

from cxflow.hooks import AbstractHook


class RethinkDBHook(AbstractHook):
    """
    Stores the training configuration and progress to RethinkDB.
    
    The access is configured by the credentials file, which must be protected. The configuration directly via config
    is not supported as the configs are commonly in git.
    
    -------------------------------------------------------
    The saved document structure:
    -------------------------------------------------------
    {
        id: RethinkDB id
        timestamp: document creation timestamp 
        config: the cxflow config (after CLI application)
        training: [
            {
                timestamp: epoch update timestamp
                epoch_id: id of the epoch
                epoch_data: the `epoch_data` object 
            }, 
            ...
        ]
    }
    -------------------------------------------------------
    
    -------------------------------------------------------
    Example usage in config
    -------------------------------------------------------
     - class: RethinkDBHook
       credentials_file: local_config/rethinkdb.json
       table: example_table
    -------------------------------------------------------
       
    -------------------------------------------------------
    Where `local_config/rethinkdb.json` is a text file containing:
    -------------------------------------------------------
    {
      "host": "database.url.com",
      "port": 28016,
      "user": "JohnWick",
      "password": "SuperPassword123?",
      "db": "example_db"
    }
    -------------------------------------------------------

    """

    def __init__(self, output_dir: str, credentials_file: str, table: str, config_file: str='config.yaml',
                 rethink_key_file: str='rethink_key.json', **kwargs):
        """
        Save training config to database `configs`. Save the unique id generated by the database.
        :param output_dir: output (logging) directory
        :param config: the whole configuration file
        :param credentials_file: path to JSON credentials file which contains fields: host, port, user, password, db.
                                 This file must not be included in git.
        :param table: database table in which the results will be stored
        """

        super().__init__(output_dir=output_dir, **kwargs)

        self._table = table

        with open(credentials_file, 'r') as file:
            self._credentials = json.load(file)

        with open(path.join(output_dir, config_file), 'r') as config_f:
            config = yaml.load(config_f)

        logging.debug('Creating training document in the db')
        with r.connect(**self._credentials) as conn:
            response = r.table(self._table)\
                        .insert({'config': config, 'training': [], 'timestamp': r.expr(datetime.now(pytz.utc))})\
                        .run(conn)
            if response['errors'] > 0:
                logging.error('Error: %s', response['errors'])
                return
            if response['inserted'] != 1:
                logging.error('Inserted unexpected number of documents: %s', response['inserted'])
                return
            self._rethink_id = response['generated_keys'][0]
            logging.debug('Created config: %s', self._rethink_id)

        rethink_id_file = path.join(output_dir, rethink_key_file)
        logging.info('Saving document id to `%s`', rethink_id_file)
        with open(rethink_id_file, 'w') as file:
            json.dump({'rethink_id': self._rethink_id}, file)

    @staticmethod
    def _numpy_dict_to_jsonable_dict(data):
        """Make a dict containing numpy arrays/scalars JSON serializable."""

        result = {}
        for key, value in data.items():
            if isinstance(value, np.ndarray):
                result[key] = value.tolist()
            elif isinstance(value, list):
                result[key] = value
            elif np.isscalar(value):
                if isinstance(value, np.float_):
                    result[key] = float(value)
                elif isinstance(value, np.int_):
                    result[key] = int(value)
            elif isinstance(value, dict):  # works for OrderedDict as well
                result[key] = RethinkDBHook._numpy_dict_to_jsonable_dict(value)
            else:
                raise ValueError('Unsupported JSON type `{}` (key `{}`)'.format(type(value), key))
        return result

    def after_epoch(self, epoch_id: int, epoch_data: AbstractHook.EpochData, **kwargs) -> None:
        logging.info('Rethink: after epoch %d', epoch_id)

        with r.connect(**self._credentials) as conn:
            item = {'timestamp': r.expr(datetime.now(pytz.utc)),
                    'epoch_id': epoch_id,
                    'epoch_data': RethinkDBHook._numpy_dict_to_jsonable_dict(epoch_data)}

            response = r.table(self._table)\
                        .get(self._rethink_id)\
                        .update({'training': r.row['training'].append(item)})\
                        .run(conn)

            if response['errors'] > 0:
                logging.error('Error: %s', response['errors'])
                return
            if response['replaced'] != 1:
                logging.error('Modified unexpected number of documents: %s instead of 1', response['replaced'])
                return
            logging.debug('Appended train. progress to: %s', self._rethink_id)
