''' A package for implementing Caffe-like network structure using Owl APIs.

The package implements Caffe-like network structure with some minor differences. It uses Caffe-defined
protobuf as core data structure so Caffe users could easily adapt to this. The package serves the purpose
of:

1. Quick deployment of neural network training using configure file.
2. Demonstrate the power of ``owl`` package (it takes only several hundreds LOC to implement Caffe and run it on dataflow engine).
'''

import numpy as np
import math
import Queue

import owl
import owl.elewise as ele
import owl.conv as co
from caffe import *

from netio import LMDBDataProvider
from netio import ImageListDataProvider
from netio import ImageWindowDataProvider

class ComputeUnit(object):
    ''' Interface for each compute unit.

    In ``owl.net``, the network is graph (in fact a DAG) that is composed of ``ComputeUnit`` s.
    ``ComputeUnit`` is a wrap-up of Caffe's ``layer`` abstraction, but is more
    general and flexible in its function sigature.

    :ivar caffe.LayerParameter params: layer parameter in Caffe's proto structure
    :ivar str name: name of the unit; the name must be unique
    :ivar btm_names: names of the bottom units
    :vartype btm_names: list str
    :ivar top_names: names of the top units
    :vartype top_names: list str
    :ivar list int out_shape:

    .. note::
        ``params``, ``name``, ``btm_names`` and ``top_names`` will be parsed from Caffe's network
        description file. ``out_shape`` should be set in :py:meth:`compute_size`

    '''
    def __init__(self, params):
        self.params = params
        self.name = params.name
        self.btm_names = []
        self.top_names = []
        self.out = None
        self.out_shape = None
        self.rec_on_ori = None
        self.stride_on_ori = None
        self.start_on_ori = None
    def __str__(self):
        return 'N/A unit'
    def compute_size(self, from_btm, to_top):
        ''' Calculate the output size of this unit

        This function will be called before training during the ``compute_size`` phase.
        The ``compute_size`` phase is a feed-forward-like phase, during which each ``ComputeUnit``, rather than
        calculating the output tensor but calculating the output size (list int) for the top units. The
        size is usually used to calculate the weight and bias size for initialization.

        :param dict from_btm: input size from bottom units
        :param dict to_top: output size to top units

        .. seealso::
            :py:meth:`FullyConnection.compute_size`
            :py:meth:`ConvConnection.compute_size`
            :py:meth:`Net.compute_size`
        '''
        pass
    def forward(self, from_btm, to_top, phase):
        ''' Function for forward propagation

        This function will be called during forward-propagation. The function
        should take input in ``from_btm``, perform customized computation, and then
        put the result in ``to_top``. Both ``from_btm`` and ``to_top`` are ``dict`` type
        where key is a ``str`` of name of the bottom/top units and value is an ``owl.NArray``
        served as input or output of the function.

        :param dict from_btm: input from bottom units
        :param dict to_top: output for top units
        :param str phase: name of the phase of the running. Currently either ``"TRAIN"`` or ``"TEST"``
        '''
        pass
    def backward(self, from_top, to_btm, phase):
        ''' Function for backward propagation

        This function will be called during backward-propagation. Similar to :py:meth:`forward`,
        The function should take input in ``from_top``, perform customized computation, and then
        put the result in ``to_btm``. The function also need to calculate the gradient (if any) and
        save them to the ``weightgrad`` field (see :py:meth:WeightedComputeUnit.weight_update).

        :param dict from_top: input from top units
        :param dict to_btm: output for top units
        :param str phase: name of the phase of the running. Currently either ``"TRAIN"`` or ``"TEST"``
        '''
        pass
    def weight_update(self, base_lr, base_weight_decay, momentum, batch_size):
        ''' Function for weight update

        This function will be called during weight update.

        :param float base_lr: base learning rate
        :param float base_weight_decay: base weight decay
        :param float momentum: momentum value
        :param int batch_size: the size of the current minibatch
        '''
        pass

class ComputeUnitSimple(ComputeUnit):
    ''' An auxiliary class for :py:class:`ComputeUnit` that will only have one input unit and one output unit.
    '''
    def __init__(self, params):
        super(ComputeUnitSimple, self).__init__(params)
    def compute_size(self, from_btm, to_top):
        ''' Set the ``out_shape`` as the same shape of the input. Inherited classes could override this function.
        '''
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori']
        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori']
        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
    def forward(self, from_btm, to_top, phase):
        ''' Transform the interface from multiple input/output to only one input/output function :py:meth:`ff`.
        '''
        to_top[self.top_names[0]] = self.ff(from_btm[self.btm_names[0]], phase)
        self.out = to_top[self.top_names[0]]
    def ff(self, act, phase):
        ''' Function for forward-propagation

        :param owl.NArray act: the activation from the bottom unit
        :param str phase: name of the phase of the running. Currently either ``"TRAIN"`` or ``"TEST"``
        :return: the activation of this unit
        :rtype: owl.NArray
        '''
        pass
    def backward(self, from_top, to_btm, phase):
        ''' Transform the interface from multiple input/output to only one input/output function :py:meth:`bp`.
        '''
        to_btm[self.btm_names[0]] = self.bp(from_top[self.top_names[0]], phase)
    def bp(self, sen, phase):
        ''' Function for backward-propagation

        :param owl.NArray sen: the sensitivity (or error derivative to the input) from the top unit
        :return: the sensitivity of this unit
        :rtype: owl.NArray
        '''
        pass

class WeightedComputeUnit(ComputeUnitSimple):
    ''' An auxiliary class for :py:class:`ComputeUnit` with weights

    :ivar owl.NArray weight: weight tensor
    :ivar owl.NArray weightdelta: momentum of weight
    :ivar owl.NArray weightgrad: gradient of weight
    :ivar owl.NArray bias: bias tensor
    :ivar owl.NArray biasdelta: momentum of bias
    :ivar owl.NArray biasgrad: gradient of bias
    :ivar float lr_mult_w: learning rate multiplier for the weight of this unit
    :ivar float lr_mult_b: bias learning rate multiplier for the bias of this unit
    :ivar float decay_mult_w: decay multiplier for the weight of this unit
    :ivar float decay_mult_b: decay multiplier for the bias of this unit
    '''
    def __init__(self, params):
        super(WeightedComputeUnit, self).__init__(params)
        # weights and bias
        self.weight = None
        self.weightdelta = None
        self.weightgrad = None
        self.bias = None
        self.biasdelta = None
        self.biasgrad = None

        self.in_shape = None
        self.fan_in = None
        self.fan_out = None

        # blob learning rate and weight decay
        if len(params.param) >= 1:
            self.lr_mult_w = params.param[0].lr_mult
            self.decay_mult_w = params.param[0].decay_mult
        else:
            self.lr_mult_w = 1
            self.decay_mult_w = 1

        if len(params.param) >= 2:
            self.lr_mult_b = params.param[1].lr_mult
            self.decay_mult_b = params.param[1].decay_mult
        else:
            self.lr_mult_b = 1
            self.decay_mult_b = 0

        #self.blobs_lr = params.blobs_lr
        #self.weight_decay = params.weight_decay
        #if len(self.blobs_lr) == 0:
            #self.blobs_lr = [1,1]
        #if len(self.weight_decay) == 0:
            #self.weight_decay = [1, 0]

    def compute_size(self, from_btm, to_top):
        pass

    def init_weights_with_filler(self):
        ''' Init weights & bias. The function will be called during weight initialization.

        Currently, four types of initializers are supported: ``"constant", "gaussian", "uniform", "xavier"``.
        '''
        #init weight
        npweights = None
        if self.weight_filler.type == "constant":
            npweights = np.ones(self.wshape, dtype = np.float32) * self.weight_filler.value
        elif self.weight_filler.type == "gaussian":
            npweights = np.random.normal(self.weight_filler.mean, self.weight_filler.std, self.wshape)
        elif self.weight_filler.type == "uniform":
            npweights = np.random.uniform(self.weight_filler.min, self.weight_filler.max, self.wshape)
        elif self.weight_filler.type == "xavier":
            scale = np.sqrt(float(3)/self.fan_in)
            npweights = np.random.uniform(-scale, scale, self.wshape)
        self.weight = owl.from_numpy(npweights.astype(np.float32)).reshape(self.wshape)

        #init bias
        npwbias = None
        if self.bias_filler.type == "constant":
            npbias = np.ones(self.bshape, dtype = np.float32) * self.bias_filler.value
        elif self.bias_filler.type == "gaussian":
            npbias = np.random.normal(self.bias_filler.mean, self.bias_filler.std, self.bshape)
        elif self.bias_filler.type == "uniform":
            npbias = np.random.uniform(self.bias_filler.min, self.bias_filler.max, self.bshape)
        elif self.bias_filler.type == "xavier":
            scale = np.sqrt(float(3)/self.fan_in)
            npbias = np.random.uniform(-scale, scale, self.bshape)
        self.bias = owl.from_numpy(npbias.astype(np.float32)).reshape(self.bshape)

    def weight_update(self, base_lr, base_weight_decay, momentum, batch_size):
        ''' Update the weight & bias

        Using following formula:

        ``$_delta = momentum * $_delta - (base_lr * $_lr / batch_size) * $_grad - (base_lr * $_lr * base_wd * $_wd) * $``

        , where ``$`` could be either ``weight`` or ``bias``.
        '''
        if self.weightdelta == None:
            self.weightdelta = owl.zeros(self.weightgrad.shape)

        self.weightdelta = momentum * self.weightdelta \
                        - (base_lr * self.lr_mult_w / batch_size) * self.weightgrad \
                        - (base_lr * self.lr_mult_w * base_weight_decay * self.decay_mult_w) * self.weight

        self.weight = self.weight + self.weightdelta
        self.weightgrad = None

        if self.biasdelta == None:
            self.biasdelta = owl.zeros(self.biasgrad.shape)

        self.biasdelta = momentum * self.biasdelta \
                        - (base_lr * self.lr_mult_b / batch_size) * self.biasgrad \
                        - (base_lr * self.lr_mult_b * base_weight_decay * self.decay_mult_b) * self.bias
        self.bias = self.bias + self.biasdelta
        self.biasgrad = None

class LinearUnit(ComputeUnitSimple):
    ''' Compute unit for linear transformation
    '''
    def ff(self, x, phase):
        return x
    def bp(self, y, phase):
        return y
    def __str__(self):
        return 'linear'

class SigmoidUnit(ComputeUnitSimple):
    ''' Compute unit for Sigmoid non-linearity
    '''
    def ff(self, x, phase):
        return ele.sigm(x)
    def bp(self, y, phase):
        return ele.sigm_back(y)
    def __str__(self):
        return 'sigmoid'

class ReluUnit(ComputeUnitSimple):
    ''' Compute unit for RELU non-linearity
    '''
    def ff(self, x, phase):
        self.ff_x = x
        return ele.relu(x)
    def bp(self, y, phase):
        return ele.relu_back(y, self.ff_x)
    def __str__(self):
        return 'relu'

class TanhUnit(ComputeUnitSimple):
    ''' Compute unit for Hyperbolic Tangine non-linearity
    '''
    def ff(self, x, phase):
        return ele.tanh(x)
    def bp(self, y, phase):
        return ele.tanh_back(y)
    def __str__(self):
        return 'tanh'

class PoolingUnit(ComputeUnitSimple):
    ''' Compute unit for Pooling

    .. note::
        The input and output is of size ``[HWCN]``:

        - ``H``: image height
        - ``W``: image width
        - ``C``: number of image channels (feature maps)
        - ``N``: size of minibatch
    '''
    def __init__(self, params):
        super(PoolingUnit, self).__init__(params)
        self.ppa = params.pooling_param
        if self.ppa.pool == PoolingParameter.PoolMethod.Value('MAX'):
            pool_ty = co.pool_op.max
        elif self.ppa.pool == PoolingParameter.PoolMethod.Value('AVE'):
            pool_ty = co.pool_op.average
        self.pooler = co.Pooler(self.ppa.kernel_size, self.ppa.kernel_size,
                                self.ppa.stride, self.ppa.stride,
                                self.ppa.pad, self.ppa.pad,
                                pool_ty)

    def compute_size(self, from_btm, to_top):
        self.out_shape = from_btm[self.btm_names[0]]['out_shape'][:]
        ori_height = self.out_shape[0]
        ori_width = self.out_shape[1]
        self.out_shape[0] = int(np.ceil(float(self.out_shape[0] + 2 * self.ppa.pad - self.ppa.kernel_size) / self.ppa.stride)) + 1
        self.out_shape[1] = int(np.ceil(float(self.out_shape[1] + 2 * self.ppa.pad - self.ppa.kernel_size) / self.ppa.stride)) + 1
        if self.ppa.pad:
            if (self.out_shape[0] - 1) * self.ppa.stride >= ori_height + self.ppa.pad:
                self.out_shape[0] = self.out_shape[0] - 1
                self.out_shape[1] = self.out_shape[1] - 1
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = self.out_shape[:]
        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori'] * self.ppa.stride
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori'] + (self.ppa.kernel_size - 1) * from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori'] - self.ppa.pad * from_btm[self.btm_names[0]]['stride_on_ori']
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']

    def ff(self, x, phase):
        self.ff_x = x
        self.ff_y = self.pooler.ff(x)
        return self.ff_y
    def bp(self, y, phase):
        return self.pooler.bp(y, self.ff_y, self.ff_x)
    def __str__(self):
        return 'pooling'

class DropoutUnit(ComputeUnitSimple):
    ''' Compute unit for dropout
    '''
    def __init__(self, params):
        super(DropoutUnit, self).__init__(params)
        self.scale = 1.0 / (1.0 - self.params.dropout_param.dropout_ratio)
        self.keep_ratio = 1 - self.params.dropout_param.dropout_ratio
    def ff(self, x, phase):
        ''' Foward function of dropout

        The dropout mask will not be multiplied if under ``"TEST"`` mode.
        '''
        if phase == "TRAIN":
            self.dropmask = owl.randb(x.shape, self.keep_ratio)
            return ele.mult(x, self.dropmask)*self.scale
        else:
            return x
        #for gradient test
        #return x
    def bp(self, y, phase):
        if phase == "TRAIN":
            return ele.mult(y, self.dropmask)*self.scale
        else:
            return y
    def __str__(self):
        return 'dropout'

class SoftmaxUnit(ComputeUnit):
    ''' Compute unit for softmax
    '''
    def __init__(self, params):
        super(SoftmaxUnit, self).__init__(params)
        self.loss_weight = params.loss_weight

    def compute_size(self, from_btm, to_top):
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori']
        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori']

        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']

    def forward(self, from_btm, to_top, phase):
        to_top[self.top_names[0]] = co.softmax(from_btm[self.btm_names[0]], co.soft_op.instance)
        self.ff_y = to_top[self.top_names[0]]
        #turn label into matrix form
        nplabel = np.zeros([self.ff_y.shape[1], self.ff_y.shape[0]], dtype=np.float32)
        self.strlabel = from_btm[self.btm_names[1]]

        for i in range(len(self.strlabel)):
            nplabel[i, self.strlabel[i]] = 1
        self.y = owl.from_numpy(nplabel)
        self.out = self.ff_y

    def backward(self, from_top, to_btm, phase):
        if len(self.loss_weight) == 1:
            to_btm[self.btm_names[0]] = (self.ff_y - self.y)*self.loss_weight[0]
        else:
            to_btm[self.btm_names[0]] = (self.ff_y - self.y)

    def getloss(self):
        ''' Get the loss of the softmax (cross entropy)
        '''
        lossmat = ele.mult(ele.ln(self.ff_y), self.y)
        res = lossmat.sum(0).sum(1).to_numpy()
        return -res[0][0] / lossmat.shape[1]

    def __str__(self):
        return 'softmax'

class AccuracyUnit(ComputeUnit):
    ''' Compute unit for calculating accuracy

    .. note::
        In terms of Minerva's lazy evaluation, the unit is a **non-lazy** one since it gets the actual
        contents (accuracy) out of an ``owl.NArray``.
    '''
    def __init__(self, params):
        super(AccuracyUnit, self).__init__(params)
        self.acc = 0
        self.batch_size = 0
        self.top_k = params.accuracy_param.top_k

    def compute_size(self, from_btm, to_top):
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori']
        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori']

        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']

    def forward(self, from_btm, to_top, phase):
        if self.top_k == 1:
            predict = from_btm[self.btm_names[0]].max_index(0)
            ground_truth = owl.from_numpy(from_btm[self.btm_names[1]]).reshape(predict.shape)
            self.batch_size = from_btm[self.btm_names[0]].shape[1]
            correct = (predict - ground_truth).count_zero()
            self.acc = correct * 1.0 / self.batch_size
        elif self.top_k == 5:
            predict = from_btm[self.btm_names[0]].to_numpy()
            top_5 = np.argsort(predict, axis=1)[:,::-1]
            ground_truth = from_btm[self.btm_names[1]]
            self.batch_size = np.shape(ground_truth)[0]
            correct = 0
            for i in range(self.batch_size):
                for t in range(5):
                    if ground_truth[i] == top_5[i,t]:
                        correct += 1
                        break
            self.acc = correct * 1.0 / self.batch_size
        else:
            assert(FALSE)

    def backward(self, from_top, to_btm, phase):
        pass

    def __str__(self):
        return 'accuracy'

class LRNUnit(ComputeUnitSimple):
    ''' Compute unit for LRN
    '''
    def __init__(self, params):
        super(LRNUnit, self).__init__(params)
        self.lrner = co.Lrner(params.lrn_param.local_size, params.lrn_param.alpha, params.lrn_param.beta)
        self.scale = None
    def ff(self, x, phase):
        self.ff_x = x
        self.scale = owl.zeros(x.shape)
        self.ff_y = self.lrner.ff(x, self.scale)
        return self.ff_y
    def bp(self, y, phase):
        return self.lrner.bp(self.ff_x, self.ff_y, self.scale, y)
    def __str__(self):
        return 'lrn'

class ConcatUnit(ComputeUnit):
    ''' Compute unit for concatenation

    Concatenate input arrays along the dimension specified by Caffe's ``concat_dim_caffe``
    '''
    def __init__(self, params):
        super(ConcatUnit, self).__init__(params)
        self.concat_dim_caffe = params.concat_param.concat_dim
        self.slice_count = []

    def compute_size(self, from_btm, to_top):
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori']
        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori']

        self.concat_dim = len(from_btm[self.btm_names[0]]['out_shape']) - 1 - self.concat_dim_caffe
        for i in range(1, len(self.btm_names)):
            to_top[self.top_names[0]]['out_shape'][self.concat_dim] = to_top[self.top_names[0]]['out_shape'][self.concat_dim] + from_btm[self.btm_names[i]]['out_shape'][self.concat_dim]
        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']

    def forward(self, from_btm, to_top, phase):
        narrays = []
        self.concat_dim = len(from_btm[self.btm_names[0]].shape) - 1 - self.concat_dim_caffe
        for i in range(len(self.btm_names)):
            narrays.append(from_btm[self.btm_names[i]])
            self.slice_count.append(from_btm[self.btm_names[i]].shape[self.concat_dim])
        to_top[self.top_names[0]] = owl.concat(narrays, self.concat_dim)
        self.out = to_top[self.top_names[0]]

    def backward(self, from_top, to_btm, phase):
        st_off = 0
        for i in range(len(self.btm_names)):
            to_btm[self.btm_names[i]] = owl.slice(from_top[self.top_names[0]],
                                                  self.concat_dim,
                                                  st_off,
                                                  self.slice_count[i])
            st_off += self.slice_count[i]
    def __str__(self):
        return 'concat'

class FullyConnection(WeightedComputeUnit):
    ''' Compute unit for traditional fully connected layer
    '''
    def __init__(self, params):
        super(FullyConnection, self).__init__(params)
        self.inner_product_param = params.inner_product_param
        self.weight_filler = params.inner_product_param.weight_filler
        self.bias_filler = params.inner_product_param.bias_filler

    def compute_size(self, from_btm, to_top):
        ''' Compute the output size and also weight and bias size
        The weight size is ``[top_shape[0], btm_shape[0]]``; the bias size is ``[top_shape[0], 1]``
        (assume both ``top`` and ``btm`` are 2-dimensional array)
        '''
        shp = from_btm[self.btm_names[0]]['out_shape'][:]
        if len(shp) > 2:
            self.in_shape = [np.prod(shp[0:-1], dtype=np.int32), shp[-1]]
        else:
            self.in_shape = shp
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = self.in_shape[:]
        to_top[self.top_names[0]]['out_shape'][0] = self.inner_product_param.num_output
        to_top[self.top_names[0]]['out_shape'][1] = 1
        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.wshape = [self.out_shape[0], self.in_shape[0]]
        self.bshape = [self.out_shape[0], 1]

        if len(shp) > 2:
            #last layer is conv layer
            self.rec_on_ori = from_btm[self.btm_names[0]]['rec_on_ori'] + (shp[0] - 1) * from_btm[self.btm_names[0]]['stride_on_ori']
            self.stride_on_ori = self.rec_on_ori
        else:
            self.rec_on_ori = from_btm[self.btm_names[0]]['rec_on_ori']
            self.stride_on_ori = from_btm[self.btm_names[0]]['stride_on_ori']

        to_top[self.top_names[0]]['rec_on_ori'] = self.rec_on_ori
        to_top[self.top_names[0]]['stride_on_ori'] = self.stride_on_ori
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
        #set fan_in fan_out
        self.fan_out = self.inner_product_param.num_output
        self.fan_in = np.prod(from_btm[self.btm_names[0]]['out_shape'][0:len(from_btm[self.btm_names[0]]['out_shape'])])

    def ff(self, act, phase):
        shp = act.shape
        if len(shp) > 2:
            a = act.reshape([np.prod(shp[0:-1], dtype=np.int32), shp[-1]])
        else:
            a = act
        self.ff_act = act # save ff value
        if self.weight == None:
            self.init_weights_with_filler()
        return self.weight * a + self.bias

    def bp(self, sen, phase):
        shp = self.ff_act.shape
        if len(shp) > 2:
            a = self.ff_act.reshape([np.prod(shp[0:-1], dtype=np.int32), shp[-1]])
        else:
            a = self.ff_act
        self.weightgrad = sen * a.trans()
        self.biasgrad = sen.sum(1)
        s = self.weight.trans() * sen
        if len(shp) > 2:
            s = s.reshape(shp)
        return s
    def __str__(self):
        return 'fc'

class ConvConnection(WeightedComputeUnit):
    ''' Convolution operation

    .. note::
        The input and output is of size ``[HWCN]``:

        - ``H``: image height
        - ``W``: image width
        - ``C``: number of image channels (feature maps)
        - ``N``: size of minibatch

    '''
    def __init__(self, params):
        super(ConvConnection, self).__init__(params)
        self.conv_params = params.convolution_param
        self.convolver = co.Convolver(self.conv_params.pad,
                self.conv_params.pad, self.conv_params.stride, self.conv_params.stride)
        self.num_output = params.convolution_param.num_output
        self.group = params.convolution_param.group

        #TODO: hack, we don't want to slice agian to use it into bp as a parameter
        self.group_data = []
        self.group_filter = []
        self.group_bias = []
        self.weight_filler = params.convolution_param.weight_filler
        self.bias_filler = params.convolution_param.bias_filler

    def compute_size(self, from_btm, to_top):
        ''' Compute the output size and also weight and bias size

        .. note::
            The weight(kernel) size is ``[HWCiCo]``; bias shape is ``[Co]``:

            - ``H``: kernel_height
            - ``W``: kernel_width
            - ``Ci``: number of input channels
            - ``Co``: number of output channels
        '''
        self.in_shape = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = from_btm[self.btm_names[0]]['out_shape'][:]
        to_top[self.top_names[0]]['out_shape'][0] = (to_top[self.top_names[0]]['out_shape'][0] + 2 * self.conv_params.pad - self.conv_params.kernel_size) / self.conv_params.stride + 1
        to_top[self.top_names[0]]['out_shape'][1] = (to_top[self.top_names[0]]['out_shape'][1] + 2 * self.conv_params.pad - self.conv_params.kernel_size) / self.conv_params.stride + 1
        to_top[self.top_names[0]]['out_shape'][2] = self.num_output
        self.out_shape = to_top[self.top_names[0]]['out_shape'][:]
        self.wshape = [self.conv_params.kernel_size,
                       self.conv_params.kernel_size,
                       self.in_shape[2],
                       self.num_output]
        self.bshape = [self.out_shape[2]]

        to_top[self.top_names[0]]['stride_on_ori'] = from_btm[self.btm_names[0]]['stride_on_ori'] * self.conv_params.stride
        to_top[self.top_names[0]]['rec_on_ori'] = from_btm[self.btm_names[0]]['rec_on_ori'] + (self.conv_params.kernel_size - 1) * from_btm[self.btm_names[0]]['stride_on_ori']
        to_top[self.top_names[0]]['start_on_ori'] = from_btm[self.btm_names[0]]['start_on_ori'] - self.conv_params.pad * from_btm[self.btm_names[0]]['stride_on_ori']
        self.stride_on_ori = to_top[self.top_names[0]]['stride_on_ori']
        self.start_on_ori = to_top[self.top_names[0]]['start_on_ori']
        self.rec_on_ori = to_top[self.top_names[0]]['rec_on_ori']
        #set fan_in fan_out
        self.fan_out = self.conv_params.kernel_size * self.conv_params.kernel_size * self.conv_params.num_output
        self.fan_in = self.conv_params.kernel_size * self.conv_params.kernel_size * from_btm[self.btm_names[0]]['out_shape'][2]

    def ff(self, act, phase):
        ''' Feed-forward of convolution

        .. warning::
            Currently multi-group convolution (as in AlexNet paper) is not supported. One could walk around it by
            using a bigger convolution with number of feature maps doubled.
        '''
        if self.group == 1:
            self.ff_act = act
            if self.weight == None:
                self.init_weights_with_filler()
            return self.convolver.ff(act, self.weight, self.bias)
        else:
            #currently doesn't support multi-group
            assert(False)

    def bp(self, sen, phase):
        ''' Backward propagation of convolution

        .. warning::
            Currently multi-group convolution (as in AlexNet paper) is not supported. One could walk around it by
            using a bigger convolution with number of feature maps doubled.
        '''
        if self.group == 1:
            self.weightgrad = self.convolver.weight_grad(sen, self.ff_act, self.weight)
            self.biasgrad = self.convolver.bias_grad(sen)
            return self.convolver.bp(sen, self.ff_act, self.weight)
        else:
            #currently doesn't support multi-group
            assert(False)

    def __str__(self):
        return 'conv'

class DataUnit(ComputeUnit):
    ''' The base class of dataunit.

    :ivar dp: dataprovider, different kind of dp load data from different formats
    :ivar generator: the iterator produced by dataprovider
    '''

    def __init__(self, params, num_gpu):
        super(DataUnit, self).__init__(params)

    def compute_size(self, from_btm, to_top):
        pass

    def forward(self, from_btm, to_top, phase):
        ''' Feed-forward of data unit will get a batch of a fixed batch_size from data provider.

        .. note::

            Phase indicates whether it's training or testing. Usualy, the data augmentation operation for training involves some randomness, while testing doesn't

        '''

        if self.generator == None:
            self.generator = self.dp.get_mb(phase)

        while True:
            try:
                (samples, labels) = next(self.generator)
                if len(labels) == 0:
                    (samples, labels) = next(self.generator)
            except StopIteration:
                print 'Have scanned the whole dataset; start from the begginning agin'
                self.generator = self.dp.get_mb(phase)
                continue
            break

        to_top[self.top_names[0]] = owl.from_numpy(samples).reshape(
                [self.crop_size, self.crop_size, 3, samples.shape[0]])
        #may have multiplier labels
        for i in range (1, len(self.top_names)):
            to_top[self.top_names[i]] = labels[:,i - 1]

        #the output of datalayer is the data not label
        self.out = to_top[self.top_names[0]]

    def backward(self, from_top, to_btm, phase):
        # no bp pass
        pass
    def __str__(self):
        return 'data'

class LMDBDataUnit(DataUnit):
    ''' DataUnit load from LMDB.

    :ivar caffe.LayerParameter params: lmdb data layer param defined by Caffe, params.data_param contains information about data source, parmas.transform_param mainly defines data augmentation operations

    '''

    def __init__(self, params, num_gpu):
        super(LMDBDataUnit, self).__init__(params, num_gpu)
        if params.include[0].phase == Phase.Value('TRAIN'):
            self.dp = LMDBDataProvider(params.data_param, params.transform_param, num_gpu)
        else:
            self.dp = LMDBDataProvider(params.data_param, params.transform_param, 1)
        self.params = params
        self.crop_size = params.transform_param.crop_size
        self.generator = None
        self.out = None
        self.multiview = False

    def compute_size(self, from_btm, to_top):
        self.out_shape = [self.params.transform_param.crop_size,
                          self.params.transform_param.crop_size,
                          3, 1]
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = self.out_shape[:]
        to_top[self.top_names[0]]['rec_on_ori'] = 1
        to_top[self.top_names[0]]['stride_on_ori'] = 1
        to_top[self.top_names[0]]['start_on_ori'] = 0
        self.rec_on_ori = 1
        self.stride_on_ori = 1
        self.start_on_ori = 0

    def forward(self, from_btm, to_top, phase):
        ''' Feed-forward operation may vary according to phase.

        .. note::

            LMDB data provider now support multi-view testing, if multiview == True, it will produce concequtive 10 batches of different views of the same original image
        '''
        if self.generator == None:
            if self.multiview == False:
                self.generator = self.dp.get_mb(phase)
            #multiview test
            else:
                self.generator = self.dp.get_multiview_mb()
        while True:
            try:
                (samples, labels) = next(self.generator)
                if len(labels) == 0:
                    (samples, labels) = next(self.generator)
            except StopIteration:
                print 'Have scanned the whole dataset; start from the begginning agin'
                if self.multiview == False:
                    self.generator = self.dp.get_mb(phase)
                #multiview test
                else:
                    self.generator = self.dp.get_multiview_mb()
                continue
            break
        to_top[self.top_names[0]] = owl.from_numpy(samples).reshape(
                [self.crop_size, self.crop_size, 3, samples.shape[0]])
        for i in range (1, len(self.top_names)):
            to_top[self.top_names[i]] = labels[:,i - 1]
        #to_top[self.top_names[0]] = owl.zeros([self.crop_size, self.crop_size, 3, 256])
        #for i in range (1, len(self.top_names)):
            #to_top[self.top_names[i]] = np.ones(256)
        self.out = to_top[self.top_names[0]]

    def __str__(self):
        return 'lmdb_data'

class ImageDataUnit(DataUnit):
    ''' DataUnit load from raw images.
    :ivar caffe.LayerParameter params: image data layer param defined by Caffe, this is often used when data is limited. Loading from original image will be slower than loading from LMDB
    '''

    def __init__(self, params, num_gpu):
        super(ImageDataUnit, self).__init__(params, num_gpu)
        if params.include[0].phase == Phase.Value('TRAIN'):
            self.dp = ImageListDataProvider(params.image_data_param, params.transform_param, num_gpu)
        else:
            self.dp = ImageListDataProvider(params.image_data_param, params.transform_param, 1)
        self.params = params
        self.crop_size = params.transform_param.crop_size
        self.generator = None

    def compute_size(self, from_btm, to_top):
        self.out_shape = [self.params.transform_param.crop_size,
                          self.params.transform_param.crop_size,
                          3, 1]
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = self.out_shape[:]
        to_top[self.top_names[0]]['rec_on_ori'] = 1
        to_top[self.top_names[0]]['stride_on_ori'] = 1
        to_top[self.top_names[0]]['start_on_ori'] = 0
        self.rec_on_ori = 1
        self.stride_on_ori = 1
        self.start_on_ori = 0

    def __str__(self):
        return 'image_data'

class ImageWindowDataUnit(DataUnit):
    ''' DataUnit load from image window patches.
    :ivar caffe.LayerParameter params: image window data layer param defined by Caffe, this is often used when data is limited and object bounding box is given

    '''

    def __init__(self, params, num_gpu):
        super(ImageWindowDataUnit, self).__init__(params, num_gpu)
        if params.include[0].phase == Phase.Value('TRAIN'):
            self.dp = ImageWindowDataProvider(params.window_data_param, num_gpu)
        else:
            self.dp = ImageWindowDataProvider(params.window_data_param, 1)
        self.params = params
        self.crop_size = params.window_data_param.crop_size
        self.generator = None

    #reset generator
    def reset_generator(self):
        if self.params.include[0].phase == Phase.Value('TRAIN'):
            self.generator = self.dp.get_mb('TRAIN')
        else:
            self.generator = self.dp.get_mb('TEST')

    def compute_size(self, from_btm, to_top):
        self.out_shape = [self.params.window_data_param.crop_size,
                          self.params.window_data_param.crop_size,
                          3, 1]
        to_top[self.top_names[0]] = dict()
        to_top[self.top_names[0]]['out_shape'] = self.out_shape[:]
        to_top[self.top_names[0]]['rec_on_ori'] = 1
        to_top[self.top_names[0]]['stride_on_ori'] = 1
        to_top[self.top_names[0]]['start_on_ori'] = 0
        self.rec_on_ori = 1
        self.stride_on_ori = 1
        self.start_on_ori = 0

    def __str__(self):
        return 'window_data'

class Net:
    ''' The class for neural network structure

    The Net is basically a graph (DAG), of which each node is a :py:class:`ComputeUnit`.

    :ivar units: all the ``ComputeUnit`` s.
    :vartype units: list owl.net.ComputeUnit
    :ivar adjacent: the adjacent list (units are represented by their name)
    :vartype adjacent: list list str
    :ivar reverse_adjacent: the reverse adjacent list (units are represented by their name)
    :vartype reverse_adjacent: list list str
    :ivar dict name_to_uid: a map from units' name to the unit object
    :ivar loss_uids: all the units for computing loss
    :vartype loss_uids: list int
    :ivar accuracy_uids: all the units for calculating accuracy
    :vartype accuracy_uids: list int
    '''
    def __init__(self):
        self.units = []
        self.adjacent = []
        self.reverse_adjacent = []
        self.base_lr = 0
        self.base_weight_decay = 0
        self.momentum = 0
        self.name_to_uid = {}
        self.loss_uids = []
        self.accuracy_uids = []

    def add_unit(self, unit):
        ''' Method for adding units into the graph

        :param owl.net.ComputeUnit unit: the unit to add
        '''
        uid = len(self.units)
        self.units.append(unit)
        self.adjacent.append([])
        self.reverse_adjacent.append([])
        if not unit.name in self.name_to_uid:
            self.name_to_uid[unit.name] = []
        self.name_to_uid[unit.name].append(uid)
        return uid

    def connect(self, u1, u2):
        ''' Method for connecting two units

        :param str u1: name of the bottom unit
        :param str u2: name of the top unit
        '''
        self.adjacent[u1].append(u2)
        self.reverse_adjacent[u2].append(u1)

    def get_units_by_name(self, name):
        ''' Get ``ComputeUnit`` object by its name

        :param str name: unit name
        :return: the compute unit object of that name
        :rtype: owl.net.ComputeUnit
        '''
        return [self.units[uid] for uid in self.name_to_uid[name]]

    def get_loss_units(self):
        ''' Get all ``ComputeUnit`` object for loss

        :return: all compute unit object for computing loss
        :rtype: list owl.net.ComputeUnit
        '''
        return [self.units[uid] for uid in self.loss_uids]

    def get_accuracy_units(self):
        ''' Get all ``ComputeUnit`` object for accuracy

        :return: all compute unit object for computing accuracy
        :rtype: list owl.net.ComputeUnit
        '''
        return [self.units[uid] for uid in self.accuracy_uids]

    def get_data_unit(self, phase = 'TRAIN'):
        ''' Get the ``ComputeUnit`` object for data loading

        :param str phase: phase name of the run
        :return: the compute unit object for loading data
        :rtype: owl.net.ComputeUnit
        '''
        data_units = self.name_to_uid['data']
        for du in data_units:
            if not self._is_excluded(du, phase):
                return self.units[du]

    def get_weighted_unit_ids(self):
        ''' Get ids for all :py:class:owl.net.WeightedComputeUnit

        :return: ids of all weighted compute unit
        :rtype: list int
        '''
        weights_id = []
        for i in xrange(len(self.units)):
            if isinstance(self.units[i], WeightedComputeUnit):
                weights_id.append(i)
        return weights_id

    def _is_excluded(self, unit, phase):
        p = self.units[unit].params
        return phase != None and len(p.include) != 0 and p.include[0].phase != Phase.Value(phase)

    def _toporder(self, phase = None):
        depcount = [len(inunits) for inunits in self.reverse_adjacent]
        queue = Queue.Queue()
        # remove dep from excluded units
        for unit in range(len(depcount)):
            if self._is_excluded(unit, phase):
                for l in self.adjacent[unit]:
                    depcount[l] -= 1
        # find start units
        for unit in range(len(depcount)):
            count = depcount[unit]
            if count == 0:
                queue.put(unit)
        # run
        while not queue.empty():
            unit = queue.get()
            if self._is_excluded(unit, phase):
                continue
            yield unit
            for l in self.adjacent[unit]:
                depcount[l] -= 1
                if depcount[l] == 0:
                    queue.put(l)

    def _reverse_toporder(self, phase = None):
        depcount = [len(outunits) for outunits in self.adjacent]
        queue = Queue.Queue()
        # remove dep from excluded units
        for unit in range(len(depcount)):
            if self._is_excluded(unit, phase):
                for l in self.reverse_adjacent[unit]:
                    depcount[l] -= 1
        # find start units
        for unit in range(len(depcount)):
            count = depcount[unit]
            if count == 0:
                queue.put(unit)
        # run
        while not queue.empty():
            unit = queue.get()
            if self._is_excluded(unit, phase):
                continue
            yield unit
            for l in self.reverse_adjacent[unit]:
                depcount[l] -= 1
                if depcount[l] == 0:
                    queue.put(l)

    def compute_size(self, phase = 'TRAIN'):
        ''' Perform the compute_size phase before running
        '''
        unit_to_tops = [{} for name in self.units]
        for u in self._toporder(phase):
            from_btm = {}
            for btm in self.reverse_adjacent[u]:
                from_btm.update(unit_to_tops[btm])
            self.units[u].compute_size(from_btm, unit_to_tops[u])
        '''
        for u in self._toporder(phase):
            print self.units[u].name
            print self.units[u].out_shape
            print self.units[u].rec_on_ori
            print self.units[u].stride_on_ori
            print self.units[u].start_on_ori
        exit(0)
        '''

    def forward(self, phase = 'TRAIN'):
        ''' Perform the forward pass
        '''
        unit_to_tops = [{} for name in self.units]
        for u in self._toporder(phase):
            from_btm = {}
            for btm in self.reverse_adjacent[u]:
                from_btm.update(unit_to_tops[btm])
            self.units[u].forward(from_btm, unit_to_tops[u], phase)

    def forward_check(self):
        ''' Check forward function, use the same batch of data, remove random
        '''
        unit_to_tops = [{} for name in self.units]
        for u in self._toporder('TEST'):
            from_btm = {}
            for btm in self.reverse_adjacent[u]:
                from_btm.update(unit_to_tops[btm])
            self.units[u].forward(from_btm, unit_to_tops[u], 'CHECK')

    def backward(self, phase = 'TRAIN'):
        ''' Perform the backward pass
        '''
        unit_to_btms = [{} for name in self.units]
        for u in self._reverse_toporder(phase):
            from_top = {}
            for top in self.adjacent[u]:
                for keys in unit_to_btms[top]:
                    if keys in from_top:
                        from_top[keys] += unit_to_btms[top][keys]
                    else:
                        from_top[keys] = unit_to_btms[top][keys]
            self.units[u].backward(from_top, unit_to_btms[u], phase)

    def update(self, uid):
        ''' Update weights of one compute unit of the given uid

        :param int uid: id of the compute unit to update
        '''
        self.units[uid].weight_update(self.current_lr,
                                      self.base_weight_decay,
                                      self.momentum,
                                      self.batch_size)

    def weight_update(self):
        ''' Update weights for all units
        '''
        for i in range(len(self.units)):
            self.update(i)

    def __str__(self):
        ret = 'digraph G {\n'
        for uid in range(len(self.units)):
            ret += 'n' + str(uid) + ' [label="' + self.units[uid].name + '"]\n'
        for uid in range(len(self.units)):
            for nuid in self.adjacent[uid]:
                ret += 'n' + str(uid) + ' -> n' + str(nuid) + '\n'
        return ret + '}\n'


#!/usr/bin/env python
from astropy.io import ascii
import numpy as np
import matplotlib.pyplot as plt
import util
import os
import EDGE as edge
import starparam

#Set up path to where EDGE/MODULES resides
modulespath = os.path.dirname(os.path.realpath(util.__file__)) + '/'
#Set up path to where EDGE/COMMON resides
commonpath = os.path.realpath(modulespath + '../COMMON') + '/'
#Set up path to where you'd like to put your output files
outpath = './'

# interactive Av determination? if False, it just uses input Av
inter = True

# object name for labeling output file
obj='gmaur'

#spectral type; should be a whole number
sptin='K5'

# uncomment extinction law to use, default is mcclure
#law='mathis'
#law='HD29647'
#law='CCM89'
law='mcclure'

# uncomment which color table to use
table='kh95'
#table='pm13'

# uncomment which isochrones to use
#isochrone = 'baraffe'
isochrone = 'siess'

# uncomment to turn on/off HR diagram display
HR = True

# uncomment Rv to use
r=3.1
#r=5.0

# Av
avin=0.8

# distance in pc
distance=160

# If you have an obs fits file created with EDGE, you can
# input the photometry using it. You might need to change
# some code below to select the correct photometry with
# the names you gave them
input_obs = True
if input_obs:
    c = 2.99793e10
    source_obs = edge.loadObs(obj,outpath)
    xu = util.convertJy_to_Mag(source_obs.photometry['UBVRI']['lFl'][0]*
    source_obs.photometry['UBVRI']['wl'][0]*1e-4/c*1e23,'U')
    xb = util.convertJy_to_Mag(source_obs.photometry['UBVRI']['lFl'][1]*
    source_obs.photometry['UBVRI']['wl'][1]*1e-4/c*1e23,'B')
    xv = util.convertJy_to_Mag(source_obs.photometry['UBVRI']['lFl'][2]*
    source_obs.photometry['UBVRI']['wl'][2]*1e-4/c*1e23,'V')
    xr = util.convertJy_to_Mag(source_obs.photometry['UBVRI']['lFl'][3]*
    source_obs.photometry['UBVRI']['wl'][3]*1e-4/c*1e23,'R')
    xi = util.convertJy_to_Mag(source_obs.photometry['UBVRI']['lFl'][4]*
    source_obs.photometry['UBVRI']['wl'][4]*1e-4/c*1e23,'I')
    xj = util.convertJy_to_Mag(source_obs.photometry['2MASS']['lFl'][0]*
    source_obs.photometry['2MASS']['wl'][0]*1e-4/c*1e23,'J')
    xh = util.convertJy_to_Mag(source_obs.photometry['2MASS']['lFl'][1]*
    source_obs.photometry['2MASS']['wl'][1]*1e-4/c*1e23,'H')
    xk = util.convertJy_to_Mag(source_obs.photometry['2MASS']['lFl'][2]*
    source_obs.photometry['2MASS']['wl'][2]*1e-4/c*1e23,'K')
    xl = np.nan
    xm = np.nan
    irac1 = util.convertJy_to_Mag(source_obs.photometry['IRAC']['lFl'][0]*
    source_obs.photometry['IRAC']['wl'][0]*1e-4/c*1e23,'IRAC3.6')
    irac2 = util.convertJy_to_Mag(source_obs.photometry['IRAC']['lFl'][1]*
    source_obs.photometry['IRAC']['wl'][1]*1e-4/c*1e23,'IRAC4.5')
    irac3 = util.convertJy_to_Mag(source_obs.photometry['IRAC']['lFl'][2]*
    source_obs.photometry['IRAC']['wl'][2]*1e-4/c*1e23,'IRAC5.8')
    irac4 = util.convertJy_to_Mag(source_obs.photometry['IRAC']['lFl'][3]*
    source_obs.photometry['IRAC']['wl'][3]*1e-4/c*1e23,'IRAC8.0')
    mips1 = util.convertJy_to_Mag(source_obs.photometry['MIPS']['lFl'][0]*
    source_obs.photometry['MIPS']['wl'][0]*1e-4/c*1e23,'MIPS24')

else: #If not, input it by hand
    # input photometry in magnitudes, when missing a value np.nan
    # need to list at least J-band ('xj') for SpT less than ~G
    # for earlier type stars need to list at least V-band ('xv')
    # to get Mdot need U-band ('xu') in addition to above
    # put 99.00 for magnitudes with no value
    xu=13.6
    xb=12.5
    xv=11.4
    xr=10.4
    xi=9.74
    xj=8.90
    xh=8.47
    xk=8.58
    xl=np.nan
    xm=np.nan
    irac1=8.04
    irac2=7.88
    irac3=7.68
    irac4=7.07
    mips1=2.48

# what to call output file
outputfile=outpath+'starparam.'+obj+'.'+law+'.rv'+str(r)+'.av'+str(avin)+'.'+table

# FOR TEMPLATE PHOTOSPHERE

# do you want an output template photosphere (useful when fitting for Av)
calcphot = True

# what to call output photosphere file
photfile=outpath+'photosphere.'+obj+'.'+law+'.rv'+str(r)+'.av'+str(avin)+'.'+table

# wavelengths over which to interpolate template photosphere
# use 'wlfile_standard.ent' if you're using the photosphere in SED models
photfilewl=commonpath+'wavelengths/'+'wlfile_standard.ent'

#
starparam.starparam(obj, sptin, avin, distance, law, table, isochrone, HR,
calcphot, inter, r,
xu, xb, xv, xr, xi, xj, xh, xk, xl, xm, irac1, irac2, irac3, irac4, mips1,
outpath, outputfile, photfile, commonpath, photfilewl)


"""
Author: Hans Pinckaers
MIT License
"""
import copy
import math
import os
from dataclasses import dataclass
from itertools import repeat
from typing import NamedTuple, Union, List

import numpy as np
import torch
import torch.autograd
import torch.backends
import torch.nn.functional

from torch._six import container_abcs
from torch.nn.modules.conv import _ConvNd
from torch.nn.modules.utils import _pair
from torch.utils.cpp_extension import load

from tqdm import tqdm

# from torch.nn.grad import _grad_input_padding

if '1.6' in torch.__version__: # type:ignore
    def forward_amp_decorator(func):
        return torch.cuda.amp.custom_fwd(func)  # type:ignore
    def backward_amp_decorator(func):
        return torch.cuda.amp.custom_bwd(func)  # type:ignore
    from torch.cuda.amp import autocast
else:
    def forward_amp_decorator(func):
        return func
    def backward_amp_decorator(func):
        return func

# Load and compile cpp code to call cudnn conv2d backward function
dirname = os.path.dirname(__file__)
filename = os.path.join(dirname, "cpp_functions.cpp")
cpp_functions = load(name="cpp_functions", sources=[filename], verbose=False)

# inspired by torch/nn/modules/utils.py
def _ntuple(n):
    def parse(x, default=0):
        if isinstance(x, container_abcs.Iterable):
            if len(x) == n:
                return x
            elif len(x) == n-1:
                return tuple([default, *x])
            else:
                return tuple(repeat(x[0], n))
        return tuple(repeat(x, n))
    return parse

_triple = _ntuple(3)

# Utility named tuples, makes code more readable
class Sides(NamedTuple):
    left: int
    top: int
    right: int
    bottom: int

@dataclass
class Box:
    y: int
    height: int
    x: int
    width: int
    sides: Union[Sides, None]

class IOShape(NamedTuple):
    batch: int
    channels: int
    height: int
    width: int

@dataclass
class Lost:
    top: int
    left: int
    bottom: int
    right: int

    def __str__(self):
        return 'Lost(top:%2.1f, left:%2.1f, bottom:%2.1f, right:%2.1f)' \
            % (self.top, self.left, self.bottom, self.right)

class StreamingConv2dF(torch.autograd.Function):
    @staticmethod
    @forward_amp_decorator
    def forward(ctx, inpt, weight, bias, stride, padding, dilation, groups, grad_lost, seen_indices, output_stride, input_loc):
        ctx.save_for_backward(inpt, weight, bias)
        ctx.stride = stride
        ctx.padding = padding
        ctx.dilation = dilation
        ctx.groups = groups
        ctx.grad_lost = grad_lost
        ctx.seen_indices = seen_indices
        ctx.output_stride = output_stride
        ctx.input_loc = input_loc
        return torch.nn.functional.conv2d(inpt, weight, bias, stride, padding, dilation, groups)

    @staticmethod
    @backward_amp_decorator
    def backward(ctx, grad_output):
        inpt, weight, bias = ctx.saved_variables
        grad = grad_weight = grad_bias = None

        stride = ctx.stride
        padding = ctx.padding
        dilation = ctx.dilation
        groups = ctx.groups
        sides = ctx.input_loc.sides  # Type: Sides
        seen_indices = ctx.seen_indices
        grad_lost = ctx.grad_lost  # Type: Lost
        output_stride = ctx.output_stride
        grad_bias = None
        kernel_size = weight.shape[-1]

        if ctx.needs_input_grad[0]:
            # TODO: performance improvements possible by only backpropping valid input
            # grad_input_padding = _grad_input_padding(grad_output, inpt.shape, stride, padding, (weight.shape[2], weight.shape[3]))
            # TODO: use this!?
            grad_in = cpp_functions.backward_input(inpt.shape, grad_output, weight.to(inpt.dtype), padding,
                                                   stride, dilation, groups,
                                                   torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic)
        else:
            grad_in = None

        grad = grad_output

        lost_top = grad_lost.top if not sides.top else 0
        lost_bottom = grad_lost.bottom if not sides.bottom else 0
        lost_left = grad_lost.left if not sides.left else 0
        lost_right = grad_lost.right if not sides.right else 0

        valid_grad = grad[:, :, lost_top:grad.shape[H_DIM] - lost_bottom,
                          lost_left:grad.shape[W_DIM] - lost_right]

        stride, kernel_size, padding = _triple(stride), _triple(kernel_size), _triple(padding)

        output_stride = output_stride * torch.tensor(stride)
        input_loc = ctx.input_loc

        # Move the location according to how many pixels have been trimmed
        # this will be the location of the valid gradient of this layer in relation
        # to the actual gradient in a normal backpass
        data_loc_y = int(input_loc.y // output_stride[1]) + lost_top
        data_loc_x = int(input_loc.x // output_stride[2]) + lost_left

        data_loc = Box(data_loc_y, 0,
                       data_loc_x, 0,
                       input_loc.sides)

        # Calculate which part of the gradient is 'new'
        old_value_indices = seen_indices
        new_output_box, updated_total_indices = StreamingCNN._new_value_indices(valid_grad.shape,
                                                                                data_loc,
                                                                                old_value_indices)

        # Update inplace
        seen_indices.y = updated_total_indices.y
        seen_indices.height = updated_total_indices.height
        seen_indices.x = updated_total_indices.x
        seen_indices.width = updated_total_indices.width
        seen_indices.sides = updated_total_indices.sides

        if new_output_box.height > 0 and new_output_box.width > 0:
            relevant_grad = valid_grad[:, :,
                                       new_output_box.y:new_output_box.y + new_output_box.height,
                                       new_output_box.x:new_output_box.x + new_output_box.width]

            input_y = (new_output_box.y + lost_top) * stride[1]
            input_x = (new_output_box.x + lost_left) * stride[2]

            # Accounting for padding:
            # the kernel locations are relative to the padded input, inpt[0] is not padded
            # this means that the corresponding input of the grad_loc is module.padding shifted to the left
            # we account for this:
            input_y -= padding[1]
            input_x -= padding[2]
            input_x = max(0, input_x)
            input_y = max(0, input_y)

            relevant_input_height = relevant_grad.shape[H_DIM] * stride[1] + (kernel_size[1] - 1)
            relevant_input_width = relevant_grad.shape[W_DIM] * stride[2] + (kernel_size[2] - 1)

            relevant_input = inpt[:, :,
                                  input_y:input_y + relevant_input_height,
                                  input_x:input_x + relevant_input_width]

            # If layer has padding we need to pad based on if the current tile
            # is at the sides of the input.
            if (padding[0] > 0 or padding[1] > 0 or padding[2] > 0) and \
                    (sides.top or sides.left or sides.right or sides.bottom):
                # The size of the tile should remain equal.
                crop_bottom = padding[1] if sides.top else 0
                crop_right = padding[2] if sides.left else 0
                relevant_input = inpt[:, :,
                                      input_y:input_y + relevant_input_height - crop_bottom,
                                      input_x:input_x + relevant_input_width - crop_right]

                relevant_input = torch.nn.functional.pad(relevant_input, [padding[2] if sides.left else 0,
                                                                          padding[2] if sides.right else 0,
                                                                          padding[1] if sides.top else 0,
                                                                          padding[1] if sides.bottom else 0])

            # Calculate the kernel gradients with the new unseen gradient values
            relevant_grad = relevant_grad.contiguous()

            grad_weight = cpp_functions.backward(weight.shape,
                                                 relevant_grad.to(weight.dtype),
                                                 relevant_input.to(weight.dtype),
                                                 (0, 0),  # padding
                                                 stride[1:3], dilation, groups,
                                                 torch.backends.cudnn.benchmark,  # benchmark
                                                 torch.backends.cudnn.deterministic)  # deterministic

            if bias is not None:
                grad_bias = relevant_grad[0].sum((1, 2))

            del relevant_input
            del relevant_grad
        else:
            # if self.verbose and not hasattr(self, '_inefficient_tile_shape_warning'):
            # print("Warning: no new gradient values found. Tile size could be too small.")
            # self._inefficient_tile_shape_warning = True
            grad_weight = torch.zeros_like(weight)
            if bias is None: grad_bias = None
            else: grad_bias = torch.zeros_like(bias)

        if bias is not None:
            return grad_in, grad_weight, grad_bias, None, None, None, None, None, None, None, None,
        else:
            return grad_in, grad_weight, None, None, None, None, None, None, None, None, None,

conv2d = StreamingConv2dF.apply  # type:ignore

class StreamingConv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros'):
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(StreamingConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False, _pair(0), groups, bias, padding_mode)
        self.grad_lost = Lost(0, 0, 0, 0)
        self.tile_output_box = Box(0, 0, 0, 0, None)
        self.reset()

    def reset(self):
        self.seen_indices = Box(0, 0, 0, 0, None)
        self.input_loc = Box(0, 0, 0, 0, None)

    def forward(self, input):
        return conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups,
                      self.grad_lost, self.seen_indices, self.output_stride, self.input_loc)

B_DIM = 0
C_DIM = 1
H_DIM = 2
W_DIM = 3

class StreamingCNN(object):
    '''Initialize Streaming CNN helper class. After initialization use the
    forward() and backward() function of this class to stream.
    Pseudocode example:

    ```python
    sCNN = StreamingCNN(stream_layers, tile_shape=(1, 3, 600, 600))
    str_output = sCNN.forward(image)
    final_output = final_layers(str_output)
    loss = criterion(final_output, labels)
    loss.backward()
    sCNN.backward(image, str_output.grad)
    ```

    Hooks are used to perform streaming, to use the stream_layers without
    streaming you can disable StreamingCNN with the disable() function.
    Subsequently, enable() enables it again. Streaming gets enabled by default
    after initialization.
    '''
    def __init__(self, stream_module, tile_shape, verbose=False, deterministic=False,
                 saliency=False, gather_gradients=False, replace_non_linearity=True,
                 eps=1e-5, copy_to_gpu=True, dtype=None, statistics_on_cpu=False,
                 normalize_on_gpu=False, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
                 state_dict=None):
        '''
        Parameters:
            stream_module (torch.nn.Module): module containing the to be streamed layers
            tile_shape (tuple, NCHW): size of the to be streamed tiles
            verbose (bool): will log various debugging relevant information (default is False)
            deterministic (bool): whether to use the deterministic algorithms for cudnn
            saliency (bool): will gather the gradients of the input image (saliency map)
            gather_gradients (bool): will gather the gradients of the feature maps
            eps (float): epsilon error to compare floating values
        '''
        global H_DIM, W_DIM
        self.stream_module = stream_module
        self.verbose = verbose
        self.deterministic = deterministic
        self.eps = eps
        self.device = next(stream_module.parameters()).device
        self.dtype = next(stream_module.parameters()).dtype
        if dtype is not None: self.dtype = dtype
        self.tile_shape = tile_shape
        self.gather_input_gradient = saliency
        self.gather_gradient = gather_gradients
        self.replace_non_linearity = replace_non_linearity
        self.copy_to_gpu = copy_to_gpu
        self.statistics_on_cpu = statistics_on_cpu

        self.mean = torch.tensor(mean).cuda()[:, None, None]
        self.std = torch.tensor(std).cuda()[:, None, None]
        self.should_normalize = normalize_on_gpu

        self._tile_output_shape = None
        self._module_stats = {}
        self._backward_seen_indices = {}
        self._saved_tensors = {}
        self._current_tile_input_loc = None
        self._hooks = []

        if state_dict is None:
            self._configure()
        else:
            self.load_state_dict(state_dict)

    def _configure(self):
        if self.replace_non_linearity: self.convert_modules_model(self.stream_module)
        self.convert_modules_model(self.stream_module, from_mod=torch.nn.BatchNorm2d, to_mod=torch.nn.Sequential)

        # Save current model and cudnn flags, since we need to change them and restore later
        state_dict = self._save_parameters()
        old_deterministic_flag, old_benchmark_flag = self._set_cudnn_flags_to_determistic()
        self._reset_parameters_to_constant()

        # Add hooks to each layer to gather statistics
        self._add_hooks_for_statistics()

        # We need to temporary store statistics per layer to keep track of the
        # total output stride at each layer
        self._stats_per_grad_fn = {}

        # TODO; temp hack for tile sizes too big on gpu,
        # we need float32 precision
        if self.statistics_on_cpu:
            self.stream_module = self.stream_module.cpu()
            self.device = torch.device('cpu')  # type:ignore

        # Create all-ones tile
        tile = torch.ones(self.tile_shape, dtype=self.dtype, requires_grad=True, device=self.device)

        self._gather_forward_statistics(tile)
        if self.verbose: print('')
        self._gather_backward_statistics(tile)

        # TODO; temp hack for tile sizes too big on gpu,
        if self.statistics_on_cpu:
            self.stream_module = self.stream_module.cuda()
            self.device = torch.device('cuda')  # type:ignore

        # Remove all hooks and add hooks for correcting gradients
        # during streaming
        self._remove_hooks()
        self._add_hooks_for_streaming()
        self._restore_parameters(state_dict)
        self._convert_modules_for_streaming(self.stream_module)
        if self.replace_non_linearity: self.convert_modules_model(self.stream_module, back=True)

        # Remove temporary data
        self._saved_tensors = {}
        del self._stats_per_grad_fn

        # Zero the gradients
        for param in self.stream_module.parameters():
            if param.grad is not None: param.grad.data.zero_()

        self._set_cudnn_flags(old_deterministic_flag, old_benchmark_flag)
        del state_dict

    def _gather_backward_statistics(self, tile):
        # Forward pass with grads enabled
        torch.set_grad_enabled(True)
        output = self.stream_module(tile)

        # Gather backward statistics
        self._tile_output_shape = output.shape

        gradient = torch.zeros(*output.shape, dtype=self.dtype, device=self.device)
        gradient[:, :,
                 self.tile_output_lost.top:output.shape[H_DIM] - self.tile_output_lost.bottom,
                 self.tile_output_lost.left:output.shape[W_DIM] - self.tile_output_lost.right] = 1
        output.backward(gradient=gradient)

        # Calculate the output stride of the whole stream_module
        p_stats = self._prev_stats(output)
        if p_stats: self.output_stride = p_stats['output_stride'] * torch.tensor(p_stats['stride'])
        else: self.output_stride = torch.tensor([1, 1, 1])

        self.tile_gradient_lost = self._non_max_border_amount(tile.grad)

        if self.verbose:
            print('\n', 'Input gradient lost', self.tile_gradient_lost)

    def _gather_forward_statistics(self, tile):
        torch.set_grad_enabled(False)
        output = self.stream_module(tile)
        self.tile_output_lost = self._non_max_border_amount(output)
        if self.verbose: print('\n', 'Output lost', self.tile_output_lost)

    def convert_modules_model(self, module, from_mod=torch.nn.ReLU6, to_mod=torch.nn.ReLU, back=False):
        mod = module
        if not back and isinstance(module, from_mod):
            mod = to_mod()
            # mod.previous_mod = module
        if back and isinstance(module, to_mod):
            mod = module.previous_mod
        for name, child in module.named_children():
            mod.add_module(name, self.convert_modules_model(child, from_mod, to_mod))
        del module
        return mod

    def _convert_modules_for_streaming(self, module):
        mod = module
        if isinstance(module, torch.nn.Conv2d):
            if module in self._module_stats:
                mod = StreamingConv2d(module.in_channels, module.out_channels, module.kernel_size, module.stride, module.padding, module.dilation, module.groups, module.bias is not None)
                mod = mod.to(module.weight.device)
                mod = mod.to(module.weight.dtype)

                mod.weight.requires_grad = module.weight.requires_grad
                if module.bias is not None:
                    mod.bias.requires_grad = module.bias.requires_grad

                mod.load_state_dict(module.state_dict())  # copy params
                mod.grad_lost = self._module_stats[module]['grad_lost']
                mod.output_stride = self._module_stats[module]['output_stride']
                self._module_stats[mod] = self._module_stats[module]
                del self._module_stats[module]
        for name, child in module.named_children():
            mod.add_module(name, self._convert_modules_for_streaming(child))
        del module
        return mod

    def _reset_converted_modules(self, module):
        mod = module
        if isinstance(module, StreamingConv2d):
            mod = torch.nn.Conv2d(module.in_channels, module.out_channels, module.kernel_size, module.stride, module.padding, module.dilation, module.groups, module.bias is not None)
            mod = mod.to(module.weight.device)
            mod = mod.to(module.weight.dtype)

            mod.weight.requires_grad = module.weight.requires_grad
            if module.bias is not None:
                mod.bias.requires_grad = module.bias.requires_grad

            mod.load_state_dict(module.state_dict())  # copy params
            self._module_stats[mod] = self._module_stats[module]
            del self._module_stats[module]
        for name, child in module.named_children():
            mod.add_module(name, self._reset_converted_modules(child))
        del module
        return mod

    def _reset_parameters_to_constant(self):
        for mod in self.stream_module.modules():
            if isinstance(mod, (torch.nn.Conv2d)):
                # to counter loating precision errors, we assign 1 to the weights and
                # normalize the output after the conv.
                torch.nn.init.constant_(mod.weight, 1)
                if mod.bias is not None:
                    torch.nn.init.constant_(mod.bias, 0)

        for m in self.stream_module.modules():
            if isinstance(m, torch.nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
                m.eval()

    def _set_cudnn_flags(self, deterministic_flag, benchmark_flag):
        torch.backends.cudnn.deterministic = deterministic_flag
        torch.backends.cudnn.benchmark = benchmark_flag

    def _set_cudnn_flags_to_determistic(self):
        deterministic_flag = torch.backends.cudnn.deterministic
        benchmark_flag = torch.backends.cudnn.benchmark
        self._set_cudnn_flags(True, False)
        return deterministic_flag, benchmark_flag

    def _save_parameters(self):
        state_dict = self.stream_module.state_dict()
        state_dict = copy.deepcopy(state_dict)
        return state_dict

    def _restore_parameters(self, state_dict):
        self.stream_module.load_state_dict(state_dict)

    def _non_max_border_amount(self, tensor):
        # Sum over the channels, useful for networks that treat certain channels
        # different (e.g., DenseNet)
        if tensor.dim() > 3: tensor = torch.sum(tensor, dim=1)[0]
        tensor = tensor / tensor.max()  # normalize
        tensor = (tensor > tensor.max() * (1-self.eps))
        non_zero = tensor.nonzero()
        top, left = non_zero.min(dim=0)[0]
        # for bottom and right we need to substract -1: correct index 3 is actually the 4th pixel
        bottom, right = torch.tensor([*tensor.size()], dtype=torch.long, device=self.device) - non_zero.max(dim=0)[0] - 1
        return Lost(int(top), int(left), int(bottom), int(right))

    def forward(self, image, result_on_cpu=False):
        """Perform forward pass with streaming.

        Parameters:
            image (torch.Tensor): CHW the image to stream
        """
        # The input image is likely quite small in terms of channels, for
        # performance reasons it is beneficial to copy to the GPU as a whole
        # instead of tile-by-tile.
        image = image
        if self.copy_to_gpu:
            image = image.to(self.device, non_blocking=True)

        tile_width, tile_height = self.tile_shape[W_DIM], self.tile_shape[H_DIM]

        # Size of valid output of a tile
        valid_output_height = (self._tile_output_shape[H_DIM] - self.tile_output_lost.top - self.tile_output_lost.bottom)
        valid_output_width = (self._tile_output_shape[W_DIM] - self.tile_output_lost.left - self.tile_output_lost.right)

        # We will keep track which part of the output of the whole image we
        # already filled with valid values from tile output.
        already_filled = Box(0, 0, 0, 0, None)

        # Calculate size of output that we would get by inferencing the
        # whole image.
        output_height = (image.shape[H_DIM] - self.tile_shape[H_DIM]) // self.output_stride[1] + self._tile_output_shape[H_DIM]
        output_width = (image.shape[W_DIM] - self.tile_shape[W_DIM]) // self.output_stride[2] + self._tile_output_shape[W_DIM]

        if result_on_cpu:
            device = torch.device('cpu')
        else:
            device = self.device
        output = torch.empty((image.shape[0], self._tile_output_shape[1], output_height, output_width), dtype=self.dtype, device=device).fill_(999)

        n_rows = math.ceil(float(output_height) / float(valid_output_height))
        n_cols = math.ceil(float(output_width) / float(valid_output_width))

        if image.shape[W_DIM] <= tile_width: n_cols = 1
        if image.shape[H_DIM] <= tile_height: n_rows = 1

        if self.gather_input_gradient:
            self.saliency_map = torch.zeros(image.shape, dtype=self.dtype, device=self.device)

        if self.verbose: print('Number of tiles in forward:', n_rows * n_cols)
        if self.verbose: iterator = tqdm(range(n_rows))
        else: iterator = range(n_rows)

        with torch.no_grad():
            for row in iterator:
                for col in range(n_cols):
                    # Coordinates of the output w.r.t. the output of full image
                    output_y = row * valid_output_height
                    output_x = col * valid_output_width

                    # Check if we are at borders, since we can not create
                    # overlap here and should not crop values.
                    sides_top = True if row == 0 else False
                    sides_left = True if col == 0 else False

                    sides_bottom = True if output_y * self.output_stride[1] + self.tile_shape[H_DIM] >= image.shape[H_DIM] else False
                    sides_right = True if output_x * self.output_stride[2] + self.tile_shape[W_DIM] >= image.shape[W_DIM] else False
                    sides = Sides(sides_left, sides_top, sides_right, sides_bottom)

                    # These values are used to crop invalid output values
                    lost = self._get_tile_lost_for_sides(sides)

                    # Since we need to stay at multiples of output stride we
                    # need to keep that into account when we are at the bottom
                    # and right side of the output.
                    if sides_bottom: output_y = (image.shape[H_DIM] - self.tile_shape[H_DIM]) // self.output_stride[1]
                    if sides_right:  output_x = (image.shape[W_DIM] - self.tile_shape[W_DIM]) // self.output_stride[2]

                    output_y = output_y if not sides.top else 0
                    output_x = output_x if not sides.left else 0
                    output_loc = Box(output_y + lost.top, -1, output_x + lost.left, -1, sides)

                    # Coordinates of the input w.r.t. the output of full image
                    tile_y = output_y * self.output_stride[1]
                    tile_x = output_x * self.output_stride[2]

                    # Extract tile and perform forward pass
                    tile = image[:, :,
                                 tile_y:tile_y + tile_height,
                                 tile_x:tile_x + tile_width]

                    # normalize on gpu for speed in dataloader
                    # does this reduce speed significantly?
                    if not self.copy_to_gpu:
                        tile = tile.to(self.device, non_blocking=True)

                    if self.should_normalize: tile = self._normalize_on_gpu(tile)
                    tile_output = self.stream_module(tile)

                    trimmed_output = tile_output[:, :,
                                                 lost.top:tile_output.shape[H_DIM] - lost.bottom,
                                                 lost.left:tile_output.shape[W_DIM] - lost.right]

                    new_output_box, updated_total_indices = self._new_value_indices(trimmed_output.shape, output_loc, already_filled)
                    already_filled = updated_total_indices

                    relevant_output = trimmed_output[:, :,
                                                     new_output_box.y:updated_total_indices.y + new_output_box.height,
                                                     new_output_box.x:new_output_box.x + new_output_box.width]

                    output[:, :, int(updated_total_indices.y):int(updated_total_indices.height), int(updated_total_indices.x - new_output_box.width):int(updated_total_indices.x)] = relevant_output

                    del tile

            assert sides_bottom and sides_right, "It seems like we could not reconstruct all output"  #type:ignore

        # mem management
        del relevant_output  # type:ignore
        del image
        self._saved_tensors = {}

        return output

    def backward(self, image, grad):
        """Perform backward pass with streaming.

        Parameters:
            image (torch.Tensor): the image (expects NCHW) that was used in the forward pass
            grad (torch.Tensor): this should be the gradient of the output of
                the stream_layers.
        """
        # The input image is likely quite small in terms of channels, for
        # performance reasons it is beneficial to copy to the GPU as a whole
        # instead of tile-by-tile.
        image = image
        if self.copy_to_gpu:
            image = image.to(self.device, non_blocking=True)
        grad = grad

        height = image.shape[H_DIM]
        width = image.shape[W_DIM]

        tile_height = self.tile_shape[H_DIM]
        tile_width = self.tile_shape[W_DIM]
        grad_lost = self.tile_gradient_lost

        output_height = self._tile_output_shape[H_DIM]
        output_width = self._tile_output_shape[W_DIM]

        valid_grad_height = (tile_height - grad_lost.top - grad_lost.bottom) // self.output_stride[1]
        valid_grad_height *= self.output_stride[1]
        valid_grad_width = (tile_width - grad_lost.left - grad_lost.right) // self.output_stride[2]
        valid_grad_width *= self.output_stride[2]

        n_rows = math.ceil(float(height - grad_lost.top - grad_lost.bottom) / float(valid_grad_height))
        n_cols = math.ceil(float(width - grad_lost.left - grad_lost.right) / float(valid_grad_width))

        if self.verbose:
            ideal_tile_size = height / float(n_rows) + grad_lost.top + grad_lost.bottom
            next_ideal_tile_size = height / float(n_rows - 1) + grad_lost.top + grad_lost.bottom
            print(ideal_tile_size, n_rows*n_cols, next_ideal_tile_size)

        if image.shape[W_DIM] <= tile_width: n_cols = 1
        if image.shape[H_DIM] <= tile_height: n_rows = 1

        if self.gather_gradient:
            self.gradients = {}
        self._inputs = {}
        self._backward_seen_indices = {}

        if self.verbose: print('Number of tiles in backprop:', n_rows * n_cols)
        if self.verbose: iterator = tqdm(range(n_rows))
        else: iterator = range(n_rows)

        for row in iterator:
            for col in range(n_cols):
                # Since we determine output (gradient) coordinates based on input
                # coordinates. We need to divide by output stride.
                output_y = row * valid_grad_height // self.output_stride[1]
                output_x = col * valid_grad_width // self.output_stride[2]

                sides_top = True if row == 0 else False
                sides_left = True if col == 0 else False

                sides_bottom = True if output_y + output_height >= grad.shape[H_DIM] else False
                sides_right = True if output_x + output_width >= grad.shape[W_DIM] else False
                sides = Sides(sides_left, sides_top, sides_right, sides_bottom)

                # We are doing a forward pass
                lost = self._get_tile_lost_for_sides(sides)

                # If the tile is at the bottom or right side of the input image
                # than we need to shift back so that the tile fits (does not go
                # over the border)
                if sides_bottom: output_y = max(grad.shape[H_DIM] - output_height, 0)
                if sides_right: output_x = max(grad.shape[W_DIM] - output_width, 0)

                input_y = output_y * self.output_stride[1]
                input_x = output_x * self.output_stride[2]

                input_loc = Box(input_y, tile_height, input_x, tile_width, sides)

                tile = image[:, :,
                             input_y:input_y + tile_height,
                             input_x:input_x + tile_width]

                gradient = grad[:, :,
                                output_y:output_y + output_height,
                                output_x:output_x + output_width]

                self._saved_tensors = {}

                # Trim output and gradient
                trimmed_grad = gradient[:, :,
                                        lost.top:gradient.shape[H_DIM] - lost.bottom,
                                        lost.left:gradient.shape[W_DIM] - lost.right]

                if not self.copy_to_gpu:
                    tile = tile.to(self.device, non_blocking=True)

                for mod in self.stream_module.modules():
                    if isinstance(mod, StreamingConv2d):
                        mod.input_loc = input_loc

                # normalize on gpu for speed in dataloader
                # does this reduce speed significantly?
                if self.should_normalize: tile = self._normalize_on_gpu(tile)

                if self.dtype == torch.float16:
                    with autocast():
                        tile_output = self.stream_module(tile)
                else:
                    tile_output = self.stream_module(tile)

                del tile # memory management

                trimmed_output = tile_output[:, :,
                         lost.top:tile_output.shape[H_DIM] - lost.bottom,
                         lost.left:tile_output.shape[W_DIM] - lost.right]

                # Do backward pass, fix gradient in hooks
                trimmed_output = trimmed_output.to(self.device, non_blocking=True)

                # Sometimes when training with variable input shapes,
                # the gradient size is a bit too big
                if trimmed_grad.shape[H_DIM] != trimmed_output.shape[H_DIM] or \
                   trimmed_grad.shape[W_DIM] != trimmed_output.shape[W_DIM]:
                    assert image.shape[H_DIM] < self.tile_shape[H_DIM] or \
                        image.shape[W_DIM] < self.tile_shape[W_DIM]
                    trimmed_grad = trimmed_grad[:, :,
                                                0:trimmed_output.shape[H_DIM],
                                                0:trimmed_output.shape[W_DIM]]

                trimmed_output.backward(trimmed_grad)

                # Memory management
                del tile_output
                del trimmed_grad
                del trimmed_output

        # Memory management
        self._saved_tensors = {}
        self._current_tile_input_loc = None

        for mod in self.stream_module.modules():
            if isinstance(mod, StreamingConv2d):
                mod.input_loc = None
                mod.reset()

        assert sides_right and sides_bottom, "It seems like we could not reconstruct all output"  # type:ignore

    def _get_tile_lost_for_sides(self, sides):
        lost_top = self.tile_output_lost.top if not sides.top else 0
        lost_bottom = self.tile_output_lost.bottom if not sides.bottom else 0
        lost_left = self.tile_output_lost.left if not sides.left else 0
        lost_right = self.tile_output_lost.right if not sides.right else 0
        lost = Lost(lost_top, lost_left, lost_bottom, lost_right)
        return lost

    def _normalize_on_gpu(self, tile):
        tile_norm = tile.to(self.dtype)
        del tile
        tile_norm.div_(255)
        tile_norm.sub_(self.mean)
        tile_norm.div_(self.std)
        tile = tile_norm
        return tile

    def disable(self):
        """Disable the streaming hooks"""
        self._remove_hooks()
        self._reset_converted_modules(self.stream_module)

    def enable(self):
        """Enable the streaming hooks"""
        self._remove_hooks()
        self._add_hooks_for_streaming()
        self._convert_modules_for_streaming(self.stream_module)

    def _add_hooks_for_statistics(self):
        def forw_lambda(module, inpt, outpt):
            self._forward_gather_statistics_hook(module, inpt, outpt)

        def back_lambda(module, grad_in, grad_out):
            return self._backward_gather_statistics_hook(module, grad_in, grad_out)

        self._add_hooks(forward_hook=forw_lambda, backward_hook=back_lambda)

    def _add_hooks_for_streaming(self):
        if self.gather_input_gradient:
            def back_lambda(module, grad_in, grad_out):
                return self._backward_saliency_hook(module, grad_in, grad_out)

            for mod in self.stream_module.modules():
                if isinstance(mod, (torch.nn.Conv2d)):
                    if mod.in_channels == 3:
                        back_handle = mod.register_backward_hook(back_lambda)
                        self._hooks.append(back_handle)

    def _add_hooks(self, forward_hook, backward_hook,
                   forward_modules=(torch.nn.Conv2d, torch.nn.MaxPool2d, torch.nn.AvgPool2d),
                   back_modules=(torch.nn.Conv2d, torch.nn.MaxPool2d)):
        for mod in self.stream_module.modules():
            if isinstance(mod, forward_modules):
                forw_handle = mod.register_forward_hook(forward_hook)
                self._hooks.append(forw_handle)
                if back_modules and isinstance(mod, back_modules):
                    back_handle = mod.register_backward_hook(backward_hook)
                    self._hooks.append(back_handle)

    def _remove_hooks(self):
        for hook in self._hooks:
            hook.remove()

    def _forward_gather_statistics_hook(self, module, inpt, output):
        stride, kernel_size, _ = _triple(module.stride), _triple(module.kernel_size), _triple(module.padding)

        if not torch.is_grad_enabled():  # type:ignore
            # Convert strided convolutions/pooling to average pool
            if isinstance(module, (torch.nn.MaxPool2d)) or \
                    (stride[0] > 1 and stride[0] > kernel_size[0]) or \
                    (stride[1] > 1 and stride[1] > kernel_size[1]) or \
                    (stride[2] > 1 and stride[2] > kernel_size[2]):
                # Pytorch documentation is explicitely against changing output in a forward hook
                # However, since we do not really need the graph or gradients to be correct
                # it shouldn't harm.
                if module.padding != 0:
                    padding = module.padding
                    if not isinstance(module.padding, tuple):
                        padding = [module.padding, module.padding]
                    padded_input = torch.nn.functional.pad(inpt[0], [padding[1], padding[1], padding[0], padding[0]])
                else:
                    padded_input = inpt[0]

                new_output = torch.nn.functional.avg_pool2d(padded_input, kernel_size[1:], stride[1:])
                new_output = torch.sum(new_output, dim=1)[0]
                new_output = (new_output > (1-self.eps) * new_output.max())
                new_output = new_output.expand_as(output[0])
                output[0] = new_output.type(self.dtype)

            # Sum all dimensions (useful for DenseNet like networks)
            lost = self._non_max_border_amount(output)

            # Make output between 0-1 again, so the values do not explode
            output.fill_(0)
            output[:,:,lost.top:output[0, 0].shape[0] - lost.bottom,
                     lost.left:output[0, 0].shape[1] - lost.right] = 1

            module_stats = {'lost': lost, 'stride': stride, 'module': module}
            if self.verbose: print(module, "\n", module_stats['lost'])

            self._saved_tensors[module] = inpt
            self._module_stats[module] = module_stats
        else:
            module_stats = self._module_stats[module]

            p_stats = self._prev_stats(output)
            if p_stats: output_stride = p_stats['output_stride'] * torch.tensor(p_stats['stride'])
            else: output_stride = torch.tensor([1, 1, 1])

            module_stats['output_stride'] = output_stride.clone().detach()

            self._stats_per_grad_fn[output.grad_fn] = module_stats
            self._module_stats[module] = module_stats

    def _backward_gather_statistics_hook(self, module, grad_in, grad_out):
        stride, kernel_size, _ = _triple(module.stride), _triple(module.kernel_size), _triple(module.padding)

        if grad_in[0] is not None:
            # We sum over the channels to deal with networks that do different operations
            # on groups of channels
            f_grad = torch.sum(grad_in[0], dim=1)[0]

            if isinstance(module, (torch.nn.MaxPool2d)):
                # MaxPool shifts indices around, which break the calculation to
                # find valid gradient values. To fix this we do an average pool
                # with the same kernel-size and stride and repeat using the stride.
                inpt = self._saved_tensors[module]
                padded_inpt = inpt[0]
                if module.padding != 0:
                    padded_inpt = torch.nn.functional.pad(inpt[0], [module.padding, module.padding,
                                                                    module.padding, module.padding], value=-1)

                new_outpt = torch.nn.functional.avg_pool2d(padded_inpt, kernel_size[1:], stride[1:])[0]
                new_outpt = torch.sum(new_outpt, dim=0)

                f_grad = torch.sum(grad_out[0], dim=1)[0]
                f_grad = f_grad * new_outpt
                f_grad = f_grad.cpu()
                f_grad = np.repeat(f_grad, stride[0], axis=0)
                f_grad = np.repeat(f_grad, stride[1], axis=1)
                grad = np.zeros(grad_in[0].shape[2:])
                grad[:f_grad.shape[0], :f_grad.shape[1]] = f_grad
                f_grad = torch.from_numpy(grad)
                f_grad = f_grad.to(self.device)

            grad_lost = self._non_max_border_amount(grad_out[0])

            if self.verbose: print(module, "\n", grad_lost)
            self._module_stats[module]['grad_lost'] = grad_lost

            valid_grad = (f_grad > (1-self.eps) * f_grad.max())

            # When kernel_size > stride we have some _overlap_ of gradients,
            # this overlap makes extra positions in the input gradient invalid
            if (stride[0] > 1 and kernel_size[0] > stride[0]) or \
                    (stride[1] > 1 and kernel_size[1] > stride[1]) or \
                    (stride[2] > 1 and kernel_size[2] > stride[2]):
                valid_lost = self._non_max_border_amount(f_grad)
                valid_grad.fill_(0)
                overlap_rows = kernel_size[1] - stride[1]
                overlap_cols = kernel_size[2] - stride[2]
                valid_grad[valid_lost.top + overlap_rows:
                           valid_grad.shape[0] - valid_lost.bottom - overlap_rows,
                           valid_lost.left + overlap_cols:
                           valid_grad.shape[1] - valid_lost.right - overlap_cols] = 1

            new_grad_in = valid_grad[None].expand(grad_in[0].shape[1], *valid_grad.shape)[None]
            new_grad_in = (new_grad_in.type(self.dtype) * 10 - 1)
            new_grad_in_lost = self._non_max_border_amount(new_grad_in)
            return (new_grad_in, *grad_in[1:])

    def _backward_saliency_hook(self, module: StreamingConv2d, grad_in, grad_out, is_bias=False, change_grad=True):
        stride: List[int] = _triple(module.stride)  # type:ignore

        # Trim gradient of invalid values
        sides = module.input_loc.sides
        grad_lost = module.grad_lost  # type: Lost

        lost_top = grad_lost.top if not sides.top else 0
        lost_bottom = grad_lost.bottom if not sides.bottom else 0
        lost_left = grad_lost.left if not sides.left else 0
        lost_right = grad_lost.right if not sides.right else 0
        lost = Lost(lost_top, lost_left, lost_bottom, lost_right)

        # Calculate which part of the gradient is 'new'
        new_output_box = module.tile_output_box
        updated_total_indices = module.seen_indices

        if module.in_channels == 3:
            valid_grad_in = grad_in[0][:, :,
                                       lost.top*stride[0]:grad_in[0].shape[2] - lost.bottom*stride[0],
                                       lost.left*stride[1]:grad_in[0].shape[3] - lost.right*stride[1]]

            relevant_input_grad = valid_grad_in[:, :,
                                                new_output_box.y*stride[0]:
                                                new_output_box.y*stride[0] + new_output_box.height*stride[0],
                                                new_output_box.x*stride[1]:
                                                new_output_box.x*stride[1] + new_output_box.width*stride[1]]

            self.saliency_map[:, :,
                              updated_total_indices.y * stride[0]:
                              updated_total_indices.height * stride[0],
                              updated_total_indices.x * stride[1] - relevant_input_grad.shape[3]:
                              updated_total_indices.x * stride[1]] = relevant_input_grad.detach().cpu()

            del relevant_input_grad
            del valid_grad_in

    @staticmethod
    def _new_value_indices(data_shape, data_indices, old_value_indices):
        """
        This helper functions assumes we reconstruct feature maps and
        gradients in tiles from top-left to bottom-right. Using current tile
        index and old_value_indices it finds the relative indices of `data`
        which are unique for this tile (not earlier seen in other tiles).
        """
        rel_top, rel_bottom, rel_left, rel_right = 0, 0, 0, 0

        old_values_y = old_value_indices.y
        old_values_x = old_value_indices.x
        old_values_height = old_value_indices.height

        # Check if new row
        if data_indices.x == 0:
            old_values_y = old_values_height
            old_values_height = data_indices.y + data_shape[H_DIM]
            old_values_x = 0

        # Check x-axis:
        # If this gradient is exactly on the border of old_value_indices
        # everything is new.
        if data_indices.x == old_values_x:
            rel_left = 0
            rel_right = data_shape[W_DIM]

        # If data_indices has some overlap with old_value_indices, trim unique
        # indices.
        else:
            assert old_values_x - data_indices.x >= 0, "Misses data in x-axis!"
            rel_left = old_values_x - data_indices.x
            rel_right = data_shape[W_DIM]

        # Check y-axis:
        # Equal to column logic (see above)
        if data_indices.y == old_values_y:
            rel_top = 0
            rel_bottom = data_shape[H_DIM]
        else:
            assert old_values_y - data_indices.y >= 0, "We miss data in y-axis"
            rel_top = old_values_y - data_indices.y
            rel_bottom = data_shape[H_DIM]

        # Update old-value-indices
        old_values_x += (rel_right - rel_left)

        assert rel_top >= 0, f"We miss data in y-axis before: {data_indices}"
        assert rel_left >= 0, f"We miss data in x-axis before: {data_indices}"

        new_value_indices = Box(rel_top, rel_bottom - rel_top, rel_left, rel_right - rel_left, None)
        old_value_indices = Box(int(old_values_y), int(old_values_height), int(old_values_x), 0, None)

        return new_value_indices, old_value_indices

    def _prev_stats(self, tensor):
        prev = tensor.grad_fn
        prev_stats = None
        while True:
            if prev in self._stats_per_grad_fn:
                prev_stats = self._stats_per_grad_fn[prev]
                break
            if hasattr(prev, 'next_functions') and len(prev.next_functions) > 0:
                prev = prev.next_functions[0][0]
            else:
                break
        return prev_stats

    def state_dict(self):
        named_stats = {
            'net_stats': {}
        }
        for name, module in self.stream_module.named_modules():
            if module in self._module_stats:
                named_stats['net_stats'][name] = self._module_stats[module]
        named_stats['output_stride'] = self.output_stride
        named_stats['tile_output_lost'] = self.tile_output_lost  # type:ignore
        named_stats['tile_gradient_lost'] = self.tile_gradient_lost  # type:ignore
        named_stats['tile_output_shape'] = self._tile_output_shape  # type:ignore
        return named_stats

    def load_state_dict(self, state):
        self.disable()

        self.output_stride = state['output_stride']
        self.tile_output_lost = state['tile_output_lost']
        self.tile_gradient_lost = state['tile_gradient_lost']
        self._tile_output_shape = state['tile_output_shape']

        for name, module in self.stream_module.named_modules():
            if name in state['net_stats']:
                self._module_stats[module] = state['net_stats'][name]

        self.enable()