#!/usr/bin/env python3

import libspn as spn
import tensorflow as tf
from enum import Enum


class Dataset(Enum):
    IMAGE = 0
    MNIST = 1


class Model(Enum):
    DISCRETE_DENSE = 0


class Learner(Enum):
    EM_SUP = 0
    EM_UNSUP = 1


class WeightInit(Enum):
    ONE = 1
    RANDOM = 2


class SpnModel(spn.App):

    def __init__(self):
        super().__init__("Model")

    def _define_build_args(self, parser):
        parser.add_argument('model', type=str,
                            choices=[i.name.lower() for i in Model],
                            help="model to build")

        model_params = parser.add_argument_group(title="model parameters")
        model_params.add_argument('--num-classes', type=int, default=1,
                                  help="number of classes assumed by the model; "
                                  "set to 1 to use a 1-class model")
        model_params.add_argument('--weight-init', type=str, default='random',
                                  choices=[i.name.lower() for i in WeightInit],
                                  help='type of weight initialization')
        model_params.add_argument('--dense-decomps', type=int, default=1,
                                  help='dense generator: number of decompositions')
        model_params.add_argument('--dense-subsets', type=int, default=2,
                                  help='dense generator: number of subsets')
        model_params.add_argument('--dense-mixtures', type=int, default=4,
                                  help='dense generator: number of mixtures')
        model_params.add_argument('--dense-input-dist', type=str, default='raw',
                                  choices=[i.name.lower() for i in spn.DenseSPNGenerator.InputDist],
                                  help='dense generator: type of input distributions')
        model_params.add_argument('--dense-input-mixtures', type=int, default=3,
                                  help='dense generator: number of input mixtures')
        model_params.add_argument('--dense-seed', type=int, default=None,
                                  help='seed used to generate the random structure')

    def _define_load_args(self, parser):
        parser.add_argument('path', type=str,
                            help="path to the model file")

    def _define_train_args(self, parser):
        parser.add_argument('learner', type=str,
                            choices=[i.name.lower() for i in Learner],
                            help="learner to use")
        parser.add_argument('dataset', type=str,
                            choices=[i.name.lower() for i in Dataset],
                            help="training dataset to use")
        parser.add_argument('path', type=str, nargs='?',
                            help="path to the training data; "
                            "a folder, a glob, or a coma-separated list of files "
                            "see specific dataset docs for details; no path is "
                            "needed for MNIST")

        data_params = parser.add_argument_group(title="data parameters")
        data_params.add_argument('--classes', type=str, default=None,
                                 help="coma-separated list of labels of classes to use; "
                                 "if not provided, all classes will be used for some datasets")
        data_params.add_argument('--num-epochs', type=int, default=100,
                                 help="number of training epochs")
        data_params.add_argument('--batch-size', type=int, default=100,
                                 help="size of a batch used for training")
        data_params.add_argument('--allow_smaller_batch', type=spn.App.bool_arg,
                                 default=False, help="allow smaller final batch")
        data_params.add_argument('--num-threads', type=int, default=1,
                                 help="number of threads enqueuing the data queue.")
        data_params.add_argument('--shuffle', type=spn.App.bool_arg, default=False,
                                 help="shuffle the data every epoch")
        data_params.add_argument('--min-after-dequeue', type=int, default=None,
                                 help="min number of elements in the data queue "
                                 "after each dequeue used for shuffling; not used "
                                 "for image or MNIST datasets")
        data_params.add_argument('--seed', type=int, default=None,
                                 help="seed used for shuffling")

        image_params = parser.add_argument_group(title="image data parameters")
        image_params.add_argument('--image-format', type=str, default='float',
                                  choices=[i.name.lower() for i in spn.ImageFormat],
                                  help="image format")
        image_params.add_argument('--ratio', type=int, default=1,
                                  help="downsample images by the given ratio")
        image_params.add_argument('--crop', type=int, default=0,
                                  help="crop image border pixels")
        image_params.add_argument('--accurate', type=spn.App.bool_arg, default=True,
                                  help="more accurate JPEG decompression")

        mnist_cifar_params = parser.add_argument_group(
            title="MNIST/CIFAR data parameters")
        mnist_cifar_params.add_argument('--mnist-subset', type=str, default='all',
                                        choices=[i.name.lower()
                                                 for i in spn.MNISTDataset.Subset],
                                        help="subset to use")
        mnist_cifar_params.add_argument('--cifar-subset', type=str, default='all',
                                        choices=[i.name.lower()
                                                 for i in spn.CIFAR10Dataset.Subset],
                                        help="subset to use")

        learn_params = parser.add_argument_group(title="learning parameters")
        learn_params.add_argument('--value-inference', type=str, default='marginal',
                                  choices=[i.name.lower() for i in spn.InferenceType],
                                  help='type of inference during EM upwards pass')
        learn_params.add_argument('--init-accum', type=int, default=20,
                                  help='initial accumulator value')
        learn_params.add_argument('--smoothing-val', type=float, default=0.0,
                                  help='additive smoothing value')
        learn_params.add_argument('--smoothing-decay', type=float, default=0.2,
                                  help='additive smoothing decay')
        learn_params.add_argument('--smoothing-min', type=float, default=0.0,
                                  help='additive smoothing min value')
        learn_params.add_argument('--stop-condition', type=float, default=0.05,
                                  help='min likelihood change between epochs')

    def _define_test_args(self, parser):
        parser.add_argument('dataset', type=str,
                            choices=[i.name.lower() for i in Dataset],
                            help="test dataset to use")
        parser.add_argument('path', type=str, nargs='?',
                            help="path to the test data; "
                            "a folder, a glob, or a coma-separated list of files "
                            "see specific dataset docs for details; no path is "
                            "needed for MNIST")

        data_params = parser.add_argument_group(title="data parameters")
        data_params.add_argument('--classes', type=str, default=None,
                                 help="coma-separated list of labels of classes to use; "
                                 "if not provided, all classes will be used for some datasets")
        data_params.add_argument('--batch-size', type=int, default=100,
                                 help="size of a batch used for training")
        data_params.add_argument('--num-threads', type=int, default=1,
                                 help="number of threads enqueuing the data queue.")

        image_params = parser.add_argument_group(title="image data parameters")
        image_params.add_argument('--image-format', type=str, default='float',
                                  choices=[i.name.lower() for i in spn.ImageFormat],
                                  help="image format")
        image_params.add_argument('--ratio', type=int, default=1,
                                  help="downsample images by the given ratio")
        image_params.add_argument('--crop', type=int, default=0,
                                  help="crop image border pixels")
        image_params.add_argument('--accurate', type=spn.App.bool_arg, default=True,
                                  help="more accurate JPEG decompression")

        mnist_cifar_params = parser.add_argument_group(
            title="MNIST/CIFAR data parameters")
        mnist_cifar_params.add_argument('--mnist-subset', type=str, default='all',
                                        choices=[i.name.lower()
                                                 for i in spn.MNISTDataset.Subset],
                                        help="subset to use")
        mnist_cifar_params.add_argument('--cifar-subset', type=str, default='all',
                                        choices=[i.name.lower()
                                                 for i in spn.CIFAR10Dataset.Subset],
                                        help="subset to use")

    def _define_mpe_args(self, parser):
        pass

    def _define_save_args(self, parser):
        parser.add_argument('path', type=str,
                            help="path to the model file")
        parser.add_argument('--pretty', type=spn.App.bool_arg,
                            default=False, help="use pretty printing")

    def define_args(self, parser, commands):
        self._define_load_args(
            commands.add_parser('load', app=self,
                                help='load a model'))
        self._define_build_args(
            commands.add_parser('build', app=self,
                                help='build model structure'))
        self._define_train_args(
            commands.add_parser('train', app=self,
                                help='train a model'))
        self._define_test_args(
            commands.add_parser('test', app=self,
                                help='get p(labels|data)'))
        self._define_mpe_args(
            commands.add_parser('mpe', app=self,
                                help='get mpe state of data variables'))
        self._define_save_args(
            commands.add_parser('save', app=self,
                                help='save a model'))

    def process_args(self, parser, args):
        if (args.build is None) == (args.load is None):
            parser.error("build xor load command must be used")

        if (args.build and not args.train):
            parser.error("a new model must be trained")

        ##########################
        # Build command
        ##########################
        if args.build:
            # Positional
            args.build.model = Model[args.build.model.upper()]
            # Model
            args.build.dense_input_dist = spn.DenseSPNGenerator.InputDist[
                args.build.dense_input_dist.upper()]
            args.build.weight_init = WeightInit[args.build.weight_init.upper()]

        ##########################
        # Train command
        ##########################
        if args.train:
            # Positional
            args.train.dataset = Dataset[args.train.dataset.upper()]
            if args.train.dataset is not Dataset.MNIST and args.train.path is None:
                parser.error("PATH is needed for this dataset")
            # Data
            if args.train.classes is not None:
                args.train.classes = [i.strip() for i in
                                      args.train.classes.split(',')]
            args.train.image_format = spn.ImageFormat[args.train.image_format.upper()]
            if args.train.ratio < 1:
                parser.error("RATIO must be >=1")
            if args.train.crop < 0:
                parser.error("CROP cannot be negative")
            args.train.mnist_subset = spn.MNISTDataset.Subset[args.train.mnist_subset.upper()]
            args.train.cifar_subset = spn.CIFAR10Dataset.Subset[args.train.cifar_subset.upper()]
            # Learning
            args.train.value_inference = spn.InferenceType[args.train.value_inference.upper()]

        ##########################
        # Test command
        ##########################
        if args.test:
            # Positional
            args.test.dataset = Dataset[args.test.dataset.upper()]
            if args.test.dataset is not Dataset.MNIST and args.test.path is None:
                parser.error("PATH is needed for this dataset")
            # Data
            if args.test.classes is not None:
                args.test.classes = [i.strip() for i in
                                     args.test.classes.split(',')]
            args.test.image_format = spn.ImageFormat[args.test.image_format.upper()]
            if args.test.ratio < 1:
                parser.error("RATIO must be >=1")
            if args.test.crop < 0:
                parser.error("CROP cannot be negative")
            args.test.mnist_subset = spn.MNISTDataset.Subset[args.test.mnist_subset.upper()]
            args.test.cifar_subset = spn.CIFAR10Dataset.Subset[args.test.cifar_subset.upper()]

    def _prepare_data(self, args):
        if args.train:
            if args.train.dataset == Dataset.MNIST:
                self.train_set = spn.MNISTDataset(
                    subset=args.train.mnist_subset,
                    format=args.train.image_format,
                    num_epochs=args.train.num_epochs,
                    batch_size=args.train.batch_size,
                    shuffle=args.train.shuffle,
                    ratio=args.train.ratio,
                    crop=args.train.crop,
                    num_threads=args.train.num_threads,
                    allow_smaller_final_batch=args.train.allow_smaller_batch,
                    classes=args.train.classes,
                    seed=args.train.seed)
            elif args.train.dataset == Dataset.IMAGE:
                self.train_set = spn.ImageDataset(
                    image_files=args.train.path,
                    format=args.train.image_format,
                    num_epochs=args.train.num_epochs,
                    batch_size=args.train.batch_size,
                    shuffle=args.train.shuffle,
                    ratio=args.train.ratio,
                    crop=args.train.crop,
                    accurate=args.train.accurate,
                    num_threads=args.train.num_threads,
                    allow_smaller_final_batch=args.train.allow_smaller_batch,
                    classes=args.train.classes,
                    seed=args.train.seed)
            elif args.train.dataset == Dataset.CIFAR10:
                self.train_set = spn.CIFAR10Dataset(
                    data_dir=args.train.path,
                    subset=args.train.cifar_subset,
                    format=args.train.image_format,
                    num_epochs=args.train.num_epochs,
                    batch_size=args.train.batch_size,
                    shuffle=args.train.shuffle,
                    ratio=args.train.ratio,
                    crop=args.train.crop,
                    min_after_dequeue=args.train.min_after_dequeue,
                    num_threads=args.train.num_threads,
                    allow_smaller_final_batch=args.train.allow_smaller_batch,
                    classes=args.train.classes,
                    seed=args.train.seed)

        if args.test:
            if args.test.dataset == Dataset.MNIST:
                self.test_set = spn.MNISTDataset(
                    subset=args.test.mnist_subset,
                    format=args.test.image_format,
                    num_epochs=1,
                    batch_size=args.test.batch_size,
                    shuffle=False,
                    ratio=args.test.ratio,
                    crop=args.test.crop,
                    num_threads=args.test.num_threads,
                    allow_smaller_final_batch=True,
                    classes=args.test.classes)
            elif args.test.dataset == Dataset.IMAGE:
                self.test_set = spn.ImageDataset(
                    image_files=args.test.path,
                    format=args.test.image_format,
                    num_epochs=1,
                    batch_size=args.test.batch_size,
                    shuffle=False,
                    ratio=args.test.ratio,
                    crop=args.test.crop,
                    accurate=args.test.accurate,
                    num_threads=args.test.num_threads,
                    allow_smaller_final_batch=True,
                    classes=args.test.classes)
            elif args.test.dataset == Dataset.CIFAR10:
                self.test_set = spn.CIFAR10Dataset(
                    data_dir=args.test.path,
                    subset=args.test.cifar_subset,
                    format=args.test.image_format,
                    num_epochs=1,
                    batch_size=args.test.batch_size,
                    shuffle=False,
                    ratio=args.test.ratio,
                    crop=args.test.crop,
                    num_threads=args.test.num_threads,
                    allow_smaller_final_batch=True,
                    classes=args.test.classes)

    def _load_model(self, args, sess):
        self.model = spn.Model.load_from_json(args.load.path,
                                              load_param_vals=True,
                                              sess=sess)

    def _build_model(self, args):
        # Decode weight_init_value
        if args.build.weight_init == WeightInit.ONE:
            weight_init_value = 1
        else:
            weight_init_value = tf.initializers.random_uniform(0, 1)

        # Build model
        if args.build.model == Model.DISCRETE_DENSE:
            self.model = spn.DiscreteDenseModel(
                num_classes=args.build.num_classes,
                num_decomps=args.build.dense_decomps,
                num_subsets=args.build.dense_subsets,
                num_mixtures=args.build.dense_mixtures,
                input_dist=args.build.dense_input_dist,
                num_input_mixtures=args.build.dense_input_mixtures,
                weight_initializer=weight_init_value)
            self.model.build(num_vars=self.train_set.num_vars,
                             num_vals=self.train_set.num_vals,
                             seed=args.build.dense_seed)

    def _train_model(self, args, sess):
        init = spn.initialize_weights(self.model.root)
        sess.run(init)

        pass
        # num_vars=int(self.train_samples.shape[1]),
        # num_vals=self.train_num_vals,

        # TODO: Distinguish between em_sup and em_unsup
        # TODO: train, test etc., feed appropriate data through feed_dict
        # model.learn(value_inference_type=args.learn.value_inference,
        #             init_accum_value=args.learn.init_accum,
        #             additive_smoothing_value=args.learn.smoothing_val,
        #             additive_smoothing_decay=args.learn.smoothing_decay,
        #             additive_smoothing_min=args.learn.smoothing_min,
        #             stop_condition=args.learn.stop_condition)

        # Save TF graph
        # self.print2("Saving TensorFlow graph")
        # writer = tf.summary.FileWriter("./logs", self._root.tf_graph)
        # writer.flush()

    def _save_model(self, args, sess):
        self.model.save_to_json(args.save.path, pretty=args.save.pretty,
                                # Save params only if the model was trained
                                save_param_vals=args.train is not None)

    def _test_model(self, args, sess):
        pass

    def run(self, args):
        # Prepare data first, we use dataset info in _build_model
        self._prepare_data(args)

        # Run tasks in session
        with spn.session() as (sess, run):
            # Get model
            if args.build:
                self._build_model(args)
            elif args.load:
                self._load_model(args, sess)
            # Train
            if args.train:
                self._train_model(args, sess)
            # Save
            if args.save:
                self._save_model(args, sess)
            if args.test:
                self._test_model(args, sess)

        # All done
        self.print2("Done!")


if __name__ == '__main__':
    app = SpnModel()
    app.main()
