from django.db import models
from . import cgan as cg
from time import time

import tensorflow as tf
import numpy as np
import sys

sys.path.append('../')
from utils import image_utils as iu


SEED = 1337
is_model_loaded = False

dirs = {
    'output': './cgan/gen_img/',
    'model_path': './cgan/cgan_model/',
}

np.random.seed(SEED)

# Global TF obj
tf_sess = None
cgan_model = None


class CGANInference:

    @staticmethod
    def load_model(verbose=True):
        global is_model_loaded
        global tf_sess, cgan_model

        if not is_model_loaded:
            # GPU configure
            config = tf.ConfigProto()
            config.gpu_options.allow_growth = True

            tf_sess = tf.Session(config=config)
            cgan_model = cg.CGAN(tf_sess)  # CGAN Model
            if verbose:
                print("[+] CGAN Model loaded!")

            # initializing
            tf_sess.run(tf.global_variables_initializer())

            # Load model & Graph & Weights
            ckpt = tf.train.get_checkpoint_state(dirs['model_path'])
            if ckpt and ckpt.model_checkpoint_path:
                # Restores from checkpoint
                cgan_model.saver.restore(tf_sess, ckpt.model_checkpoint_path)

                saved_global_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])
                if verbose:
                    print("[+] global step : %d" % saved_global_step, " successfully loaded")

                is_model_loaded = True
            else:
                if verbose:
                    print('[-] No checkpoint file found :( pre-trained model is needed!')

                is_model_loaded = False

    @staticmethod
    def gen_image(input_label=1, verbose=True):
        global is_model_loaded
        global tf_sess, cgan_model

        num_label = np.eye(10)[input_label]  # one-hot encoding
        num_label = num_label.reshape(-1, 10)

        # Training G model with sample image and noise
        sample_z = np.random.uniform(-1., 1., [1, cgan_model.z_dim]).astype(np.float32)
        samples = tf_sess.run(cgan_model.g,
                              feed_dict={
                                  cgan_model.c: num_label,
                                  cgan_model.z: sample_z,
                                  cgan_model.do_rate: 0.0,
                              })

        samples = np.reshape(samples, [-1, 28, 28, 1])
        if verbose:
            print("[*] samples shape : ", samples.shape)

        # Export image generated by model G
        sample_image_height = 1
        sample_image_width = 1
        sample_dir = dirs['output'] + "gen.png"  # .format(input_label)

        # Generated image save
        iu.save_images(samples,
                       size=[sample_image_height, sample_image_width],
                       image_path=sample_dir)
        if verbose:
            print("[*] generated image (%d) is saved" % input_label)

    def __init__(self, input_label=0, verbose=True):
        self.success = False
        self.verbose = verbose

        self.load_model(verbose=self.verbose)  # Loading CGAN Model
        if is_model_loaded:
            self.start_time = time()
            self.gen_image(input_label, verbose=self.verbose)
            self.end_time = time()

            if self.verbose:
                print("[*] Total {:.4f}s is taken".format(self.end_time - self.start_time))

            self.success = True
        else:
            self.success = False

    def __str__(self):
        return "CGAN Model"
