import logging
import os
from datetime import datetime
from inspect import signature, Parameter
from pathlib import Path
from pprint import pprint
from textwrap import dedent
from typing import Optional, Union

import fire
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, TerminateOnNaN
from tensorflow.keras import Model

from spellnn import models
from spellnn.data import alphabet
from spellnn.data.alphabet import get_chars
from spellnn.data.processing import DataProcessor
from spellnn.data.util import nb_lines
from spellnn.layers.mapping import CharMapping

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # FATAL
logging.getLogger('tensorflow').setLevel(logging.FATAL)


class Gym:
    def __init__(self):
        self.train_dataset: Optional[tf.data.Dataset] = None
        self.valid_dataset: Optional[tf.data.Dataset] = None
        self.char2int: Optional[CharMapping] = None
        self.model: Optional[Model] = None
        self.nb_train_samples: int = 0
        self.nb_valid_samples: int = 0
        self.batch_size = 0

    def construct_dataset(self, path: str, locale: str, batch_size: int = 32, validation_split: float = 0.3):
        pprint(locals())
        all_chars = [alphabet.START, alphabet.END] + get_chars(locale)
        char_weights = [0.5 if c.isalpha() and c.islower() else
                        0.2 if c.isalpha() else
                        0.1 if c not in {alphabet.START, alphabet.END} else
                        0 for c in all_chars]
        self.char2int = CharMapping(chars=all_chars, include_unknown=True)
        data_processor = DataProcessor(locale=locale, char2id=self.char2int,
                                       alphabet=all_chars, alphabet_weighs=char_weights)

        print('Calculating number of lines in the file...', end=' ')
        all_samples = nb_lines(path)
        print(all_samples)

        self.batch_size = batch_size
        self.nb_train_samples = int((1 - validation_split) * all_samples)
        self.nb_valid_samples = all_samples - self.nb_train_samples

        dataset = tf.data.TextLineDataset(path)
        self.train_dataset = dataset.take(self.nb_train_samples)
        self.train_dataset = self.train_dataset.shuffle(10 * batch_size, seed=42, reshuffle_each_iteration=True)
        self.train_dataset = self.train_dataset.batch(batch_size, drop_remainder=True)
        self.train_dataset = self.train_dataset.map(
            lambda b: tf.numpy_function(func=data_processor.process_batch, inp=[b], Tout=['int32', 'int32', 'int32']))
        self.train_dataset = self.train_dataset.map(lambda enc_in, dec_in, targ: ((enc_in, dec_in), targ))
        self.train_dataset = self.train_dataset.repeat()

        self.valid_dataset = dataset.skip(self.nb_train_samples)
        self.valid_dataset = self.valid_dataset.shuffle(10 * batch_size, seed=42, reshuffle_each_iteration=True)
        self.valid_dataset = self.valid_dataset.batch(batch_size, drop_remainder=True)
        self.valid_dataset = self.valid_dataset.map(
            lambda b: tf.numpy_function(func=data_processor.process_batch, inp=[b], Tout=['int32', 'int32', 'int32']))
        self.valid_dataset = self.valid_dataset.map(lambda enc_in, dec_in, targ: ((enc_in, dec_in), targ))
        self.valid_dataset = self.valid_dataset.repeat()
        return self

    def create_model(self, name):
        arguments = signature(getattr(models, name).__init__)
        arguments = {k: v.default for k, v in arguments.parameters.items()
                     if v.default is not Parameter.empty and k != 'self'}
        arguments['nb_symbols'] = len(self.char2int)
        arg_str = ', '.join([f'{k}=' + str(v) if type(v) != str else f'{k}=' '"' + str(v) + '"'
                             for k, v in arguments.items()])
        # print(arg_str)
        exec(dedent(f'''
        def create({arg_str}):
            self.model = {name}(**locals())
            return self
        create.__name__ = {name}.__name__
        create.__doc__ = {name}.__init__.__doc__
        setattr(self, create.__name__, create)
        '''), {'self': self, name: getattr(models, name), arg_str: arg_str})
        return getattr(self, name)

    def train(self, epochs: int, monitor_metric='val_acc', patience: int = 5,
              steps_per_epoch: Union[int, str] = 'auto', validation_steps: Union[int, str] = 'auto',
              log_dir: str = 'logs',
              use_multiprocessing: bool = False):
        pprint(locals())
        log_dir = Path(log_dir).joinpath(datetime.now().replace(microsecond=0).isoformat())
        model_path = Path(log_dir).joinpath('checkpoints').joinpath('best-model.h5py')
        model_path = str(model_path)

        if steps_per_epoch == 'auto':
            steps_per_epoch = self.nb_train_samples // self.batch_size
        if validation_steps == 'auto':
            validation_steps = self.nb_valid_samples // self.batch_size

        self.model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['acc'])
        history = self.model.fit_generator(
            self.train_dataset.as_numpy_iterator(), steps_per_epoch=steps_per_epoch,
            validation_data=self.valid_dataset.as_numpy_iterator(), validation_steps=validation_steps,
            epochs=epochs,
            use_multiprocessing=use_multiprocessing, workers=os.cpu_count() - 1,
            callbacks=[
                TerminateOnNaN(),
                TensorBoard(log_dir=log_dir),
                ModelCheckpoint(model_path, monitor=monitor_metric, verbose=1, save_best_only=True),
                EarlyStopping(monitor=monitor_metric, patience=patience),
            ])
        return history.history


if __name__ == '__main__':
    cli = Gym()
    fire.Fire(cli)

[((122, 4, 122, 18), 'fire.Fire', 'fire.Fire', ({(122, 14, 122, 17): 'cli'}, {}), '(cli)', False, 'import fire\n'), ((23, 0, 23, 31), 'logging.getLogger', 'logging.getLogger', ({(23, 18, 23, 30): '"""tensorflow"""'}, {}), "('tensorflow')", False, 'import logging\n'), ((43, 24, 43, 74), 'spellnn.layers.mapping.CharMapping', 'CharMapping', (), '', False, 'from spellnn.layers.mapping import CharMapping\n'), ((44, 25, 45, 88), 'spellnn.data.processing.DataProcessor', 'DataProcessor', (), '', False, 'from spellnn.data.processing import DataProcessor\n'), ((48, 22, 48, 36), 'spellnn.data.util.nb_lines', 'nb_lines', ({(48, 31, 48, 35): 'path'}, {}), '(path)', False, 'from spellnn.data.util import nb_lines\n'), ((55, 18, 55, 47), 'tensorflow.data.TextLineDataset', 'tf.data.TextLineDataset', ({(55, 42, 55, 46): 'path'}, {}), '(path)', True, 'import tensorflow as tf\n'), ((38, 53, 38, 70), 'spellnn.data.alphabet.get_chars', 'get_chars', ({(38, 63, 38, 69): 'locale'}, {}), '(locale)', False, 'from spellnn.data.alphabet import get_chars\n'), ((81, 13, 88, 12), 'textwrap.dedent', 'dedent', ({(81, 20, 88, 11): 'f"""\n        def create({arg_str}):\n            self.model = {name}(**locals())\n            return self\n        create.__name__ = {name}.__name__\n        create.__doc__ = {name}.__init__.__doc__\n        setattr(self, create.__name__, create)\n        """'}, {}), '(\n    f"""\n        def create({arg_str}):\n            self.model = {name}(**locals())\n            return self\n        create.__name__ = {name}.__name__\n        create.__doc__ = {name}.__init__.__doc__\n        setattr(self, create.__name__, create)\n        """\n    )', False, 'from textwrap import dedent\n'), ((60, 22, 60, 117), 'tensorflow.numpy_function', 'tf.numpy_function', (), '', True, 'import tensorflow as tf\n'), ((68, 22, 68, 117), 'tensorflow.numpy_function', 'tf.numpy_function', (), '', True, 'import tensorflow as tf\n'), ((96, 18, 96, 31), 'pathlib.Path', 'Path', ({(96, 23, 96, 30): 'log_dir'}, {}), '(log_dir)', False, 'from pathlib import Path\n'), ((110, 61, 110, 75), 'os.cpu_count', 'os.cpu_count', ({}, {}), '()', False, 'import os\n'), ((112, 16, 112, 32), 'tensorflow.keras.callbacks.TerminateOnNaN', 'TerminateOnNaN', ({}, {}), '()', False, 'from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, TerminateOnNaN\n'), ((113, 16, 113, 44), 'tensorflow.keras.callbacks.TensorBoard', 'TensorBoard', (), '', False, 'from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, TerminateOnNaN\n'), ((114, 16, 114, 99), 'tensorflow.keras.callbacks.ModelCheckpoint', 'ModelCheckpoint', (), '', False, 'from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, TerminateOnNaN\n'), ((115, 16, 115, 72), 'tensorflow.keras.callbacks.EarlyStopping', 'EarlyStopping', (), '', False, 'from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, TerminateOnNaN\n'), ((97, 21, 97, 34), 'pathlib.Path', 'Path', ({(97, 26, 97, 33): 'log_dir'}, {}), '(log_dir)', False, 'from pathlib import Path\n'), ((96, 41, 96, 55), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n')]