from __future__ import absolute_import
from __future__ import unicode_literals

import time
import socket
import logging

from ._compat import bytes_types, string_types
from ._compat import struct_l
from .version import __version__

try:
    import ssl
except ImportError:
    ssl = None  # pyflakes.ignore

try:
    from .snappy_socket import SnappySocket
except ImportError:
    SnappySocket = None  # pyflakes.ignore

try:
    import simplejson as json
except ImportError:
    import json  # pyflakes.ignore

import tornado.iostream
import tornado.ioloop

try:
    from tornado.simple_httpclient import _default_ca_certs as default_ca_certs
except ImportError:
    # Tornado < 4
    from tornado.simple_httpclient import _DEFAULT_CA_CERTS

    def default_ca_certs():
        return _DEFAULT_CA_CERTS

from nsq import event, protocol
from .deflate_socket import DeflateSocket

logger = logging.getLogger(__name__)

# states
INIT = 'INIT'
DISCONNECTED = 'DISCONNECTED'
CONNECTING = 'CONNECTING'
CONNECTED = 'CONNECTED'

DEFAULT_USER_AGENT = 'pynsq/%s' % __version__

class AsyncConn(event.EventedMixin):
    """
    Low level object representing a TCP connection to nsqd.

    When a message on this connection is requeued and the requeue delay
    has not been specified, it calculates the delay automatically by an
    increasing multiple of ``requeue_delay``.

    Generates the following events that can be listened to with
    :meth:`nsq.AsyncConn.on`:

     * ``connect``
     * ``close``
     * ``error``
     * ``identify``
     * ``identify_response``
     * ``auth``
     * ``auth_response``
     * ``heartbeat``
     * ``ready``
     * ``message``
     * ``response``
     * ``backoff``
     * ``resume``

    :param host: the host to connect to

    :param port: the post to connect to

    :param timeout: the timeout for read/write operations (in seconds)

    :param heartbeat_interval: the amount of time (in seconds) to negotiate
        with the connected producers to send heartbeats (requires nsqd 0.2.19+)

    :param requeue_delay: the base multiple used when calculating requeue delay
        (multiplied by # of attempts)

    :param tls_v1: enable TLS v1 encryption (requires nsqd 0.2.22+)

    :param tls_options: dictionary of options to pass to `ssl.wrap_socket()
        <http://docs.python.org/2/library/ssl.html#ssl.wrap_socket>`_ as
        ``**kwargs``

    :param snappy: enable Snappy stream compression (requires nsqd 0.2.23+)

    :param deflate: enable deflate stream compression (requires nsqd 0.2.23+)

    :param deflate_level: configure the deflate compression level for this
        connection (requires nsqd 0.2.23+)

    :param output_buffer_size: size of the buffer (in bytes) used by nsqd
        for buffering writes to this connection

    :param output_buffer_timeout: timeout (in ms) used by nsqd before
        flushing buffered writes (set to 0 to disable).  **Warning**:
        configuring clients with an extremely low (``< 25ms``)
        ``output_buffer_timeout`` has a significant effect on ``nsqd``
        CPU usage (particularly with ``> 50`` clients connected).

    :param sample_rate: take only a sample of the messages being sent
        to the client. Not setting this or setting it to 0 will ensure
        you get all the messages destined for the client.
        Sample rate can be greater than 0 or less than 100 and the client
        will receive that percentage of the message traffic.
        (requires nsqd 0.2.25+)

    :param user_agent: a string identifying the agent for this client
        in the spirit of HTTP (default: ``<client_library_name>/<version>``)
        (requires nsqd 0.2.25+)

    :param auth_secret: a byte string passed when using nsq auth
        (requires nsqd 1.0+)

    :param msg_timeout: the amount of time (in seconds) that nsqd will wait
        before considering messages that have been delivered to this
        consumer timed out (requires nsqd 0.2.28+)

    :param hostname: a string identifying the host where this client runs
        (default: ``<hostname>``)
    """
    def __init__(
            self,
            host,
            port,
            timeout=1.0,
            heartbeat_interval=30,
            requeue_delay=90,
            tls_v1=False,
            tls_options=None,
            snappy=False,
            deflate=False,
            deflate_level=6,
            user_agent=DEFAULT_USER_AGENT,
            output_buffer_size=16 * 1024,
            output_buffer_timeout=250,
            sample_rate=0,
            io_loop=None,
            auth_secret=None,
            msg_timeout=None,
            hostname=None):
        assert isinstance(host, string_types)
        assert isinstance(port, int)
        assert isinstance(timeout, float)
        assert isinstance(tls_options, (dict, None.__class__))
        assert isinstance(deflate_level, int)
        assert isinstance(heartbeat_interval, int) and heartbeat_interval >= 1
        assert isinstance(requeue_delay, int) and requeue_delay >= 0
        assert isinstance(output_buffer_size, int) and output_buffer_size >= 0
        assert isinstance(output_buffer_timeout, int) and output_buffer_timeout >= 0
        assert isinstance(sample_rate, int) and sample_rate >= 0 and sample_rate < 100
        assert isinstance(auth_secret, bytes_types + (None.__class__,))
        assert tls_v1 and ssl or not tls_v1, \
            'tls_v1 requires Python 2.6+ or Python 2.5 w/ pip install ssl'
        assert msg_timeout is None or (isinstance(msg_timeout, (float, int)) and msg_timeout > 0)

        self.state = INIT
        self.host = host
        self.port = port
        self.timeout = timeout
        self.last_recv_timestamp = time.time()
        self.last_msg_timestamp = time.time()
        self.in_flight = 0
        self.rdy = 0
        self.rdy_timeout = None
        # for backwards compatibility when interacting with older nsqd
        # (pre 0.2.20), default this to their hard-coded max
        self.max_rdy_count = 2500
        self.tls_v1 = tls_v1
        self.tls_options = tls_options
        self.snappy = snappy
        self.deflate = deflate
        self.deflate_level = deflate_level
        self.hostname = hostname
        if self.hostname is None:
            self.hostname = socket.gethostname()
        self.short_hostname = self.hostname.split('.')[0]
        self.heartbeat_interval = heartbeat_interval * 1000
        self.msg_timeout = int(msg_timeout * 1000) if msg_timeout else None
        self.requeue_delay = requeue_delay
        self.io_loop = io_loop
        if not self.io_loop:
            self.io_loop = tornado.ioloop.IOLoop.instance()

        self.output_buffer_size = output_buffer_size
        self.output_buffer_timeout = output_buffer_timeout
        self.sample_rate = sample_rate
        self.user_agent = user_agent

        self._authentication_required = False  # tracking server auth state
        self.auth_secret = auth_secret

        self.socket = None
        self.stream = None
        self._features_to_enable = []

        self.last_rdy = 0
        self.rdy = 0

        self.callback_queue = []

        super(AsyncConn, self).__init__()

    @property
    def id(self):
        return str(self)

    def __str__(self):
        return self.host + ':' + str(self.port)

    def connected(self):
        return self.state == CONNECTED

    def connecting(self):
        return self.state == CONNECTING

    def closed(self):
        return self.state in (INIT, DISCONNECTED)

    def connect(self):
        if not self.closed():
            return

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.settimeout(self.timeout)
        self.socket.setblocking(0)

        self.stream = tornado.iostream.IOStream(self.socket, io_loop=self.io_loop)
        self.stream.set_close_callback(self._socket_close)
        self.stream.set_nodelay(True)

        self.state = CONNECTING
        self.on(event.CONNECT, self._on_connect)
        self.on(event.DATA, self._on_data)

        self.stream.connect((self.host, self.port), self._connect_callback)

    def _connect_callback(self):
        self.state = CONNECTED
        self.stream.write(protocol.MAGIC_V2)
        self._start_read()
        self.trigger(event.CONNECT, conn=self)

    def _read_bytes(self, size, callback):
        try:
            self.stream.read_bytes(size, callback)
        except IOError:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.ConnectionClosedError('Stream is closed'),
            )

    def _start_read(self):
        self._read_bytes(4, self._read_size)

    def _socket_close(self):
        self.state = DISCONNECTED
        self.trigger(event.CLOSE, conn=self)

    def close(self):
        self.stream.close()

    def _read_size(self, data):
        try:
            size = struct_l.unpack(data)[0]
        except Exception:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.IntegrityError('failed to unpack size'),
            )
            return
        self._read_bytes(size, self._read_body)

    def _read_body(self, data):
        try:
            self.trigger(event.DATA, conn=self, data=data)
        except Exception:
            logger.exception('uncaught exception in data event')
        self._start_read()

    def send(self, data):
        self.stream.write(data)

    def upgrade_to_tls(self, options=None):
        assert ssl, 'tls_v1 requires Python 2.6+ or Python 2.5 w/ pip install ssl'

        # in order to upgrade to TLS we need to *replace* the IOStream...
        #
        # first remove the event handler for the currently open socket
        # so that when we add the socket to the new SSLIOStream below,
        # it can re-add the appropriate event handlers.
        self.io_loop.remove_handler(self.socket.fileno())

        opts = {
            'cert_reqs': ssl.CERT_REQUIRED,
            'ca_certs': default_ca_certs()
        }
        opts.update(options or {})
        self.socket = ssl.wrap_socket(self.socket, ssl_version=ssl.PROTOCOL_TLSv1,
                                      do_handshake_on_connect=False, **opts)

        self.stream = tornado.iostream.SSLIOStream(self.socket, io_loop=self.io_loop)
        self.stream.set_close_callback(self._socket_close)

        # now that the IOStream has been swapped we can kickstart
        # the SSL handshake
        self.stream._do_ssl_handshake()

    def upgrade_to_snappy(self):
        assert SnappySocket, 'snappy requires the python-snappy package'

        # in order to upgrade to Snappy we need to use whatever IOStream
        # is currently in place (normal or SSL)...
        #
        # first read any compressed bytes the existing IOStream might have
        # already buffered and use that to bootstrap the SnappySocket, then
        # monkey patch the existing IOStream by replacing its socket
        # with a wrapper that will automagically handle compression.
        existing_data = self.stream._consume(self.stream._read_buffer_size)
        self.socket = SnappySocket(self.socket)
        self.socket.bootstrap(existing_data)
        self.stream.socket = self.socket

    def upgrade_to_deflate(self):
        # in order to upgrade to DEFLATE we need to use whatever IOStream
        # is currently in place (normal or SSL)...
        #
        # first read any compressed bytes the existing IOStream might have
        # already buffered and use that to bootstrap the DefalteSocket, then
        # monkey patch the existing IOStream by replacing its socket
        # with a wrapper that will automagically handle compression.
        existing_data = self.stream._consume(self.stream._read_buffer_size)
        self.socket = DeflateSocket(self.socket, self.deflate_level)
        self.socket.bootstrap(existing_data)
        self.stream.socket = self.socket

    def send_rdy(self, value):
        try:
            self.send(protocol.ready(value))
        except Exception as e:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.SendError('failed to send RDY %d' % value, e),
            )
            return False
        self.last_rdy = value
        self.rdy = value
        return True

    def _on_connect(self, **kwargs):
        identify_data = {
            'short_id': self.short_hostname, # TODO remove when deprecating pre 1.0 support
            'long_id': self.hostname, # TODO remove when deprecating pre 1.0 support
            'client_id': self.short_hostname,
            'hostname': self.hostname,
            'heartbeat_interval': self.heartbeat_interval,
            'feature_negotiation': True,
            'tls_v1': self.tls_v1,
            'snappy': self.snappy,
            'deflate': self.deflate,
            'deflate_level': self.deflate_level,
            'output_buffer_timeout': self.output_buffer_timeout,
            'output_buffer_size': self.output_buffer_size,
            'sample_rate': self.sample_rate,
            'user_agent': self.user_agent
        }
        if self.msg_timeout:
            identify_data['msg_timeout'] = self.msg_timeout
        self.trigger(event.IDENTIFY, conn=self, data=identify_data)
        self.on(event.RESPONSE, self._on_identify_response)
        try:
            self.send(protocol.identify(identify_data))
        except Exception as e:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.SendError('failed to bootstrap connection', e),
            )

    def _on_identify_response(self, data, **kwargs):
        self.off(event.RESPONSE, self._on_identify_response)

        if data == b'OK':
            logger.warning('nsqd version does not support feature netgotiation')
            return self.trigger(event.READY, conn=self)

        try:
            data = json.loads(data.decode('utf-8'))
        except ValueError:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.IntegrityError(
                    'failed to parse IDENTIFY response JSON from nsqd - %r' %
                    data
                ),
            )
            return

        self.trigger(event.IDENTIFY_RESPONSE, conn=self, data=data)

        if self.tls_v1 and data.get('tls_v1'):
            self._features_to_enable.append('tls_v1')
        if self.snappy and data.get('snappy'):
            self._features_to_enable.append('snappy')
        if self.deflate and data.get('deflate'):
            self._features_to_enable.append('deflate')

        if data.get('auth_required'):
            self._authentication_required = True

        if data.get('max_rdy_count'):
            self.max_rdy_count = data.get('max_rdy_count')
        else:
            # for backwards compatibility when interacting with older nsqd
            # (pre 0.2.20), default this to their hard-coded max
            logger.warn('setting max_rdy_count to default value of 2500')
            self.max_rdy_count = 2500

        self.on(event.RESPONSE, self._on_response_continue)
        self._on_response_continue(conn=self, data=None)

    def _on_response_continue(self, data, **kwargs):
        if self._features_to_enable:
            feature = self._features_to_enable.pop(0)
            if feature == 'tls_v1':
                self.upgrade_to_tls(self.tls_options)
            elif feature == 'snappy':
                self.upgrade_to_snappy()
            elif feature == 'deflate':
                self.upgrade_to_deflate()
            # the server will 'OK' after these conneciton upgrades triggering another response
            return

        self.off(event.RESPONSE, self._on_response_continue)
        if self.auth_secret and self._authentication_required:
            self.on(event.RESPONSE, self._on_auth_response)
            self.trigger(event.AUTH, conn=self, data=self.auth_secret)
            try:
                self.send(protocol.auth(self.auth_secret))
            except Exception as e:
                self.close()
                self.trigger(
                    event.ERROR,
                    conn=self,
                    error=protocol.SendError('Error sending AUTH', e),
                )
            return
        self.trigger(event.READY, conn=self)

    def _on_auth_response(self, data, **kwargs):
        try:
            data = json.loads(data.decode('utf-8'))
        except ValueError:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.IntegrityError(
                    'failed to parse AUTH response JSON from nsqd - %r' % data
                ),
            )
            return

        self.off(event.RESPONSE, self._on_auth_response)
        self.trigger(event.AUTH_RESPONSE, conn=self, data=data)
        return self.trigger(event.READY, conn=self)

    def _on_data(self, data, **kwargs):
        self.last_recv_timestamp = time.time()
        frame, data = protocol.unpack_response(data)
        if frame == protocol.FRAME_TYPE_MESSAGE:
            self.last_msg_timestamp = time.time()
            self.in_flight += 1

            message = protocol.decode_message(data)
            message.on(event.FINISH, self._on_message_finish)
            message.on(event.REQUEUE, self._on_message_requeue)
            message.on(event.TOUCH, self._on_message_touch)

            self.trigger(event.MESSAGE, conn=self, message=message)
        elif frame == protocol.FRAME_TYPE_RESPONSE and data == b'_heartbeat_':
            self.send(protocol.nop())
            self.trigger(event.HEARTBEAT, conn=self)
        elif frame == protocol.FRAME_TYPE_RESPONSE:
            self.trigger(event.RESPONSE, conn=self, data=data)
        elif frame == protocol.FRAME_TYPE_ERROR:
            self.trigger(event.ERROR, conn=self, error=protocol.Error(data))

    def _on_message_requeue(self, message, backoff=True, time_ms=-1, **kwargs):
        if backoff:
            self.trigger(event.BACKOFF, conn=self)
        else:
            self.trigger(event.CONTINUE, conn=self)

        self.in_flight -= 1
        try:
            time_ms = self.requeue_delay * message.attempts * 1000 if time_ms < 0 else time_ms
            self.send(protocol.requeue(message.id, time_ms))
        except Exception as e:
            self.close()
            self.trigger(event.ERROR, conn=self, error=protocol.SendError(
                'failed to send REQ %s @ %d' % (message.id, time_ms), e))

    def _on_message_finish(self, message, **kwargs):
        self.trigger(event.RESUME, conn=self)

        self.in_flight -= 1
        try:
            self.send(protocol.finish(message.id))
        except Exception as e:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.SendError('failed to send FIN %s' % message.id, e),
            )

    def _on_message_touch(self, message, **kwargs):
        try:
            self.send(protocol.touch(message.id))
        except Exception as e:
            self.close()
            self.trigger(
                event.ERROR,
                conn=self,
                error=protocol.SendError('failed to send TOUCH %s' % message.id, e),
            )


"""
An implementation of the basestring type for Python 3

Example use:

>>> s = b'abc'
>>> assert isinstance(s, basestring)
>>> from past.types import str as oldstr
>>> s2 = oldstr(b'abc')
>>> assert isinstance(s2, basestring)

"""

import sys

from past.utils import with_metaclass, PY2

if PY2:
    str = unicode

ver = sys.version_info[:2]

class BaseBaseString(type):
    def __instancecheck__(cls, instance):
        return isinstance(instance, (bytes, str))

    def __subclasshook__(cls, thing):
        # TODO: What should go here?
        raise NotImplemented

class basestring(with_metaclass(BaseBaseString)):
    """
    A minimal backport of the Python 2 basestring type to Py3
    """

__all__ = ['basestring']


import numpy as np

from openmdao.api import ExplicitComponent

from pycycle.constants import P_REF, R_UNIVERSAL_ENG, R_UNIVERSAL_SI, MIN_VALID_CONCENTRATION

class PropsCalcs(ExplicitComponent):
    """computes, S, H, Cp, Cv, gamma, given a converged equilibirum mixture"""

    def initialize(self):
        self.options.declare('thermo', desc='thermodynamic data object', recordable=False)

    def setup(self):

        thermo = self.options['thermo']

        self.add_input('T', val=284., units="degK", desc="Temperature")
        self.add_input('P', val=1., units='bar', desc="Pressure")
        self.add_input('n', val=np.ones(thermo.num_prod),
                       desc="molar concentration of the mixtures, last element is the total molar concentration")
        self.add_input('n_moles', val=1., desc="1/molar_mass for gaseous mixture")

        ne1 = thermo.num_element + 1
        self.add_input('result_T', val=np.ones(ne1),
                       desc="result of the linear solve for T", shape=ne1)
        self.add_input('result_P', val=np.ones(ne1),
                       desc="result of the linear solve for T", shape=ne1)

        self.add_output('h', val=1., units="cal/g", desc="enthalpy")
        self.add_output('S', val=1., units="cal/(g*degK)", desc="entropy")
        self.add_output('gamma', val=1.4, lower=1.0, upper=2.0, desc="ratio of specific heats")
        self.add_output('Cp', val=1., units="cal/(g*degK)", desc="Specific heat at constant pressure")
        self.add_output('Cv', val=1., units="cal/(g*degK)", desc="Specific heat at constant volume")
        self.add_output('rho', val=0.0004, units="g/cm**3", desc="density")

        self.add_output('R', val=1., units='(N*m)/(kg*degK)', desc='Specific gas constant')
        # self.deriv_options['check_type'] = "cs"

        # partial derivs setup
        self.declare_partials('h', ['n', 'T'])
        self.declare_partials('S', ['n', 'T', 'P'])
        self.declare_partials('S', 'n_moles')
        self.declare_partials('Cp', ['n', 'T', 'result_T'])
        self.declare_partials('rho', ['T', 'P', 'n_moles'])
        self.declare_partials('gamma', ['n', 'n_moles', 'T', 'result_T', 'result_P'])
        self.declare_partials('Cv', ['n', 'n_moles', 'T', 'result_T', 'result_P'])

        self.declare_partials('R', 'n_moles', val=R_UNIVERSAL_SI)

    def compute(self, inputs, outputs):
        thermo = self.options['thermo']
        num_prod = thermo.num_prod
        num_element = thermo.num_element

        T = inputs['T']
        P = inputs['P']
        result_T = inputs['result_T']

        nj = inputs['n'][:num_prod]
        # nj[nj<0] = 1e-10 # ensure all concentrations stay non-zero
        n_moles = inputs['n_moles']

        self.dlnVqdlnP = dlnVqdlnP = -1 + inputs['result_P'][num_element]
        self.dlnVqdlnT = dlnVqdlnT = 1 - result_T[num_element]

        self.Cp0_T = Cp0_T = thermo.Cp0(T)
        Cpf = np.sum(nj*Cp0_T)

        self.H0_T = H0_T = thermo.H0(T)
        self.S0_T = S0_T = thermo.S0(T)
        self.nj_H0 = nj_H0 = nj*H0_T

        # Cpe = 0
        # for i in range(0, num_element):
        #     for j in range(0, num_prod):
        #         Cpe -= thermo.aij[i][j]*nj[j]*H0_T[j]*self.result_T[i]
        # vectorization of this for loop for speed
        Cpe = -np.sum(np.sum(thermo.aij*nj_H0, axis=1)*result_T[:num_element])
        Cpe += np.sum(nj_H0*H0_T)  # nj*H0_T**2
        Cpe -= np.sum(nj_H0)*result_T[num_element]

        outputs['h'] = np.sum(nj_H0)*R_UNIVERSAL_ENG*T

        try:
            val = (S0_T+np.log(n_moles/nj/(P/P_REF)))
        except FloatingPointError:
            P = 1e-5
            val = (S0_T+np.log(n_moles/nj/(P/P_REF)))

        outputs['S'] = R_UNIVERSAL_ENG * np.sum(nj*val)
        outputs['Cp'] = Cp = (Cpe+Cpf)*R_UNIVERSAL_ENG
        outputs['Cv'] = Cv = Cp + n_moles*R_UNIVERSAL_ENG*dlnVqdlnT**2/dlnVqdlnP

        outputs['gamma'] = -1*Cp/Cv/dlnVqdlnP

        outputs['rho'] = P/(n_moles*R_UNIVERSAL_SI*T)*100  # 1 Bar is 100 Kpa

        outputs['R'] = R_UNIVERSAL_SI*n_moles  #(m**3 * Pa)/(mol*degK)

    def compute_partials(self, inputs, J):

        thermo = self.options['thermo']
        num_prod = thermo.num_prod
        num_element = thermo.num_element

        T = inputs['T']
        P = inputs['P']
        nj = inputs['n']
        n_moles = inputs['n_moles']
        result_T = inputs['result_T']
        result_T_last = result_T[num_element]
        result_T_rest = result_T[:num_element]

        dlnVqdlnP = -1 + inputs['result_P'][num_element]
        dlnVqdlnT = 1 - result_T_last

        Cp0_T = thermo.Cp0(T)
        Cpf = np.sum(nj * Cp0_T)

        H0_T = thermo.H0(T)
        S0_T = thermo.S0(T)
        nj_H0 = nj * H0_T

        # Cpe = 0
        # for i in range(0, num_element):
        #     for j in range(0, num_prod):
        #         Cpe -= thermo.aij[i][j]*nj[j]*H0_T[j]*self.result_T[i]
        # vectorization of this for loop for speed
        Cpe = -np.sum(np.sum(thermo.aij * nj_H0, axis=1) * result_T_rest)
        Cpe += np.sum(nj_H0 * H0_T)  # nj*H0_T**2
        Cpe -= np.sum(nj_H0) * result_T_last

        Cp = (Cpe + Cpf) * R_UNIVERSAL_ENG
        Cv = Cp + n_moles * R_UNIVERSAL_ENG * dlnVqdlnT ** 2 / dlnVqdlnP

        dH0_dT = thermo.H0_applyJ(T, 1.)
        dS0_dT = thermo.S0_applyJ(T, 1.)
        dCp0_dT = thermo.Cp0_applyJ(T, 1.)
        sum_nj_R = n_moles*R_UNIVERSAL_SI

        drho_dT = P/(sum_nj_R*T**2)*100
        drho_dnmoles = -P/(n_moles**2*R_UNIVERSAL_SI*T)*100

        dCpe_dT = 2*np.sum(nj*H0_T*dH0_dT)
        # for i in range(num_element):
        #     self.dCpe_dT -= np.sum(aij[i]*nj*self.dH0_dT)*self.result_T[i]
        dCpe_dT -= np.sum(np.sum(thermo.aij*nj*dH0_dT, axis=1)*result_T_rest)
        dCpe_dT -= np.sum(nj*dH0_dT)*result_T_last

        dCpf_dT = np.sum(nj*dCp0_dT)

        J['h', 'T'] = R_UNIVERSAL_ENG*(np.sum(nj*dH0_dT)*T + np.sum(nj*H0_T))
        J['h', 'n'] = R_UNIVERSAL_ENG*T*H0_T

        J['S', 'n'] = R_UNIVERSAL_ENG*(S0_T + np.log(n_moles) - np.log(P/P_REF) - np.log(nj) - 1)
        # zero out any derivs w.r.t trace species
        _trace = np.where(nj <= MIN_VALID_CONCENTRATION+1e-20)
        J['S', 'n'][0, _trace] = 0
        J['S', 'T'] = R_UNIVERSAL_ENG*np.sum(nj*dS0_dT)
        J['S', 'P'] = -R_UNIVERSAL_ENG*np.sum(nj/P)
        J['S', 'n_moles'] = R_UNIVERSAL_ENG*np.sum(nj)/n_moles
        J['rho', 'T'] = -P/(sum_nj_R*T**2)*100
        J['rho', 'n_moles'] = -P/(n_moles**2*R_UNIVERSAL_SI*T)*100
        J['rho', 'P'] = 1/(sum_nj_R*T)*100

        dCp_dnj = R_UNIVERSAL_ENG*(Cp0_T + H0_T**2)
        for j in range(num_prod):
            for i in range(num_element):
                dCp_dnj[j] -= R_UNIVERSAL_ENG*thermo.aij[i][j]*H0_T[j]*result_T[i]
        dCp_dnj -= R_UNIVERSAL_ENG * H0_T * result_T_last
        J['Cp', 'n'] = dCp_dnj

        dCp_dresultT = np.zeros(num_element+1, dtype=inputs._data.dtype)
        # for i in range(num_element):
        #     self.dCp_dresultT[i] = -R_UNIVERSAL_ENG*np.sum(aij[i]*nj_H0)
        dCp_dresultT[:num_element] = -R_UNIVERSAL_ENG*np.sum(thermo.aij*nj_H0, axis=1)
        dCp_dresultT[num_element] = - R_UNIVERSAL_ENG*np.sum(nj_H0)
        J['Cp', 'result_T'] = dCp_dresultT

        dCp_dT = (dCpe_dT + dCpf_dT)*R_UNIVERSAL_ENG
        J['Cp', 'T'] = dCp_dT

        J['Cv', 'n'] = dCp_dnj

        dCv_dnmoles = R_UNIVERSAL_ENG*dlnVqdlnT**2/dlnVqdlnP
        J['Cv', 'n_moles'] = dCv_dnmoles
        J['Cv', 'T'] = dCp_dT

        dCv_dresultP = np.zeros((1, num_element+1), dtype=inputs._data.dtype)
        dCv_dresultP[0, -1] = -R_UNIVERSAL_ENG*n_moles*(dlnVqdlnT/dlnVqdlnP)**2
        J['Cv', 'result_P'] = dCv_dresultP

        dCv_dresultT = dCp_dresultT.copy()
        dCv_dresultT[-1] -= n_moles*R_UNIVERSAL_ENG/dlnVqdlnP*(2*dlnVqdlnT)
        dCv_dresultT_last = dCv_dresultT[-1]
        J['Cv', 'result_T'] = dCv_dresultT

        J['gamma', 'n'] = dCp_dnj*(Cp/Cv-1)/(dlnVqdlnP*Cv)
        J['gamma', 'n_moles'] = Cp/dlnVqdlnP/Cv**2*dCv_dnmoles
        J['gamma', 'T'] = dCp_dT/dlnVqdlnP/Cv*(Cp/Cv-1)

        dgamma_dresultT = np.zeros((1, num_element+1), dtype=inputs._data.dtype)
        dgamma_dresultT[0, :num_element] = 1/Cv/dlnVqdlnP*dCp_dresultT[:num_element]*(Cp/Cv-1)
        dgamma_dresultT[0, -1] = (-dCp_dresultT[-1]/Cv+Cp/Cv**2*dCv_dresultT_last)/dlnVqdlnP
        J['gamma', 'result_T'] = dgamma_dresultT

        gamma_dresultP = np.zeros((1, num_element+1), dtype=inputs._data.dtype)
        gamma_dresultP[0, num_element] = Cp/Cv/dlnVqdlnP*(dCv_dresultP[0, -1]/Cv + 1/dlnVqdlnP)
        J['gamma', 'result_P'] = gamma_dresultP

if __name__ == "__main__":

    from openmdao.api import Problem, Group, IndepVarComp

    from pycycle.cea import species_data

    thermo = species_data.Properties(species_data.co2_co_o2)

    p = Problem()
    model = p.model = Group()

    indeps = model.add_subsystem('indeps', IndepVarComp(), promotes=['*'])
    indeps.add_output('T', 2761.56784655, units='degK')
    indeps.add_output('P', 1.034210, units='bar')
    indeps.add_output('n', val=np.array([2.272e-02, 1.000e-10, 1.136e-02]))
    indeps.add_output('n_moles', val=0.0340831628675)

    indeps.add_output('result_T', val=np.array([-3.02990116, 1.95459777, -0.05024694]))
    indeps.add_output('result_P', val=np.array([0.53047724, 0.48627081, -0.00437025]))

    model.add_subsystem('calcs', PropsCalcs(thermo=thermo), promotes=['*'])

    p.setup()
    p.run_model()

    print("outputs")
    print('h', p['h'])
    print('S', p['S'])
    print('gamma', p['gamma'])
    print('Cp', p['Cp'])
    print('Cv', p['Cv'])
    print('rho', p['rho'])
    print()
    print()
    print('############################################')

    p.model.run_linearize()

    jac = p.model.get_subsystem('calcs').jacobian._subjacs
    for pair in jac:
        print(pair)
        print(jac[pair])
        print


import numpy as np
from PIL import Image
from tqdm import tqdm

import torch
from torch import nn, optim
from torch.autograd import Variable, grad
from torchvision import utils

from model import Generator, Discriminator

from datetime import datetime
import random
import copy

import os

import config
import utils
import data
import evaluate

import torch.backends.cudnn as cudnn
cudnn.benchmark = True

from torch.nn import functional as F

args   = config.get_config()
writer = None

def batch_size(reso):
    if args.gpu_count == 1:
        save_memory = False
        if not save_memory:
            batch_table = {4:128, 8:128, 16:128, 32:64, 64:32, 128:16, 256:8, 512:4, 1024:1}
        else:
            batch_table = {4:128, 8:128, 16:128, 32:32, 64:16, 128:4, 256:2, 512:2, 1024:1}
    elif args.gpu_count == 2:
        batch_table = {4:256, 8:256, 16:256, 32:128, 64:64, 128:32, 256:16, 512:8, 1024:2}
    elif args.gpu_count == 4:
        batch_table = {4:512, 8:256, 16:128, 32:64, 64:32, 128:32, 256:32, 512:16, 1024:4}
    elif args.gpu_count == 8:
        batch_table = {4:512, 8:512, 16:512, 32:256, 64:256, 128:128, 256:64, 512:32, 1024:8}
    else:
        assert(False)

    return batch_table[reso]

def batch_size_by_phase(phase):
    return batch_size(4 * 2 ** phase)

class Session:
    def __init__(self):
        # Note: 4 requirements for sampling from pre-existing models:
        # 1) Ensure you save and load both multi-gpu versions (DataParallel) or both not.
        # 2) Ensure you set the same phase value as the pre-existing model and that your local and global alpha=1.0 are set
        # 3) Sample from the g_running, not from the latest generator
        # 4) You may need to warm up the g_running by running evaluate.reconstruction_dryrun() first

        self.alpha = -1
        self.sample_i = min(args.start_iteration, 0)
        self.phase = args.start_phase

        self.generator = nn.DataParallel( Generator(args.nz+1, args.n_label).cuda() )
        self.g_running = nn.DataParallel( Generator(args.nz+1, args.n_label).cuda() )
        self.encoder   = nn.DataParallel( Discriminator(nz = args.nz+1, n_label = args.n_label, binary_predictor = args.train_mode == config.MODE_GAN).cuda() )

        print("Using ", torch.cuda.device_count(), " GPUs!")

        self.reset_opt()

        print('Session created.')

    def reset_opt(self):
        self.optimizerG = optim.Adam(self.generator.parameters(), args.lr, betas=(0.0, 0.99))
        self.optimizerD = optim.Adam(self.encoder.parameters(), args.lr, betas=(0.0, 0.99)) # includes all the encoder parameters...

    def save_all(self, path):
        torch.save({'G_state_dict': self.generator.state_dict(),
                    'D_state_dict': self.encoder.state_dict(),
                    'G_running_state_dict': self.g_running.state_dict(),
                    'optimizerD': self.optimizerD.state_dict(),
                    'optimizerG': self.optimizerG.state_dict(),
                    'iteration': self.sample_i,
                    'phase': self.phase,
                    'alpha': self.alpha},
                path)

    def load(self, path):
        checkpoint = torch.load(path)
        self.sample_i = int(checkpoint['iteration'])

        self.generator.load_state_dict(checkpoint['G_state_dict'])
        self.g_running.load_state_dict(checkpoint['G_running_state_dict'])
        self.encoder.load_state_dict(checkpoint['D_state_dict'])

        if args.reset_optimizers <= 0:
            self.optimizerD.load_state_dict(checkpoint['optimizerD'])
            self.optimizerG.load_state_dict(checkpoint['optimizerG'])
            print("Reloaded old optimizers")
        else:
            print("Despite loading the state, we reset the optimizers.")

        self.alpha = checkpoint['alpha']
        self.phase = int(checkpoint['phase'])
        if args.start_phase > 0: #If the start phase has been manually set, try to actually use it (e.g. when have trained 64x64 for extra rounds and then turning the model over to 128x128)
            self.phase = min(args.start_phase, self.phase)
            print("Use start phase: {}".format(self.phase))
        if self.phase > args.max_phase:
            print('Warning! Loaded model claimed phase {} but max_phase={}'.format(self.phase, args.max_phase))
            self.phase = args.max_phase

    def create(self):
        if args.start_iteration <= 0:
            args.start_iteration = 1
            if args.no_progression:
                self.sample_i = args.start_iteration = int( (args.max_phase + 0.5) * args.images_per_stage ) # Start after the fade-in stage of the last iteration
                args.force_alpha = 1.0
                print("Progressive growth disabled. Setting start step = {} and alpha = {}".format(args.start_iteration, args.force_alpha))
        else:
            reload_from = '{}/checkpoint/{}_state'.format(args.save_dir, str(args.start_iteration).zfill(6)) #e.g. '604000' #'600000' #latest'
            print(reload_from)
            if os.path.exists(reload_from):
                self.load(reload_from)
                print("Loaded {}".format(reload_from))
                print("Iteration asked {} and got {}".format(args.start_iteration, self.sample_i))

                if args.testonly:
                    self.generator = copy.deepcopy(self.g_running)
            else:
                assert(not args.testonly)
                self.sample_i = args.start_iteration
                print('Start from iteration {}'.format(self.sample_i))

        self.g_running.train(False)

        if args.force_alpha >= 0.0:
            self.alpha = args.force_alpha

        accumulate(self.g_running, self.generator, 0)

def setup():
    utils.make_dirs()
    if not args.testonly:
        config.log_args(args)

    if args.use_TB:
        from dateutil import tz
        from tensorboardX import SummaryWriter

        dt = datetime.now(tz.gettz('Europe/Helsinki')).strftime(r"%y%m%d_%H%M")
        global writer
        writer = SummaryWriter("{}/{}/{}".format(args.summary_dir, args.save_dir, dt))

    random.seed(args.manual_seed)
    torch.manual_seed(args.manual_seed)
    torch.cuda.manual_seed_all(args.manual_seed)

def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(1 - decay, par2[k].data)

def get_grad_penalty(discriminator, real_image, fake_image, step, alpha):
    """ Used in WGAN-GP version only. """
    eps = torch.rand(batch_size_by_phase(step), 1, 1, 1).cuda()

    if eps.size(0) != real_image.size(0) or eps.size(0) != fake_image.size(0):
        # If end-of-batch situation, we restrict other vectors to matcht the number of training images available.
        eps = eps[:real_image.size(0)]
        fake_image = fake_image[:real_image.size(0)]

    x_hat = eps * real_image.data + (1 - eps) * fake_image.data
    x_hat = Variable(x_hat, requires_grad=True)

    if args.train_mode == config.MODE_GAN: # Regular GAN mode
        hat_predict, _ = discriminator(x_hat, step, alpha, args.use_ALQ)
        grad_x_hat = grad(
            outputs=hat_predict.sum(), inputs=x_hat, create_graph=True)[0]
    else:
        hat_z = discriminator(x_hat, step, alpha, args.use_ALQ)
        # KL_fake: \Delta( e(g(Z)) , Z ) -> max_e
        KL_maximizer = KLN01Loss(direction=args.KL, minimize=False)
        KL_fake = KL_maximizer(hat_z) * args.fake_D_KL_scale
        grad_x_hat = grad(
            outputs=KL_fake.sum(), inputs=x_hat, create_graph=True)[0]

    # Push the gradients of the interpolated samples towards 1
    grad_penalty = ((grad_x_hat.view(grad_x_hat.size(0), -1)
                        .norm(2, dim=1) - 1)**2).mean()
    grad_penalty = 10 * grad_penalty
    return grad_penalty

def D_prediction_of_G_output(generator, encoder, step, alpha):
    # To use labels, enable here and elsewhere:
    #label = Variable(torch.ones(batch_size_by_phase(step), args.n_label)).cuda()
    #               label = Variable(
    #                    torch.multinomial(
    #                        torch.ones(args.n_label), args.batch_size, replacement=True)).cuda()

    myz = Variable(torch.randn(batch_size_by_phase(step), args.nz)).cuda(non_blocking=(args.gpu_count>1))
    myz = utils.normalize(myz)
    myz, label = utils.split_labels_out_of_latent(myz)

    fake_image = generator(myz, label, step, alpha)
    fake_predict, _ = encoder(fake_image, step, alpha, args.use_ALQ)

    loss = fake_predict.mean()
    return loss, fake_image

class KLN01Loss(torch.nn.Module): #Adapted from https://github.com/DmitryUlyanov/AGE

    def __init__(self, direction, minimize):
        super(KLN01Loss, self).__init__()
        self.minimize = minimize
        assert direction in ['pq', 'qp'], 'direction?'

        self.direction = direction

    def forward(self, samples):

        assert samples.nelement() == samples.size(1) * samples.size(0), '?'

        samples = samples.view(samples.size(0), -1)

        self.samples_var = utils.var(samples)
        self.samples_mean = samples.mean(0)

        samples_mean = self.samples_mean
        samples_var = self.samples_var

        if self.direction == 'pq':
            t1 = (1 + samples_mean.pow(2)) / (2 * samples_var.pow(2))
            t2 = samples_var.log()

            KL = (t1 + t2 - 0.5).mean()
        else:
            # In the AGE implementation, there is samples_var^2 instead of samples_var^1
            t1 = (samples_var + samples_mean.pow(2)) / 2
            # In the AGE implementation, this did not have the 0.5 scaling factor:
            t2 = -0.5*samples_var.log()

            KL = (t1 + t2 - 0.5).mean()

        if not self.minimize:
            KL *= -1

        return KL

def train(generator, encoder, g_running, train_data_loader, test_data_loader, session, total_steps, train_mode):
    pbar = tqdm(initial=session.sample_i, total = total_steps)

    benchmarking = False

    match_x = args.match_x
    generatedImagePool = None

    refresh_dataset   = True
    refresh_imagePool = True

    # After the Loading stage, we cycle through successive Fade-in and Stabilization stages

    batch_count = 0

    reset_optimizers_on_phase_start = False

    # TODO Unhack this (only affects the episode count statistics anyway):
    if args.data != 'celebaHQ':
        epoch_len = len(train_data_loader(1,4).dataset)
    else:
        epoch_len = train_data_loader._len['data4x4']

    if args.step_offset != 0:
        if args.step_offset == -1:
            args.step_offset = session.sample_i
        print("Step offset is {}".format(args.step_offset))
        session.phase += args.phase_offset
        session.alpha = 0.0

    while session.sample_i < total_steps:
        #######################  Phase Maintenance #######################

        steps_in_previous_phases = max(session.phase * args.images_per_stage, args.step_offset)

        sample_i_current_stage = session.sample_i - steps_in_previous_phases

        # If we can move to the next phase
        if sample_i_current_stage >= args.images_per_stage:
            if session.phase < args.max_phase: # If any phases left
                iteration_levels = int(sample_i_current_stage / args.images_per_stage)
                session.phase += iteration_levels
                sample_i_current_stage -= iteration_levels * args.images_per_stage
                match_x = args.match_x # Reset to non-matching phase
                print("iteration B alpha={} phase {} will be reduced to 1 and [max]".format(sample_i_current_stage, session.phase))

                refresh_dataset = True
                refresh_imagePool = True # Reset the pool to avoid images of 2 different resolutions in the pool

                if reset_optimizers_on_phase_start:
                    utils.requires_grad(generator)
                    utils.requires_grad(encoder)
                    generator.zero_grad()
                    encoder.zero_grad()
                    session.reset_opt()
                    print("Optimizers have been reset.")

        reso = 4 * 2 ** session.phase

        # If we can switch from fade-training to stable-training
        if sample_i_current_stage >= args.images_per_stage/2:
            if session.alpha < 1.0:
                refresh_dataset = True # refresh dataset generator since no longer have to fade
            match_x = args.match_x * args.matching_phase_x
        else:
            match_x = args.match_x

        session.alpha = min(1, sample_i_current_stage * 2.0 / args.images_per_stage) # For 100k, it was 0.00002 = 2.0 / args.images_per_stage

        if refresh_dataset:
            train_dataset = data.Utils.sample_data2(train_data_loader, batch_size(reso), reso, session)
            refresh_dataset = False
            print("Refreshed dataset. Alpha={} and iteration={}".format(session.alpha, sample_i_current_stage))
        if refresh_imagePool:
            imagePoolSize = 200 if reso < 256 else 100
            generatedImagePool = utils.ImagePool(imagePoolSize) #Reset the pool to avoid images of 2 different resolutions in the pool
            refresh_imagePool = False
            print('Image pool created with size {} because reso is {}'.format(imagePoolSize, reso))

        ####################### Training init #######################

        z = Variable( torch.FloatTensor(batch_size(reso), args.nz, 1, 1) ).cuda(non_blocking=(args.gpu_count>1))
        KL_minimizer = KLN01Loss(direction=args.KL, minimize=True)
        KL_maximizer = KLN01Loss(direction=args.KL, minimize=False)

        stats = {}

        one = torch.FloatTensor([1]).cuda(non_blocking=(args.gpu_count>1))

        try:
            real_image, _ = next(train_dataset)
        except (OSError, StopIteration):
            train_dataset = data.Utils.sample_data2(train_data_loader, batch_size(reso), reso, session)
            real_image, _ = next(train_dataset)

        ####################### DISCRIMINATOR / ENCODER ###########################

        utils.switch_grad_updates_to_first_of(encoder, generator)
        encoder.zero_grad()

        x = Variable(real_image).cuda(non_blocking=(args.gpu_count>1))
        kls = ""
        if train_mode == config.MODE_GAN:

            # Discriminator for real samples
            real_predict, _ = encoder(x, session.phase, session.alpha, args.use_ALQ)
            real_predict = real_predict.mean() \
                - 0.001 * (real_predict ** 2).mean()
            real_predict.backward(-one) # Towards 1

            # (1) Generator => D. Identical to (2) see below

            fake_predict, fake_image = D_prediction_of_G_output(generator, encoder, session.phase, session.alpha)
            fake_predict.backward(one)

            # Grad penalty

            grad_penalty = get_grad_penalty(encoder, x, fake_image, session.phase, session.alpha)
            grad_penalty.backward()

        elif train_mode == config.MODE_CYCLIC:
            e_losses = []

            # e(X)

            real_z = encoder(x, session.phase, session.alpha, args.use_ALQ)
            if args.use_real_x_KL:
                # KL_real: - \Delta( e(X) , Z ) -> max_e
                KL_real = KL_minimizer(real_z) * args.real_x_KL_scale
                e_losses.append(KL_real)

                stats['real_mean'] = KL_minimizer.samples_mean.data.mean()
                stats['real_var'] = KL_minimizer.samples_var.data.mean()
                stats['KL_real'] = KL_real.data.item()
                kls = "{0:.3f}".format(stats['KL_real'])

            # The final entries are the label. Normal case, just 1. Extract it/them, and make it [b x 1]:

            real_z, label = utils.split_labels_out_of_latent(real_z)
            recon_x = generator(real_z, label, session.phase, session.alpha)
            if args.use_loss_x_reco:
                # match_x: E_x||g(e(x)) - x|| -> min_e
                err = utils.mismatch(recon_x, x, args.match_x_metric) * match_x
                e_losses.append(err)
                stats['x_reconstruction_error'] = err.item()

            args.use_wpgan_grad_penalty = False
            grad_penalty = 0.0

            if args.use_loss_fake_D_KL:
                # TODO: The following codeblock is essentially the same as the KL_minimizer part on G side. Unify
                utils.populate_z(z, args.nz+args.n_label, args.noise, batch_size(reso))
                z = torch.squeeze(z)
                z, label = utils.split_labels_out_of_latent(z)
                fake = generator(z, label, session.phase, session.alpha).detach()

                if session.alpha >= 1.0:
                    fake = generatedImagePool.query(fake.data)

                # e(g(Z))
                egz = encoder(fake, session.phase, session.alpha, args.use_ALQ)

                # KL_fake: \Delta( e(g(Z)) , Z ) -> max_e
                KL_fake = KL_maximizer(egz) * args.fake_D_KL_scale
                e_losses.append(KL_fake)

                stats['fake_mean'] = KL_maximizer.samples_mean.data.mean()
                stats['fake_var'] = KL_maximizer.samples_var.data.mean()
                stats['KL_fake'] = -KL_fake.item()
                kls = "{0}/{1:.3f}".format(kls, stats['KL_fake'])

                if args.use_wpgan_grad_penalty:
                    grad_penalty = get_grad_penalty(encoder, x, fake, session.phase, session.alpha)

            # Update e
            if len(e_losses) > 0:
                e_loss = sum(e_losses)
                stats['E_loss'] = np.float32(e_loss.cpu().detach().numpy())
                e_loss.backward()

                if args.use_wpgan_grad_penalty:
                    grad_penalty.backward()
                    stats['Grad_penalty'] = grad_penalty.data

                #book-keeping
                disc_loss_val = e_loss.item()

        session.optimizerD.step()

        torch.cuda.empty_cache()

        ######################## GENERATOR / DECODER #############################

        if (batch_count + 1) % args.n_critic == 0:
            utils.switch_grad_updates_to_first_of(generator, encoder)

            for _ in range(args.n_generator):
                generator.zero_grad()
                g_losses = []

                if train_mode == config.MODE_GAN:
                    fake_predict, _ = D_prediction_of_G_output(generator, encoder, session.phase, session.alpha)
                    loss = -fake_predict
                    g_losses.append(loss)

                elif train_mode == config.MODE_CYCLIC: #TODO We push the z variable around here like idiots
                    def KL_of_encoded_G_output(generator, z):
                        utils.populate_z(z, args.nz+args.n_label, args.noise, batch_size(reso))
                        z, label = utils.split_labels_out_of_latent(z)
                        fake = generator(z, label, session.phase, session.alpha)

                        egz = encoder(fake, session.phase, session.alpha, args.use_ALQ)
                        # KL_fake: \Delta( e(g(Z)) , Z ) -> min_g
                        return egz, label, KL_minimizer(egz) * args.fake_G_KL_scale, z

                    egz, label, kl, z = KL_of_encoded_G_output(generator, z)

                    if args.use_loss_KL_z:
                        g_losses.append(kl) # G minimizes this KL
                        stats['KL(Phi(G))'] = kl.item()
                        kls = "{0}/{1:.3f}".format(kls, stats['KL(Phi(G))'])

                    if args.use_loss_z_reco:
                        z = torch.cat((z, label), 1)
                        z_diff = utils.mismatch(egz, z, args.match_z_metric) * args.match_z # G tries to make the original z and encoded z match
                        g_losses.append(z_diff)

                if len(g_losses) > 0:
                    loss = sum(g_losses)
                    stats['G_loss'] = np.float32(loss.cpu().detach().numpy())
                    loss.backward()

                    # Book-keeping only:
                    gen_loss_val = loss.item()

                session.optimizerG.step()

                torch.cuda.empty_cache()

                if train_mode == config.MODE_CYCLIC:
                    if args.use_loss_z_reco:
                        stats['z_reconstruction_error'] = z_diff.item()

            accumulate(g_running, generator)

            del z, x, one, real_image, real_z, KL_real, label, recon_x, fake, egz, KL_fake, kl, z_diff

            if train_mode == config.MODE_CYCLIC:
                if args.use_TB:
                    for key,val in stats.items():
                        writer.add_scalar(key, val, session.sample_i)
                elif batch_count % 100 == 0:
                    print(stats)

            if args.use_TB:
                writer.add_scalar('LOD', session.phase + session.alpha, session.sample_i)

        ########################  Statistics ########################

        b = batch_size_by_phase(session.phase)
        zr, xr = (stats['z_reconstruction_error'], stats['x_reconstruction_error']) if train_mode == config.MODE_CYCLIC else (0.0, 0.0)
        e = (session.sample_i / float(epoch_len))
        pbar.set_description(
            ('{0}; it: {1}; phase: {2}; b: {3:.1f}; Alpha: {4:.3f}; Reso: {5}; E: {6:.2f}; KL(real/fake/fakeG): {7}; z-reco: {8:.2f}; x-reco {9:.3f}; real_var {10:.4f}').format(batch_count+1, session.sample_i+1, session.phase, b, session.alpha, reso, e, kls, zr, xr, stats['real_var'])
            )
            #(f'{i + 1}; it: {iteration+1}; b: {b:.1f}; G: {gen_loss_val:.5f}; D: {disc_loss_val:.5f};'
            # f' Grad: {grad_loss_val:.5f}; Alpha: {alpha:.3f}; Reso: {reso}; S-mean: {real_mean:.3f}; KL(real/fake/fakeG): {kls}; z-reco: {zr:.2f}'))

        pbar.update(batch_size(reso))
        session.sample_i += batch_size(reso) # if not benchmarking else 100
        batch_count += 1

        ########################  Saving ########################

        if batch_count % args.checkpoint_cycle == 0:
            for postfix in {'latest', str(session.sample_i).zfill(6)}:
                session.save_all('{}/{}_state'.format(args.checkpoint_dir, postfix))

            print("Checkpointed to {}".format(session.sample_i))

        ########################  Tests ########################

        try:
            evaluate.tests_run(g_running, encoder, test_data_loader, session, writer,
            reconstruction       = (batch_count % 800 == 0),
            interpolation        = (batch_count % 800 == 0),
            collated_sampling    = (batch_count % 800 == 0),
            individual_sampling  = (batch_count % (args.images_per_stage/batch_size(reso)/4) == 0)
            )
        except (OSError, StopIteration):
            print("Skipped periodic tests due to an exception.")

    pbar.close()

def main():
    setup()
    session = Session()
    session.create()

    print('PyTorch {}'.format(torch.__version__))

    if args.train_path:
        train_data_loader = data.get_loader(args.data, args.train_path)
    else:
        train_data_loader = None

    if args.test_path:
        test_data_loader = data.get_loader(args.data, args.test_path)
    elif args.aux_inpath:
        test_data_loader = data.get_loader(args.data, args.aux_inpath)
    else:
        test_data_loader = None

    # 4 modes: Train (with data/train), test (with data/test), aux-test (with custom aux_inpath), dump-training-set

    if args.run_mode == config.RUN_TRAIN:
        train(session.generator, session.encoder, session.g_running, train_data_loader, test_data_loader,
            session  = session,
            total_steps = args.total_kimg * 1000,
            train_mode = args.train_mode)
    elif args.run_mode == config.RUN_TEST:
        if args.reconstructions_N > 0 or args.interpolate_N > 0:
            evaluate.Utils.reconstruction_dryrun(session.generator, session.encoder, test_data_loader, session=session)
        evaluate.tests_run(session.generator, session.encoder, test_data_loader, session=session, writer=writer)
    elif args.run_mode == config.RUN_DUMP:
        session.phase = args.start_phase
        data.dump_training_set(train_data_loader, args.dump_trainingset_N, args.dump_trainingset_dir, session)

if __name__ == '__main__':
    main()