from collections import OrderedDict
import numpy as np
import os
from hazel.atmosphere import General_atmosphere
from hazel.util import i0_allen
from hazel.codes import sir_code
from hazel.io import Generic_SIR_file
import scipy.interpolate as interp
from hazel.exceptions import NumericalErrorSIR
from hazel.transforms import transformed_to_physical, jacobian_transformation

try:
    from hazel.forward_nn import Forward
except:
    pass

__all__ = ['SIR_atmosphere']

# sir_parameters = OrderedDict.fromkeys('T B thetaB phiB v')

class SIR_atmosphere(General_atmosphere):
    def __init__(self, working_mode, name='', root='', verbose=0):

        super().__init__('photosphere', name=name)

        self.ff = 1.0
        self.macroturbulence = np.zeros(1)
        self.working_mode = working_mode
        self.graphnet_nlte = None
        self.root = root

        self.parameters['T'] = None
        self.parameters['vmic'] = None
        self.parameters['v'] = None
        self.parameters['Bx'] = None
        self.parameters['By'] = None
        self.parameters['Bz'] = None
        self.parameters['ff'] = None
        self.parameters['vmac'] = None

        self.nodes_location['T'] = None
        self.nodes_location['vmic'] = None
        self.nodes_location['v'] = None
        self.nodes_location['Bx'] = None
        self.nodes_location['By'] = None
        self.nodes_location['Bz'] = None
        self.nodes_location['ff'] = None
        self.nodes_location['vmac'] = None

        self.n_nodes['T'] = 0
        self.n_nodes['vmic'] = 0
        self.n_nodes['v'] = 0
        self.n_nodes['Bx'] = 0
        self.n_nodes['By'] = 0
        self.n_nodes['Bz'] = 0
        self.n_nodes['ff'] = 0
        self.n_nodes['vmac'] = 0

        self.nodes['T'] = 0
        self.nodes['vmic'] = 0
        self.nodes['v'] = 0
        self.nodes['Bx'] = 0
        self.nodes['By'] = 0
        self.nodes['Bz'] = 0
        self.nodes['ff'] = 0
        self.nodes['vmac'] = 0

        self.rf_analytical = OrderedDict()
        self.rf_analytical['T'] = None
        self.rf_analytical['vmic'] = None
        self.rf_analytical['v'] = None
        self.rf_analytical['Bx'] = None
        self.rf_analytical['By'] = None
        self.rf_analytical['Bz'] = None
        self.rf_analytical['ff'] = None
        self.rf_analytical['vmac'] = None

        self.ranges['T'] = None
        self.ranges['vmic'] = None
        self.ranges['v'] = None
        self.ranges['Bx'] = None
        self.ranges['By'] = None
        self.ranges['Bz'] = None
        self.ranges['ff'] = None
        self.ranges['vmac'] = None

        self.cycles['T'] = None
        self.cycles['vmic'] = None
        self.cycles['v'] = None
        self.cycles['Bx'] = None
        self.cycles['By'] = None
        self.cycles['Bz'] = None
        self.cycles['ff'] = None
        self.cycles['vmac'] = None

        self.epsilon['T'] = 0.01
        self.epsilon['vmic'] = 0.01
        self.epsilon['v'] = 0.01
        self.epsilon['Bx'] = 0.01
        self.epsilon['By'] = 0.01
        self.epsilon['Bz'] = 0.01
        self.epsilon['ff'] = 0.01
        self.epsilon['vmac'] = 0.01

        self.regularization['T'] = None
        self.regularization['vmic'] = None
        self.regularization['v'] = None
        self.regularization['Bx'] = None
        self.regularization['By'] = None
        self.regularization['Bz'] = None
        self.regularization['ff'] = None
        self.regularization['vmac'] = None

        self.verbose = verbose

    def list_lines(self):
        """
        List the lines available in SIR for synthesis

        """
        f = open('LINEAS', 'r')
        lines = f.readlines()
        f.close()

        print("Available lines:")
        for l in lines[:-1]:
            print(l[:-1])

    def add_active_line(self, lines, spectrum, wvl_range, verbose):
        """
        Add an active lines in this atmosphere

        Parameters
        ----------
        lines : str
            Line to activate
        spectrum : Spectrum
            Spectrum object
        wvl_range : float
            Vector containing wavelength range over which to synthesize this line

        Returns
        -------
        None

        """

        self.lines = lines
        self.wvl_range_lambda = wvl_range

        ind_low = (np.abs(spectrum.wavelength_axis - wvl_range[0])).argmin()
        ind_top = (np.abs(spectrum.wavelength_axis - wvl_range[1])).argmin()

        self.spectrum = spectrum
        self.wvl_axis = spectrum.wavelength_axis[ind_low:ind_top+1]
        self.wvl_range = np.array([ind_low, ind_top+1])

        # Check if Ca II 8542 is in the list of lines and instantiate the neural networks
        if (self.nlte):
            if 301 in self.lines:
                if self.graphnet_nlte is None:
                    path = str(__file__).split('/')
                    checkpoint = '/'.join(path[0:-1])+'/data/20211114-131045_best.prd.pth'
                    if (verbose >= 1):
                        self.logger.info('    * Reading NLTE Neural Network')
                    self.graphnet_nlte = Forward(checkpoint=checkpoint, verbose=verbose)

    def interpolate_nodes(self, log_tau, reference, nodes):
        """
        Generate a model atmosphere by interpolating the defined nodes. The interpolation
        order depends on the number of nodes.

        Parameters
        ----------
        log_tau : float
            Vector of log optical depth at 500 nm
        reference : float
            Vector with the reference atmosphere to which the nodes are added
        nodes : float
            List with the position of the nodes

        Returns
        -------
        real
            Vector with the interpolated atmosphere

        """
        n_nodes = len(nodes)
        n_depth = len(log_tau)

        if (n_nodes == 0):
            return reference, 0

        if (n_nodes == 1):
            return reference + nodes[0], n_depth//2

        # if (n_nodes >= 2):
        #     # pos = np.linspace(n_depth-1, 0, n_nodes+2, dtype=int)[1:-1]
        #     pos = np.linspace(n_depth-1, 0, n_nodes, dtype=int)
        #     f = interp.PchipInterpolator(log_tau[pos], nodes, extrapolate=True)
        #     return reference + f(log_tau), pos

        if (n_nodes == 2):
            # pos = np.linspace(0, n_depth-1, n_nodes+2, dtype=int)[1:-1]
            pos = np.linspace(0, n_depth-1, n_nodes, dtype=int)
            f = interp.interp1d(log_tau[pos], nodes, 'linear', bounds_error=False, fill_value='extrapolate')
            return reference + f(log_tau), pos

        if (n_nodes == 3):
            # pos = np.linspace(0, n_depth-1, n_nodes+2, dtype=int)[1:-1]
            pos = np.linspace(0, n_depth-1, n_nodes, dtype=int)
            f = interp.interp1d(log_tau[pos], nodes, 'quadratic', bounds_error=False, fill_value='extrapolate')
            return reference + f(log_tau), pos

        if (n_nodes > 3):
            # pos = np.linspace(n_depth-1, 0, n_nodes+2, dtype=int)[1:-1]
            pos = np.linspace(n_depth-1, 0, n_nodes, dtype=int)
            f = interp.PchipInterpolator(log_tau[pos], nodes, extrapolate=True)
            return reference + f(log_tau), pos

    def interpolate_nodes_rf(self, log_tau, reference, nodes, lower, upper):
        """
        Generate a model atmosphere by interpolating the defined nodes. The interpolation
        order depends on the number of nodes.

        Parameters
        ----------
        log_tau : float
            Vector of log optical depth at 500 nm
        reference : float
            Vector with the reference atmosphere to which the nodes are added
        nodes : float
            List with the position of the nodes

        Returns
        -------
        real
            Vector with the interpolated atmosphere

        """
        n_nodes = len(nodes)
        n_depth = len(log_tau)

        if (n_nodes == 0):
            return np.zeros(n_depth)

        if (n_nodes == 1):
            rf = np.zeros((n_nodes, n_depth))

            tmp0 = reference + nodes[0]

            # Add the Jacobian to each height
            jacobian = jacobian_transformation(tmp0, lower, upper)

            rf[0,:] = 1.0 * jacobian
            return rf

        if (n_nodes == 2):
            rf = np.zeros((n_nodes, n_depth))

            pos = np.linspace(0, n_depth-1, n_nodes+2, dtype=int)[1:-1]
            f = interp.interp1d(log_tau[pos], nodes, 'linear', bounds_error=False, fill_value='extrapolate')

            tmp0 = reference + f(log_tau)

            jacobian = jacobian_transformation(tmp0, lower, upper)

            delta = 1e-3

            for i in range(n_nodes):
                tmp_nodes = np.copy(nodes)
                tmp_nodes[i] += delta
                f = interp.interp1d(log_tau[pos], tmp_nodes, 'linear', bounds_error=False, fill_value='extrapolate')

                tmp1 = reference + f(log_tau)

                rf[i,:] = (tmp1 - tmp0) / delta * jacobian

            return rf

        if (n_nodes == 3):
            rf = np.zeros((n_nodes, n_depth))

            pos = np.linspace(0, n_depth-1, n_nodes+2, dtype=int)[1:-1]
            f = interp.interp1d(log_tau[pos], nodes, 'quadratic', bounds_error=False, fill_value='extrapolate')
            tmp0 = reference + f(log_tau)

            jacobian = jacobian_transformation(tmp0, lower, upper)

            delta = 1e-3

            for i in range(n_nodes):
                tmp_nodes = np.copy(nodes)
                tmp_nodes[i] += delta
                f = interp.interp1d(log_tau[pos], tmp_nodes, 'quadratic', bounds_error=False, fill_value='extrapolate')

                tmp1 = reference + f(log_tau)

                rf[i,:] = (tmp1 - tmp0) / delta * jacobian

            return rf

        if (n_nodes > 3):
            rf = np.zeros((n_nodes, n_depth))

            pos = np.linspace(n_depth-1, 0, n_nodes+2, dtype=int)[1:-1]
            f = interp.PchipInterpolator(log_tau[pos], nodes, extrapolate=True)

            tmp0 = reference + f(log_tau)

            jacobian = jacobian_transformation(tmp0, lower, upper)

            delta = 1e-3

            for i in range(n_nodes):
                tmp_nodes = np.copy(nodes)
                tmp_nodes[i] += delta
                f = interp.PchipInterpolator(log_tau[pos], tmp_nodes, extrapolate=True)

                tmp1 = reference + f(log_tau)

                rf[i,:] = (tmp1 - tmp0) / delta * jacobian

            return rf

    def load_reference_model(self, model_file, verbose):
        """
        Load a reference model or a model for every pixel for synthesis/inversion

        Parameters
        ----------
        model_file : str
            String with the name of the file. Extensions can currently be "1d" or "h5"
        verbose : bool
            Verbosity

        Returns
        -------
        None
        """
        extension = os.path.splitext(model_file)[1][1:]
        if (extension == '1d'):
            if (verbose >= 1):
                self.logger.info('    * Reading 1D model {0} as reference'.format(model_file))
            self.model_type = '1d'
            self.model_filename = model_file

        if (extension == 'h5'):
            if (verbose >= 1):
                self.logger.info('    * Reading 3D model {0} as reference'.format(model_file))
            self.model_type = '3d'

        self.model_handler = Generic_SIR_file(model_file)
        self.model_handler.open()
        out, ff, vmac = self.model_handler.read(pixel=0)
        self.model_handler.close()

        self.set_parameters(out, ff, vmac)

        self.t_old = np.zeros_like(self.parameters['T'])

        self.init_reference(check_borders=True)

        self.departure = np.ones((2, len(self.lines), len(self.log_tau)))

    def set_parameters(self, model, ff, vmac):
        """
        Set the parameters of the current model to those passed as argument

        Parameters
        ----------
        model_in : float
            Array with the model
        ff : float
            Value of the filling factor
        vmac : float
            Value of the macroturbulent velocity

        Returns
        -------
        None
        """

        self.log_tau = model[:,0]
        self.parameters['T'] = model[:,1]
        self.parameters['vmic'] = model[:,3]
        self.parameters['v'] = model[:,4]
        self.parameters['Bx'] = model[:,5]
        self.parameters['By'] = model[:,6]
        self.parameters['Bz'] = model[:,7]

        if (np.min(model[:,2]) > 0.0):
            self.Pe = model[:,2]
        else:
            self.Pe = -np.ones(len(self.log_tau))
            self.Pe[-1] = 1.11634e-1

        self.parameters['ff'] = ff
        self.parameters['vmac'] = vmac

        # Check that parameters are inside borders by clipping inside the interval with a border of 1e-8
        if (self.working_mode == 'inversion'):
            for k, v in self.parameters.items():
                self.parameters[k] = np.clip(v, self.ranges[k][0] + 1e-8, self.ranges[k][1] - 1e-8)

    def get_parameters(self):
        """
        Get the curent parameters as a model

        Parameters
        ----------
        None

        Returns
        -------
        model: a 6xN photspheric model
        """

        model = np.zeros((len(self.log_tau),8))
        model[:,0] = self.log_tau
        model[:,1] = self.parameters['T']
        model[:,2] = self.Pe
        model[:,3] = self.parameters['vmic']
        model[:,4] = self.parameters['v']
        model[:,5] = self.parameters['Bx']
        model[:,6] = self.parameters['By']
        model[:,7] = self.parameters['Bz']

        return model

    def nodes_to_model(self):
        """
        Transform from nodes to model

        Parameters
        ----------
        None

        Returns
        -------
        None
        """
        for k, v in self.nodes.items():
            if (self.n_nodes[k] > 0):
                self.parameters[k], self.nodes_location[k] = self.interpolate_nodes(self.log_tau, self.reference[k], self.nodes[k])
            else:
                self.parameters[k] = self.reference[k]

        self.Pe = -np.ones(len(self.log_tau))
        self.Pe[-1] = 1.11634e-1

    def model_to_nodes(self):
        """
        Transform from model to nodes

        Parameters
        ----------
        None

        Returns
        -------
        None
        """

        pass

        # for k, v in self.parameters.items():
            # if (k is not 'log_tau'):

        # self.interpolate_nodes(self.parameters['log_tau'], self.reference['T'], [1000])
        # stop()
        # for k, v in self.parameters.items():
            # if (k is not 'log_tau'):

    # def print_parameters_old(self, first=False, error=False):
    #     breakpoint()
    #     for k, v in self.nodes.items():
    #         if (self.n_nodes[k] > 0):
    #             if (k != 'ff'):
    #                 lower = self.ranges[k][0] #- self.eps_borders
    #                 upper = self.ranges[k][1] #+ self.eps_borders
    #                 nodes = transformed_to_physical(v, lower, upper)
    #                 self.logger.info('{0} -> {1}'.format(k, nodes))

    def print_parameters(self, first=False, error=False):
        for k, v in self.parameters.items():
            if (self.n_nodes[k] > 0):
                if (k != 'ff'):
                    pars = v[self.nodes_location[k]]
                    self.logger.info('{0} -> {1}'.format(k, pars))

    def synthesize(self, stokes_in, returnRF=False, nlte=False):
        """
        Carry out the synthesis and returns the Stokes parameters and the response
        functions to all physical variables at all depths

        Parameters
        ----------
        stokes_in : float
            An array of size [4 x nLambda] with the input Stokes parameter. It is irrelevant in this case
            because we assume that all SIR atmospheres have the Planck function as boundary.

        returnRF : bool, optional
            Return response functions

        Returns
        -------

        stokes : float
            Stokes parameters, with the first index containing the wavelength displacement and the remaining
                                    containing I, Q, U and V. Size (5,nLambda)
        rf: float (optional)
            Response functions to T, Pe, vmic, B, v, theta, phi, all of size (4,nLambda,nDepth), plus the RF to macroturbulence of size (4,nLambda)
                            It is not returned if returnRF=False
        """

        if (self.working_mode == 'inversion'):
            self.nodes_to_model()
            self.to_physical()

        if (returnRF):

            stokes, cmass, rf, error = sir_code.synthRF(self.index, self.n_lambda, self.log_tau, self.parameters['T'],
                self.Pe, 1e5*self.parameters['vmic'], 1e5*self.parameters['v'], self.parameters['Bx'], self.parameters['By'],
                self.parameters['Bz'], self.parameters['vmac'])

            if (error == 1):
                raise NumericalErrorSIR()

            B = np.sqrt(self.parameters['Bx']**2 + self.parameters['By']**2 + self.parameters['Bz']**2)

            thetaB = np.arccos(self.parameters['Bz'] / B)
            thetaB[B == 0] = 0.0

            phiB = np.arctan2(self.parameters['By'], self.parameters['Bx'])

            # rfn = np.zeros((150, 73))
            # pars = copy.deepcopy(self.parameters)
            # for pos in range(73):
            #     ind_sto = 0

            #     self.parameters = copy.deepcopy(pars)

            #     delta = 5e-4*self.parameters['T'][pos]
            #     self.parameters['T'][pos] += delta

            #     stokes2, rf2, error = sir_code.synthRF(self.index, self.n_lambda, self.log_tau, self.parameters['T'],
            #         self.Pe, 1e5*self.parameters['vmic'], 1e5*self.parameters['v'], self.parameters['Bx'], self.parameters['By'],
            #         self.parameters['Bz'], self.parameters['vmac'])
            #     error = 0

            #     rfn[:, pos] = (stokes2[ind_sto+1,:]-stokes[ind_sto+1,:])/delta

            # breakpoint()

            # import matplotlib.pyplot as pl
            # pl.plot(rfn[:, 40], '-o', label='numerical')

            self.rf_analytical['T'] = rf[0]+rf[1]       #OK
            self.rf_analytical['vmic'] = 1e5*rf[6]   #OK
            self.rf_analytical['v'] = 1e5*rf[3]      #OK

            # Transform SIR RFs into response functions to Bx, By and Bz.
            # To this end, we have:
            # [RF_B ]   [dBxdB     dBydB    dBzdB ][RF_Bx]
            # [RF_th] = [dBxdthB  dBydthB  dBzdthB][RF_By]
            # [RF_ph]   [dBxdphB  dBydphB  dBzdphB][RF_Bz]
            # and then invert the Jacobian

            RFB = rf[2]
            RFt = rf[4]
            RFp = rf[5]

            self.rf_analytical['Bx'] = RFB * np.sin(thetaB) * np.cos(phiB) + \
                                        RFt * np.cos(thetaB) * np.cos(phiB) / (B + 1e-6) - \
                                        RFp * np.sin(phiB) / (B * np.sin(thetaB))
            self.rf_analytical['By'] = RFB * np.sin(thetaB) * np.sin(phiB) + \
                                        RFt * np.cos(thetaB) * np.sin(phiB) / (B + 1e-6) + \
                                        RFp * np.cos(phiB) / (B * np.sin(thetaB))
            self.rf_analytical['Bz'] = RFB * np.cos(thetaB) - RFt * np.sin(thetaB) / (B + 1e-6)

            self.rf_analytical['vmac'] = rf[7][:, :, None]

            # pl.plot(self.rf_analytical['T'][ind_sto,:,pos], label='analytical')
            # pl.legend()
            # pl.show()

            # import matplotlib.pyplot as pl
            # f, ax = pl.subplots(nrows=3, ncols=3, figsize=(10,10))
            # ax = ax.flatten()
            # for i in range(9):
            #     ax[i].plot(np.log(self.rf_analytical['T'][0,i*10,:]), color=f'C{i}')
            #     ax[i].plot(np.log(rfn[i*10,:]), 'o', color=f'C{i}')
            # pl.show()

            i0 = i0_allen(np.mean(self.wvl_axis), self.spectrum.mu)

            for k, v in self.nodes.items():
                if (k != 'vmac'):
                    if (self.n_nodes[k] > 0):
                        lower = self.ranges[k][0]
                        upper = self.ranges[k][1]
                        rf = self.interpolate_nodes_rf(self.log_tau, self.reference[k], self.nodes[k], lower, upper)

                        # import matplotlib.pyplot as pl
                        # f, ax = pl.subplots(nrows=3, ncols=3, figsize=(10,10))
                        # ax = ax.flatten()
                        # for i in range(9):
                        #     ax[i].plot(rfn[i*10,:] / self.rf_analytical['T'][0,i*10,:], color=f'C{i}')
                        #     ax[i].set_ylim([0,2])
                        # pl.show()
                        # print(k)
                        # breakpoint()

                        self.rf_analytical[k] = np.einsum('ijk,lk->ijl', self.rf_analytical[k], rf) * i0

            return self.parameters['ff'] * stokes[1:,:] * i0, self.rf_analytical, error
        else:
            # stokes, error = sir_code.synth(self.index, self.n_lambda, self.log_tau, self.parameters['T'],
                # self.Pe, 1e5*self.parameters['vmic'], 1e5*self.parameters['v'], self.parameters['Bx'], self.parameters['By'],
                # self.parameters['Bz'], self.parameters['vmac'])

            # If we need to put the atmosphere in hydrostatic eq.
            if (self.working_mode == 'inversion'):
                self.Pe = sir_code.hydroeq(self.log_tau, self.parameters['T'],
                    self.Pe, 1e5*self.parameters['vmic'], 1e5*self.parameters['v'], self.parameters['Bx'], self.parameters['By'],
                    self.parameters['Bz'])

            # Check if the line is 8542 and we want NLTE. If that is the case, then evaluate the
            # neural network to return the departure coefficients

            if (nlte):
                if (self.nlte):
                    dif = (self.parameters['T'] - self.t_old)
                    if (np.max(dif) > self.t_change_departure):
                        for i, l in enumerate(self.lines):
                            if (l == 301):
                                if (self.verbose >= 4):
                                    self.logger.info('  - NLTE neural oracle')
                                n = len(self.log_tau)
                                tau = [10.0**self.log_tau[::-1]]
                                ne = self.Pe / (1.381e-16 * self.parameters['T'])
                                ne = [ne[::-1] * 1e6]                                 # in m^-3
                                tt = [self.parameters['T'][::-1]]
                                vturb = [self.parameters['vmic'][::-1] * 1e3]         # in m/s
                                vlos = [self.parameters['v'][::-1] * 1e3]             # in m/s
                                prediction = self.graphnet_nlte.predict(tau, ne, vturb, tt, vlos)
                                self.departure[0, i, :] = 10.0**prediction[0][::-1, 2]
                                self.departure[1, i, :] = 10.0**prediction[0][::-1, 4]

                            self.t_old = self.parameters['T']

            else:
                self.departure = np.ones((2, len(self.lines), len(self.log_tau)))

            stokes, cmass, rf, error = sir_code.synthRF(self.index, self.n_lambda, self.log_tau, self.parameters['T'],
                self.Pe, 1e5*self.parameters['vmic'], 1e5*self.parameters['v'], self.parameters['Bx'], self.parameters['By'],
                self.parameters['Bz'], self.parameters['vmac'], np.asfortranarray(self.departure))

            # Transform SIR RFs into response functions to Bx, By and Bz.
            # To this end, we have:
            # [RF_B ]   [dBxdB     dBydB    dBzdB ][RF_Bx]
            # [RF_th] = [dBxdthB  dBydthB  dBzdthB][RF_By]
            # [RF_ph]   [dBxdphB  dBydphB  dBzdphB][RF_Bz]
            # and then invert the Jacobian
            B = np.sqrt(self.parameters['Bx']**2 + self.parameters['By']**2 + self.parameters['Bz']**2) + 1e-6
            thetaB = np.arccos(self.parameters['Bz'] / B)
            thetaB[B == 0] = 0.0
            phiB = np.arctan2(self.parameters['By'], self.parameters['Bx'])

            self.rf_analytical['T'] = rf[0]+rf[1]       #OK
            self.rf_analytical['vmic'] = 1e5*rf[6]   #OK
            self.rf_analytical['v'] = 1e5*rf[3]      #OK

            RFB = rf[2]
            RFt = rf[4]
            RFp = rf[5]

            self.rf_analytical['Bx'] = RFB * np.sin(thetaB) * np.cos(phiB) + \
                                        RFt * np.cos(thetaB) * np.cos(phiB) / B - \
                                        RFp * np.sin(phiB) / (B * np.sin(thetaB))
            self.rf_analytical['By'] = RFB * np.sin(thetaB) * np.sin(phiB) + \
                                        RFt * np.cos(thetaB) * np.sin(phiB) / B + \
                                        RFp * np.cos(phiB) / (B * np.sin(thetaB))
            self.rf_analytical['Bz'] = RFB * np.cos(thetaB) - RFt * np.sin(thetaB) / B

            self.rf_analytical['vmac'] = rf[7][:, :, None]

            if (error == 1):
                raise NumericalErrorSIR()

            return self.parameters['ff'] * stokes[1:,:] * i0_allen(np.mean(self.wvl_axis), self.spectrum.mu), error #hsra_continuum(np.mean(self.wvl_axis)), error


# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

from audioop import bias
from os import X_OK
from re import M
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models
import numpy as np
from torch.distributions.distribution import Distribution
import torchvision.models.alexnet
from collections import OrderedDict
from itertools import chain

from domainbed.lib import wide_resnet
from domainbed.lib import gaussian_random_fields as grf
import copy
import time
import math
# from imgaug import augmenters as iaa
# from torchvision import transforms
import torch.utils.model_zoo as model_zoo

model_urls = {
    'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',
}

def remove_batch_norm_from_resnet(model):
    fuse = torch.nn.utils.fusion.fuse_conv_bn_eval
    model.eval()

    model.conv1 = fuse(model.conv1, model.bn1)
    model.bn1 = Identity()

    for name, module in model.named_modules():
        if name.startswith("layer") and len(name) == 6:
            for b, bottleneck in enumerate(module):
                for name2, module2 in bottleneck.named_modules():
                    if name2.startswith("conv"):
                        bn_name = "bn" + name2[-1]
                        setattr(bottleneck, name2,
                                fuse(module2, getattr(bottleneck, bn_name)))
                        setattr(bottleneck, bn_name, Identity())
                if isinstance(bottleneck.downsample, torch.nn.Sequential):
                    bottleneck.downsample[0] = fuse(bottleneck.downsample[0],
                                                    bottleneck.downsample[1])
                    bottleneck.downsample[1] = Identity()
    model.train()
    return model

class Identity(nn.Module):
    """An identity layer"""
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

class MLP(nn.Module):
    """Just  an MLP"""
    def __init__(self, n_inputs, n_outputs, hparams):
        super(MLP, self).__init__()
        self.input = nn.Linear(n_inputs, hparams['mlp_width'])
        self.dropout = nn.Dropout(hparams['mlp_dropout'])
        self.hiddens = nn.ModuleList([
            nn.Linear(hparams['mlp_width'], hparams['mlp_width'])
            for _ in range(hparams['mlp_depth']-2)])
        self.output = nn.Linear(hparams['mlp_width'], n_outputs)
        self.n_outputs = n_outputs

    def forward(self, x):
        x = self.input(x)
        x = self.dropout(x)
        x = F.relu(x)
        for hidden in self.hiddens:
            x = hidden(x)
            x = self.dropout(x)
            x = F.relu(x)
        x = self.output(x)
        return x

class ResNet(torch.nn.Module):
    """ResNet with the softmax chopped off and the batchnorm frozen"""
    def __init__(self, input_shape, hparams):
        super(ResNet, self).__init__()
        if hparams['resnet18']:
            self.network = torchvision.models.resnet18(pretrained=True)
            self.n_outputs = 512
        else:
            self.network = torchvision.models.resnet50(pretrained=True)
            self.n_outputs = 2048

        # self.network = remove_batch_norm_from_resnet(self.network)

        # adapt number of channels
        nc = input_shape[0]
        if nc != 3:
            tmp = self.network.conv1.weight.data.clone()

            self.network.conv1 = nn.Conv2d(
                nc, 64, kernel_size=(7, 7),
                stride=(2, 2), padding=(3, 3), bias=False)

            for i in range(nc):
                self.network.conv1.weight.data[:, i, :, :] = tmp[:, i % 3, :, :]

        # save memory
        del self.network.fc
        self.network.fc = Identity()

        self.freeze_bn()
        self.hparams = hparams
        self.dropout = nn.Dropout(hparams['resnet_dropout'])

    def forward(self, x):
        """Encode x into a feature vector of size n_outputs."""
        return self.dropout(self.network(x))

    def train(self, mode=True):
        """
        Override the default train() to freeze the BN parameters
        """
        super().train(mode)
        self.freeze_bn()

    def freeze_bn(self):
        for m in self.network.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()

class AlexNet(nn.Module):

    def __init__(self, hparams):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )

        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        # self.classifier = nn.Sequential(
        #     nn.Dropout(p=True),
        #     nn.Linear(256 * 6 * 6, 4096),
        #     nn.ReLU(inplace=True),
        #     nn.Dropout(p=True),
        #     nn.Linear(4096, 4096),
        #     nn.ReLU(inplace=True),
        #     nn.Linear(4096, 1000),
        # )

        self.n_outputs = 256 * 6 * 6
        self.hparams = hparams

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
    #     x = self.classifier(x)
        return x

class AlexNetCaffe(nn.Module):
    def __init__(self, n_classes=7, hparams=None):
        super(AlexNetCaffe, self).__init__()
        print("Using Caffe AlexNet")
        self.features = nn.Sequential(OrderedDict([
            ("conv1", nn.Conv2d(3, 96, kernel_size=11, stride=4)),
            ("relu1", nn.ReLU(inplace=True)),
            ("pool1", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)),
            ("norm1", nn.LocalResponseNorm(5, 1.e-4, 0.75)),
            ("conv2", nn.Conv2d(96, 256, kernel_size=5, padding=2, groups=2)),
            ("relu2", nn.ReLU(inplace=True)),
            ("pool2", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)),
            ("norm2", nn.LocalResponseNorm(5, 1.e-4, 0.75)),
            ("conv3", nn.Conv2d(256, 384, kernel_size=3, padding=1)),
            ("relu3", nn.ReLU(inplace=True)),
            ("conv4", nn.Conv2d(384, 384, kernel_size=3, padding=1, groups=2)),
            ("relu4", nn.ReLU(inplace=True)),
            ("conv5", nn.Conv2d(384, 256, kernel_size=3, padding=1, groups=2)),
            ("relu5", nn.ReLU(inplace=True)),
            ("pool5", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)),
        ]))
        # self.classifier = nn.Sequential(OrderedDict([
        #     ("fc6", nn.Linear(256 * 6 * 6, 4096)),
        #     ("relu6", nn.ReLU(inplace=True)),
        #     ("drop6", nn.Dropout()),
        #     ("fc7", nn.Linear(4096, 4096)),
        #     ("relu7", nn.ReLU(inplace=True)),
        #     ("drop7", nn.Dropout())]))

        # self.classifier = nn.Linear(4096, n_classes)

        self.n_outputs = 256 * 6 * 6
        self.hparams = hparams
        # self.p_logvar = nn.Sequential(nn.Linear(4096, 512),
        #                               nn.ReLU())
        # self.p_mu = nn.Sequential(nn.Linear(4096, 512),
        #                           nn.LeakyReLU())

    def get_params(self, base_lr):
        return [{"params": self.features.parameters(), "lr": 0.},
                {"params": chain(self.classifier.parameters(), self.classifier_l.parameters(), self.p_logvar.parameters(), self.p_mu.parameters()
                                 ), "lr": base_lr}]

    def is_patch_based(self):
        return False

    def forward(self, x, train=True):
        # end_points={}
        x = self.features(x*57.6)  #57.6 is the magic number needed to bring torch data back to the range of caffe data, based on used std
        x = x.view(x.size(0), -1)
        # x = self.classifier(x)

        # logvar = self.p_logvar(x)
        # mu = self.p_mu(x)
        # end_points['logvar'] = logvar
        # end_points['mu'] = mu

        # if train:
        #     x = self.reparametrize(mu, logvar)
        # else:
        #     x = mu

        # end_points['Embedding'] = x
        # x = self.classifier_l(x)
        # end_points['Predictions'] = nn.functional.softmax(input=x, dim=-1)

        # return x, end_points
        return x

    def reparametrize(self, mu, logvar, factor=0.2):
        std = logvar.div(2).exp()
        eps = std.data.new(std.size()).normal_()
        return mu + factor*std*eps

def alexnet(pretrained=True, num_classes=1000, feature_size=6, model_path=None):
    r"""AlexNet model architecture from the
    `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        num_classes (int): number of classes
        model_path (string): path to pretrain model (using official model if is None)
    """
    model = AlexNet()
    if pretrained:
        if model_path is None:
            model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
        else:
            print("Loading model {}".format(model_path))
            checkpoint = torch.load(model_path, map_location='cpu')
            model.load_state_dict(checkpoint['state_dict'])

    if feature_size == 6:
        model.classifier[6] = nn.Linear(4096, num_classes)
        # nn.init.xavier_uniform_(model.classifier[-1].weight, .1)
        # nn.init.constant_(model.classifier[-1].bias, 0.)
    else:
        print("Adpating new classifier")
        fc_size = feature_size*feature_size*256
        model.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(fc_size, fc_size),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(fc_size, fc_size),
            nn.ReLU(inplace=True),
            nn.Linear(fc_size, num_classes),
        )
        for m in model.classifier.modules():
            if type(m) == nn.Linear:
                nn.init.xavier_uniform_(m.weight)
                nn.init.constant_(m.bias, 0.)

    return model

class DigitsConvNet(nn.Module):

    def __init__(self, input_shape, hparams):
        super(DigitsConvNet, self).__init__()

        self.batch_size = hparams['batch_size']

        kernel_size = 5
        stride = 1
        padding = 0
        dim_conv1 = 64
        dim_conv2 = 128
        self.conv1 = nn.Conv2d(input_shape[0], dim_conv1, kernel_size=kernel_size, stride=stride, padding=padding)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(dim_conv1, dim_conv2, kernel_size=kernel_size, stride=stride, padding=padding)
        self.relu2 = nn.ReLU(inplace=True)
        self.mp = nn.MaxPool2d(2)
        w = int(np.floor((np.floor((input_shape[-1] - kernel_size + 2 * padding) / stride + 1) // 2 - kernel_size + 2 * padding) / stride + 1) // 2)
        self.n_outputs = w ** 2 * dim_conv2

    def forward(self, x):

        in_size = x.size(0)
        out1 = self.mp(self.relu1(self.conv1(x)))
        out2 = self.mp(self.relu2(self.conv2(out1)))
        out2 = out2.view(in_size, -1)
        return out2

class GeneralClassifier(nn.Module):
    def __init__(self, in_features, out_features, hparams):
        super(GeneralClassifier, self).__init__()
        self.is_nonlinear = hparams['nonlinear_classifier']
        self.is_proj = hparams['is_proj']
        self.dim_proj = hparams['dim_proj']
        self.is_manual_fc_dim = hparams['is_manual_fc_dim']
        self.dim_fc = hparams['dim_fc']
        self.alexnet = hparams['alexnet'] and hparams['alexnet_classifier']

        if self.alexnet:
            if self.is_proj:
                self.classifier = nn.Sequential(
                    nn.Dropout(p=hparams['alexnet_dropout']),
                    nn.Linear(in_features, 4096),
                    nn.ReLU(inplace=True),
                    nn.Dropout(p=hparams['alexnet_dropout']),
                    nn.Linear(4096, 4096),
                    nn.ReLU(inplace=True),
                )
                self.head = nn.Linear(4096, out_features)
                self.pro_head = nn.Linear(4096, self.dim_proj)
            else:
                self.classifier = nn.Sequential(
                    nn.Dropout(p=hparams['alexnet_dropout']),
                    nn.Linear(in_features, 4096),
                    nn.ReLU(inplace=True),
                    nn.Dropout(p=hparams['alexnet_dropout']),
                    nn.Linear(4096, 4096),
                    nn.ReLU(inplace=True),
                    nn.Linear(4096, out_features),
                )
        else:
            if self.is_nonlinear:
                if self.is_manual_fc_dim:
                    self.fc1 = torch.nn.Linear(in_features, self.dim_fc)
                    self.relu1 = torch.nn.ReLU()
                    self.fc2 = torch.nn.Linear(self.dim_fc, self.dim_fc)
                    self.relu2 = torch.nn.ReLU()
                    self.fc3 = torch.nn.Linear(self.dim_fc, out_features)
                    if self.is_proj:
                        self.pro_head = nn.Linear(self.dim_fc, self.dim_proj)
                else:
                    self.fc1 = torch.nn.Linear(in_features, in_features // 2)
                    self.relu1 = torch.nn.ReLU()
                    self.fc2 = torch.nn.Linear(in_features // 2, in_features // 4)
                    self.relu2 = torch.nn.ReLU()
                    self.fc3 = torch.nn.Linear(in_features // 4, out_features)
                    if self.is_proj:
                        self.pro_head = nn.Linear(in_features // 4, self.dim_proj)
            else:
                self.fc1 = torch.nn.Linear(in_features, out_features)
                if self.is_proj:
                    self.pro_head = nn.Linear(in_features, self.dim_proj)

    def forward(self, x, mode='test'):
        if self.alexnet:
            if self.is_proj:
                x = self.classifier(x)
                p = self.head(x)
            else:
                p = self.classifier(x)
        else:
            if self.is_nonlinear:
                x = self.fc1(x)
                x = self.relu1(x)
                x = self.fc2(x)
                x = self.relu2(x)
                p = self.fc3(x)
            else:
                p = self.fc1(x)
        if mode == 'test':
            return p
        elif mode == 'tsne':
            return p, x
        elif mode =='train':
            if self.is_proj:
                z = self.pro_head(x)
                z = F.normalize(z)
                return p,z
            else:
                return p,None

class MNIST_CNN(nn.Module):
    """
    Hand-tuned architecture for MNIST.
    Weirdness I've noticed so far with this architecture:
    - adding a linear layer after the mean-pool in features hurts
        RotatedMNIST-100 generalization severely.
    """
    n_outputs = 128

    def __init__(self, input_shape, hparams):
        super(MNIST_CNN, self).__init__()

        self.conv1 = nn.Conv2d(input_shape[0], 64, 3, 1, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(128, 128, 3, 1, padding=1)
        self.conv4 = nn.Conv2d(128, 128, 3, 1, padding=1)

        self.bn0 = nn.GroupNorm(8, 64)
        self.bn1 = nn.GroupNorm(8, 128)
        self.bn2 = nn.GroupNorm(8, 128)
        self.bn3 = nn.GroupNorm(8, 128)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.bn0(x)

        x = self.conv2(x)
        x = F.relu(x)
        x = self.bn1(x)

        x = self.conv3(x)
        x = F.relu(x)
        x = self.bn2(x)

        x = self.conv4(x)
        x = F.relu(x)
        x = self.bn3(x)

        x = self.avgpool(x)
        x = x.view(len(x), -1)
        return x

class ContextNet(nn.Module):
    def __init__(self, input_shape):
        super(ContextNet, self).__init__()

        # Keep same dimensions
        padding = (5 - 1) // 2
        self.context_net = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, 5, padding=padding),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, 5, padding=padding),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 1, 5, padding=padding),
        )

    def forward(self, x):
        return self.context_net(x)

def Featurizer(input_shape, hparams):
    """Auto-select an appropriate featurizer for the given input shape."""
    if hparams['digits_convnet']:
        return DigitsConvNet(input_shape, hparams)
    else:
        if len(input_shape) == 1:
            # print('MLP')
            return MLP(input_shape[0], hparams["mlp_width"], hparams)
        elif input_shape[1:3] == (28, 28):
            # print('MNIST_CNN')
            return MNIST_CNN(input_shape, hparams)
        elif input_shape[1:3] == (32, 32):
            # print('Wide_ResNet')
            if hparams['wide_resnet_type'] == 1:
                return wide_resnet.Wide_ResNet(input_shape, hparams['wide_resnet_layers'], hparams['widen_factor'], hparams['widen_dropout'])
            elif hparams['wide_resnet_type'] == 2:
                return wide_resnet.Wide_ResNet2(input_shape, hparams['wide_resnet_layers'], hparams['widen_factor'], hparams['widen_dropout'])
            else:
                NotImplementedError

        elif input_shape[1:3] == (224, 224):
            # print('ResNet')

            if hparams['alexnet']:
                if hparams['alexnet_flag'] == 1:
                    model = AlexNet(hparams)
                    model.load_state_dict(model_zoo.load_url(model_urls['alexnet']), strict=False)
                    # del model.classifier
                else:
                    model = AlexNetCaffe(hparams)
                    for m in model.modules():
                        if isinstance(m, nn.Linear):
                            nn.init.xavier_uniform_(m.weight, .1)
                            nn.init.constant_(m.bias, 0.)
                    state_dict = torch.load("./domainbed/data/alexnet/alexnet_caffe.pth.tar")
                    del state_dict["classifier.fc8.weight"]
                    del state_dict["classifier.fc8.bias"]
                    model.load_state_dict(state_dict, strict=False)
                    # del model.classifier

                return model
            else:
                return ResNet(input_shape, hparams)
        else:
            return ResNet(input_shape, hparams)
            # raise NotImplementedError

def Classifier(in_features, out_features, is_nonlinear=False):
    if is_nonlinear:
        return torch.nn.Sequential(
            torch.nn.Linear(in_features, in_features // 2),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features // 2, in_features // 4),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features // 4, out_features))
    else:
        return torch.nn.Linear(in_features, out_features)

class WholeFish(nn.Module):
    def __init__(self, input_shape, num_classes, hparams, weights=None):
        super(WholeFish, self).__init__()
        featurizer = Featurizer(input_shape, hparams)
        # classifier = Classifier(
        #     featurizer.n_outputs,
        #     num_classes,
        #     hparams['nonlinear_classifier'])
        classifier = GeneralClassifier(
            featurizer.n_outputs,
            num_classes,
            hparams)
        self.net = nn.Sequential(
            featurizer, classifier
        )
        if weights is not None:
            self.load_state_dict(copy.deepcopy(weights))

    def reset_weights(self, weights):
        self.load_state_dict(copy.deepcopy(weights))

    def forward(self, x):
        return self.net(x)

## ------------------- for PDEN ---------------------
class AdaIN2d(nn.Module):
    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm2d(num_features, affine=False)
        self.fc = nn.Linear(style_dim, num_features*2)
    def forward(self, x, s):
        h = self.fc(s)
        h = h.view(h.size(0), h.size(1), 1, 1)
        gamma, beta = torch.chunk(h, chunks=2, dim=1)
        return (1 + gamma) * self.norm(x) + beta
        #return (1+gamma)*(x)+beta

class IN2d(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm2d(num_features, affine=False)
    def forward(self, x, y=None):
        return self.norm(x)

class selfIN2d(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm2d(num_features, affine=False)
    def forward(self, x, w, b):
        return w * self.norm(x) + b

class GeneralizedAdaIN(nn.Module):
    def __init__(self, norm_opt = []):
    # def __init__(self, flag = None, num_features = 0, style = 'batch', same=True, batch_size=0, num_classes = 0, inv_prop = 0.0, std1 = 1.0, std2 = 1.0, v_std1=False, v_std2=False, norm_true=True, not_used='none'):
        super().__init__()

        self.flag_gamma = norm_opt['flag_gamma']
        self.flag_beta = norm_opt['flag_beta']

        self.num_features = norm_opt['num_features']
        self.style = norm_opt['style']
        self.same = norm_opt['same']
        self.batch_size = norm_opt['batch_size']
        self.num_classes = norm_opt['num_classes']
        self.inv_prop = norm_opt['inv_prop']
        self.std1 = norm_opt['std1']
        self.std2 = norm_opt['std2']
        self.v_std1 = norm_opt['v_std1']
        self.v_std2 = norm_opt['v_std2']
        self.norm_true = norm_opt['norm_true']
        self.not_used = norm_opt['not_used']
        self.norm = nn.InstanceNorm2d(self.num_features, affine=False)

        if self.v_std1 == 'first_gen':
            self.apply_std1 = self.std1*torch.rand(1).item()
            # print('std1:{}'.format(self.apply_std1))
        else:
            self.apply_std1 = self.std1
        if self.v_std2 == 'first_gen':
            self.apply_std2 = self.std2*torch.rand(1).item()
            # print('std2:{}'.format(self.apply_std2))
        else:
            self.apply_std2 = self.std2

        if self.same:
            self.gamma, self.beta = self.make_parameter()

    def make_parameter(self):
        if self.v_std1 == 'each_gen':
            self.apply_std1 = self.std1*torch.rand(1).item()
            # print('std11:{}'.format(self.apply_std1))
        if self.v_std2 == 'each_gen':
            self.apply_std2 = self.std2*torch.rand(1).item()
            # print('std22:{}'.format(self.apply_std2))

        if self.style == 'instance':
            num_b = self.batch_size
        elif self.style == 'class':
            num_b = self.num_classes
        elif self.style == 'batch':
            num_b = 1
        else:
            NotImplementedError

        if self.flag_gamma == 'lognormal':
            gamma = torch.exp(torch.normal(mean = 0, std = self.apply_std1 * torch.ones(num_b, self.num_features, 1, 1))).cuda() # log normal distribution
            if self.inv_prop > 0.0:
                idx = torch.rand(num_b, self.num_features, 1, 1) < self.inv_prop
                gamma[idx] = -gamma[idx]
        elif self.flag_gamma == 'normal':
            gamma = torch.normal(mean = 0, std = self.apply_std1 * torch.ones(num_b, self.num_features, 1, 1)).cuda() # normal distribution
        elif self.flag_gamma == 'uniform':
            gamma = self.apply_std1 * (2.0*torch.rand(num_b, self.num_features, 1, 1).cuda()-1) # uniform distribution

        if self.flag_beta == 'lognormal':
            beta = torch.exp(torch.normal(mean = 0, std = self.apply_std2 * torch.ones(num_b, self.num_features, 1, 1))).cuda() # log normal distribution
            if self.inv_prop > 0.0:
                idx = torch.rand(num_b, self.num_features, 1, 1) < self.inv_prop
                beta[idx] = -beta[idx]
        elif self.flag_beta == 'normal':
            beta = torch.normal(mean = 0, std = self.apply_std2 * torch.ones(num_b, self.num_features, 1, 1)).cuda() # normal distribution
        elif self.flag_beta == 'uniform':
            beta = self.apply_std2 * (2.0*torch.rand(num_b, self.num_features, 1, 1).cuda()-1) # uniform distribution

        if self.not_used == 'gamma':
            gamma[:] = 1.0
        elif self.not_used == 'beta':
            beta[:] = 0.0

        return gamma, beta

    def forward(self, x, y=None): # old: not same, instance, lognormal
        if not self.same: # re-initialization
            self.gamma, self.beta = self.make_parameter()

        if self.style == 'class':
            gamma = torch.zeros([self.batch_size, self.gamma.size(1), self.gamma.size(2), self.gamma.size(3)]).cuda()
            beta = torch.zeros([self.batch_size, self.gamma.size(1), self.gamma.size(2), self.gamma.size(3)]).cuda()

            for y_local in range(self.num_classes):
                idx = [j for j, k in enumerate(y) if k == y_local]
                if len(idx) > 0:
                    for idx_local in idx:
                        gamma[idx_local] = self.gamma[y_local,:,:,:]
                        beta[idx_local] = self.beta[y_local,:,:,:]
        else:
            gamma = self.gamma
            beta = self.beta
        if self.norm_true:
            return gamma * self.norm(x) + beta
        else:
            return gamma * x + beta

class cnnGenerator(nn.Module):
    def __init__(self, n=16, kernelsize=3, imdim=3, imsize=[192, 320], zdim=10):
        super().__init__()
        stride = (kernelsize-1)//2
        self.zdim = zdim
        self.imdim = imdim
        self.imsize = imsize

        self.conv1 = nn.Conv2d(imdim, n, kernelsize, 1, stride)
        self.conv2 = nn.Conv2d(n, 2*n, kernelsize, 1, stride)
        self.adain2 = AdaIN2d(zdim, 2*n)
        self.conv3 = nn.Conv2d(2*n, 4*n, kernelsize, 1, stride)
        self.conv4 = nn.Conv2d(4*n, imdim, kernelsize, 1, stride)

    def forward(self, x, rand=False):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        if rand:
            z = torch.randn(len(x), self.zdim).cuda()
            x = self.adain2(x, z)
        x = F.relu(self.conv3(x))
        x = torch.sigmoid(self.conv4(x))
        return x

class RandGenerator(nn.Module):
    def __init__(self, im_dim=3, im_size=[32, 32], gen_model=None, first_flag=False, gen_num=1, n_tgt=0, step_base_cnt=0, opt={}):
        super().__init__()

        self.structure = gen_model['structure']
        self.init_norm_model = gen_model['init_norm_model']
        self.deform_offset_style = gen_model['deform_offset_style']
        self.deform_offset_clamp = gen_model['deform_offset_clamp']
        self.deform_style = gen_model['deform_style']
        self.deform_same = gen_model['deform_same']
        self.deform_global = gen_model['deform_global']
        self.deform_variable = gen_model['deform_variable']
        self.deform_variable_range = gen_model['deform_variable_range']

        self.noise_model = gen_model['noise_model']
        self.noise_style = gen_model['noise_style']
        self.noise_channel = gen_model['noise_channel']
        self.noise_prop = gen_model['noise_prop']
        self.noise_scale = gen_model['noise_scale']
        self.noise_alpha = gen_model['noise_alpha']
        self.noise_first = gen_model['noise_first']

        self.augment_model = gen_model['augment_model']

        self.encoder_model = gen_model['encoder_model']
        self.encoder_style = gen_model['encoder_style']
        self.encoder_same = gen_model['encoder_same']
        self.encoder_act = gen_model['encoder_act']
        self.encoder_num = gen_model['encoder_num']

        self.norm_model = gen_model['norm_model']
        self.norm_style = gen_model['norm_style']
        self.norm_same = gen_model['norm_same']
        self.norm_std1 = gen_model['norm_std1']
        self.norm_std2 = gen_model['norm_std2']
        self.norm_not_used = gen_model['norm_not_used']
        self.norm_variable_std1 = gen_model['norm_variable_std1']
        self.norm_variable_std2 = gen_model['norm_variable_std2']
        self.norm_true = gen_model['norm_true']

        self.norm_inv_prop = gen_model['norm_inv_prop']
        self.decoder_model = gen_model['decoder_model']
        self.decoder_style = gen_model['decoder_style']
        self.decoder_act = gen_model['decoder_act']
        self.decoder_num = gen_model['decoder_num']
        self.final_act = gen_model['final_act']
        self.final_norm = gen_model['final_norm']
        self.initialization = gen_model['initialization']
        self.dim = gen_model['dim']
        self.zdim = gen_model['zdim']
        self.bias = gen_model['bias']
        self.input_normalize = gen_model['input_normalize']
        self.num_classes = gen_model['num_classes']
        self.batch_size = gen_model['batch_size']

        self.kernelsize = gen_model['kernelsize']
        self.multiconv_kernelsize = gen_model['multiconv_kernelsize']
        self.multiconv_same = gen_model['multiconv_same']

        self.variable_kernelsize = gen_model['variable_kernelsize']
        self.variable_times = gen_model['variable_times']
        self.variable_th = gen_model['variable_th']
        self.variable_limit = gen_model['variable_limit']
        self.variable_updated_flag = gen_model['variable_updated_flag']
        self.dilation = gen_model['dilation']

        self.im_size = im_size
        self.im_dim = im_dim
        self.gen_num = gen_num

        self.n_tgt = n_tgt
        self.step_base_cnt = step_base_cnt
        self.opt = opt

        self.init_norm = []
        self.encoder = []
        self.norm = []
        self.decoder = []
        self.encoder_activation = []
        self.decoder_activation = []

        self.debug_output = []
        self.debug_name = []

        if self.structure.lower() == 'pden':
            self.encoder.append(self.conv_flag(self.im_dim,   self.dim, self.bias, self.encoder_model, self.encoder_style))
            self.encoder.append(self.conv_flag(self.dim,    2*self.dim, self.bias, self.encoder_model, self.encoder_style))
            self.norm.append(AdaIN2d(self.zdim, 2*self.dim))
            self.decoder.append(self.conv_flag(2*self.dim,  4*self.dim, self.bias, self.decoder_model, self.encoder_style))
            self.decoder.append(self.conv_flag(4*self.dim, self.im_dim, self.bias, self.decoder_model, self.encoder_style))

        elif self.structure == 'randconv':

            if 'multi' in self.encoder_model:
                kernelsize = self.multiconv_kernelsize
                try:
                    kernelsize = [int(i) for i in kernelsize.split('_')]
                except:
                    kernelsize = [int(i[1:]) for i in kernelsize.split('_')]
                kernelsize = kernelsize[np.random.randint(len(kernelsize))]
                self.kernelsize = kernelsize
            self.encoder.append(self.conv_flag(self.im_dim, self.im_dim, self.bias, self.encoder_model, self.encoder_style))

        elif self.structure == 'vae':

            # if self.augment_model == 'randaug':
            #     print('.')
            #     ia.seed(4)
            #     self.aug_sequence = iaa.Affine(rotate=(-25, 25))

            #     torchvision.transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)

            #     augment_transform.append(transforms.ColorJitter(0.3, 0.3, 0.3, 0.3))
            #     tfs = transforms.Compose([
            #         iaa.Sequential([
            #             iaa.flip.Fliplr(p=0.5),
            #             iaa.flip.Flipud(p=0.5),
            #             iaa.GaussianBlur(sigma=(0.0, 0.1)),
            #             iaa.MultiplyBrightness(mul=(0.65, 1.35)),
            #         ]).augment_image,
            #         transforms.ToTensor()
            #     ])

            if self.init_norm_model == 'self_in':
                self.init_norm.append(selfIN2d(self.im_dim))
            elif self.init_norm_model == 'in': # instance norm
                self.init_norm.append(IN2d(self.im_dim))

            self.conv_dim = self.dim
            if self.decoder_num > 0:
                self.norm_dim = self.conv_dim
            else:
                self.norm_dim = self.im_dim
            if self.encoder_num == 0:
                self.norm_dim = self.im_dim

            if self.variable_kernelsize == 'none':
                if 'multi' in self.encoder_model:
                    kernelsize = self.multiconv_kernelsize
                    try:
                        kernelsize = [int(i) for i in kernelsize.split('_')]
                    except:
                        kernelsize = [int(i[1:]) for i in kernelsize.split('_')]
                    kernelsize = kernelsize[np.random.randint(len(kernelsize))]
                    self.kernelsize = kernelsize
            elif self.variable_kernelsize == 'increased_gen':
                self.kernelsize = int(self.gen_num * 2 + 1) # first kernelsize is 3
            elif self.variable_kernelsize == 'updated_epoch':
                if len(self.opt) <= 1:
                    idx = 1
                    self.kernelsize = int(idx * 2 + 1)
                else:
                    if self.variable_updated_flag == 1:
                        present_val = np.mean(self.opt['GT_logit'])
                    elif self.variable_updated_flag == 2:
                        present_val = np.mean(self.opt['predicted_logit'])
                    elif self.variable_updated_flag == 3:
                        present_val = np.mean(self.opt['acc'])

                    if present_val > self.variable_th:
                        changed_flag = True
                    else:
                        changed_flag = False

                    if changed_flag:

                        print('kernelsize is changed {}, step: {}'.format(self.kernelsize, self.step_base_cnt))
                        self.kernelsize = opt['previous_kernelsize'] + 2
                    else:
                        self.kernelsize = opt['previous_kernelsize']

                    if self.kernelsize > self.variable_limit:
                        self.kernelsize = opt['previous_kernelsize']
            elif self.variable_kernelsize == 'increased_epoch':
                idx = math.ceil(self.step_base_cnt * self.variable_times / self.n_tgt)
                if idx < 1:
                    idx = 1
                if idx > self.variable_times:
                    idx = self.variable_times # idx will be [1,self.variable_times]
                self.kernelsize = int(idx * 2 + 1)
            elif self.variable_kernelsize == 'decreased_epoch':
                idx = math.ceil(self.step_base_cnt * self.variable_times / self.n_tgt)
                if idx < 1:
                    idx = 1
                if idx > self.variable_times:
                    idx = self.variable_times
                idx = self.variable_times - idx + 1
                self.kernelsize = int(idx * 2 + 1)

            if self.encoder_model != 'none':
                if self.encoder_same: # re-initialization

                    if self.encoder_num == 1:
                        self.encoder.append(self.conv_flag(self.im_dim, self.norm_dim, self.bias, self.encoder_model, self.encoder_style))
                    elif self.encoder_num >= 2:
                        self.encoder.append(self.conv_flag(self.im_dim, self.conv_dim, self.bias, self.encoder_model, self.encoder_style))
                        for _ in range(self.encoder_num - 2):
                            self.encoder.append(self.conv_flag(self.conv_dim, self.conv_dim, self.bias, self.encoder_model, self.encoder_style))
                        self.encoder.append(self.conv_flag(self.conv_dim, self.norm_dim, self.bias, self.encoder_model, self.encoder_style))

                    # if 'deform' in self.encoder_model:
                    #     self.deform_weight = []
                    #     for i in range(self.encoder_num):
                    #         self.deform_weight.append(self.deform_weight_flag(self.encoder[i], init=True))

            if 'deform' in self.encoder_model:
                self.change_deform_offset = False
                if not self.encoder_same and ('multi' in self.encoder_model) and (not self.multiconv_same): # If kernel is changed.
                    self.change_deform_offset = True
                if not self.deform_same:
                    self.change_deform_offset = True

                if not self.change_deform_offset: # re-initialization
                    # print('deform_offset [init]')
                    self.deform_offset = []
                    for i in range(self.encoder_num):
                        self.deform_offset.append(self.deform_offset_flag(init=True))

            if self.encoder_act != 'none':
                for _ in range(self.encoder_num):
                    if self.encoder_act == 'sigmoid':
                        self.encoder_activation.append(torch.nn.Sigmoid())
                    elif self.encoder_act == 'tanh':
                        self.encoder_activation.append(torch.nn.Tanh())
                    elif self.encoder_act == 'hardtanh':
                        self.encoder_activation.append(torch.nn.Hardtanh())
                    elif self.encoder_act == 'relu':
                        self.encoder_activation.append(torch.nn.ReLU())
                    elif self.encoder_act == 'lrelu':
                        self.encoder_activation.append(torch.nn.LeakyReLU(negative_slope=0.01))

            if self.norm_model == 'adain':
                self.norm.append(AdaIN2d(self.zdim, self.norm_dim))
            elif self.norm_model == 'self_in':
                self.norm.append(selfIN2d(self.norm_dim))
            elif self.norm_model == 'in': # instance norm
                self.norm.append(IN2d(self.norm_dim))
            elif 'adain_' in self.norm_model == 'adain_normal' or self.norm_model == 'adain_uniform' or self.norm_model == 'adain_lognormal':
                if self.norm_model == 'adain_normal':
                    flag_gamma = 'normal'
                    flag_beta = 'normal'
                elif self.norm_model == 'adain_uniform':
                    flag_gamma = 'uniform'
                    flag_beta = 'uniform'
                elif self.norm_model == 'adain_lognormal':
                    flag_gamma = 'lognormal'
                    flag_beta = 'normal'
                elif self.norm_model == 'adain_uniform_gamma':
                    flag_gamma = 'uniform'
                    flag_beta = 'normal'
                elif self.norm_model == 'adain_uniform_beta':
                    flag_gamma = 'normal'
                    flag_beta = 'uniform'
                elif self.norm_model == 'adain_lognormal_both':
                    flag_gamma = 'lognormal'
                    flag_beta = 'lognormal'
                elif self.norm_model == 'adain_lognormal_uniform':
                    flag_gamma = 'lognormal'
                    flag_beta = 'uniform'

                norm_opt = {}
                norm_opt['flag_gamma'] = flag_gamma
                norm_opt['flag_beta'] = flag_beta
                norm_opt['num_features'] = self.norm_dim
                norm_opt['style'] = self.norm_style
                norm_opt['same'] = self.norm_same
                norm_opt['batch_size'] = self.batch_size
                norm_opt['num_classes'] = self.num_classes
                norm_opt['inv_prop'] = self.norm_inv_prop
                norm_opt['std1'] = self.norm_std1
                norm_opt['std2'] = self.norm_std2
                norm_opt['v_std1'] = self.norm_variable_std1
                norm_opt['v_std2'] = self.norm_variable_std2
                norm_opt['norm_true'] = self.norm_true
                norm_opt['not_used']=self.norm_not_used
                self.norm.append(GeneralizedAdaIN(norm_opt))
            #     self.norm.append(GeneralizedAdaIN(flag = 'normal', num_features = self.norm_dim, style = self.norm_style, same=self.norm_same, batch_size=self.batch_size, num_classes = self.num_classes, inv_prop = self.norm_inv_prop, std1 = self.norm_std1, std2 = self.norm_std2, v_std1 = self.norm_variable_std1, v_std2 = self.norm_variable_std2, norm_true = self.norm_true, not_used=self.norm_not_used))
            # elif self.norm_model == 'adain_uniform': # gamma: uniform, beta: uniform
            #     self.norm.append(GeneralizedAdaIN(flag = 'uniform', num_features = self.norm_dim, style = self.norm_style, same=self.norm_same, batch_size=self.batch_size, num_classes = self.num_classes, inv_prop = self.norm_inv_prop, std1 = self.norm_std1, std2 = self.norm_std2, v_std1 = self.norm_variable_std1, v_std2 = self.norm_variable_std2, norm_true = self.norm_true, not_used=self.norm_not_used))
            # elif self.norm_model == 'adain_lognormal': # gamma: lognormal, beta: normal
            #     self.norm.append(GeneralizedAdaIN(flag = 'lognormal', num_features = self.norm_dim, style = self.norm_style, same=self.norm_same, batch_size=self.batch_size, num_classes = self.num_classes, inv_prop = self.norm_inv_prop, std1 = self.norm_std1, std2 = self.norm_std2, v_std1 = self.norm_variable_std1, v_std2 = self.norm_variable_std2, norm_true = self.norm_true, not_used=self.norm_not_used))

            if self.decoder_model != 'none':
                if self.decoder_num == 1:
                    self.decoder.append(self.conv_flag(self.norm_dim, self.im_dim, self.bias, self.decoder_model, self.decoder_style))
                elif self.decoder_num >= 2:
                    self.decoder.append(self.conv_flag(self.norm_dim, self.conv_dim, self.bias, self.decoder_model, self.decoder_style))
                    for _ in range(self.decoder_num - 2): # if self.decoder_num == 3, cycle 1 time
                        self.decoder.append(self.conv_flag(self.conv_dim, self.conv_dim, self.bias, self.decoder_model, self.decoder_style))
                    self.decoder.append(self.conv_flag(self.conv_dim, self.im_dim, self.bias, self.decoder_model, self.decoder_style))

            if self.decoder_act != 'none':
                for _ in range(self.decoder_num - 1):
                    if self.decoder_act == 'sigmoid':
                        self.decoder_activation.append(torch.nn.Sigmoid())
                    elif self.decoder_act == 'tanh':
                        self.decoder_activation.append(torch.nn.Tanh())
                    elif self.decoder_act == 'hardtanh':
                        self.decoder_activation.append(torch.nn.Hardtanh())
                    elif self.decoder_act == 'relu':
                        self.decoder_activation.append(torch.nn.ReLU())
                    elif self.decoder_act == 'lrelu':
                        self.decoder_activation.append(torch.nn.LeakyReLU(negative_slope=0.01))
            if self.final_act != 'none':
                if self.final_act == 'sigmoid':
                    self.decoder_activation.append(torch.nn.Sigmoid())
                elif self.final_act == 'tanh':
                    self.decoder_activation.append(torch.nn.Tanh())
                elif self.final_act == 'hardtanh':
                    self.decoder_activation.append(torch.nn.Hardtanh())
                elif self.final_act == 'relu':
                    self.decoder_activation.append(torch.nn.ReLU())
                elif self.final_act == 'lrelu':
                    self.decoder_activation.append(torch.nn.LeakyReLU(negative_slope=0.01))

        self.init_norm_module = nn.Sequential(*self.init_norm)
        self.encoder_module = nn.Sequential(*self.encoder)
        self.encoder_activation_module = nn.Sequential(*self.encoder_activation)
        self.norm_module = nn.Sequential(*self.norm)
        self.decoder_module = nn.Sequential(*self.decoder)
        self.decoder_activation_module = nn.Sequential(*self.decoder_activation)

        if first_flag:
            if len(self.init_norm_module) > 0:
                print('*-'*10 + ' [encoder] '+'-*'*10)
                print(self.init_norm_module)
            if len(self.encoder_module) > 0:
                print('*-'*10 + ' [encoder] '+'-*'*10)
                print(self.encoder_module)
            if len(self.encoder_activation_module) > 0:
                print('*-'*10 + ' [encoder ACT] '+'-*'*10)
                print(self.encoder_activation_module)
            if len(self.norm_module) > 0:
                print('*-'*10 + ' [norm] '+'-*'*10)
                print(self.norm_module)
            if len(self.decoder_module) > 0:
                print('*-'*10 + ' [decoder] '+'-*'*10)
                print(self.decoder_module)
            if len(self.decoder_activation_module) > 0:
                print('*-'*10 + ' [decoder ACT] '+'-*'*10)
                print(self.decoder_activation_module)

    def conv_flag(self, dim1, dim2, bias, flag, style):

        if 'deform' in self.encoder_model:
            if style == 'channel':
                output = torchvision.ops.DeformConv2d(dim1, dim2, self.kernelsize, bias=bias, groups=3,dilation=self.dilation)
            elif style == 'instance':
                output = torchvision.ops.DeformConv2d(dim1*self.batch_size, dim2*self.batch_size, self.kernelsize, bias=bias, groups=self.batch_size,dilation=self.dilation)
            elif style == 'class':
                output = torchvision.ops.DeformConv2d(dim1*self.batch_size, dim2*self.batch_size, self.kernelsize, bias=bias, groups=self.batch_size,dilation=self.dilation)
            elif style == 'old_class':
                output = []
                for _ in range(self.num_classes):
                    output.append(torchvision.ops.DeformConv2d(dim1, dim2, self.kernelsize, bias=bias,dilation=self.dilation))
                output = nn.Sequential(*output)
            else:
                output = torchvision.ops.DeformConv2d(dim1, dim2, self.kernelsize, bias=bias,dilation=self.dilation)
            self.m_pad = torch.nn.ReplicationPad2d(((self.kernelsize-1)//2, (self.kernelsize-1)//2, (self.kernelsize-1)//2, (self.kernelsize-1)//2)).cuda()
        else:
            if style == 'channel':
                output = nn.Conv2d(dim1, dim2, self.kernelsize, padding=(self.kernelsize-1)//2, bias=bias, padding_mode='replicate', groups=3)
            elif style == 'instance':
                output = nn.Conv2d(dim1*self.batch_size, dim2*self.batch_size, self.kernelsize, padding=(self.kernelsize-1)//2, bias=bias, padding_mode='replicate', groups=self.batch_size)
            elif style == 'class':
                output = nn.Conv2d(dim1*self.batch_size, dim2*self.batch_size, self.kernelsize, padding=(self.kernelsize-1)//2, bias=bias, padding_mode='replicate', groups=self.batch_size)
            elif style == 'old_class':
                output = []
                for _ in range(self.num_classes):
                    output.append(nn.Conv2d(dim1, dim2, self.kernelsize, padding=(self.kernelsize-1)//2, bias=bias, padding_mode='replicate'))
                output = nn.Sequential(*output)
            else:
                output = nn.Conv2d(dim1, dim2, self.kernelsize, padding=(self.kernelsize-1)//2, bias=bias, padding_mode='replicate')

        return output

    def compute_offset(self, k_size1, k_size2, x_size1, x_size2):

        offset = torch.zeros(self.batch_size, 2*k_size1*k_size2, x_size1, x_size2).cuda()

        if self.deform_global:
            x1 = 1
            x2 = 1
        else:
            x1 = x_size1
            x2 = x_size2

        if self.deform_style == 'batch':
            b = 1
        elif self.deform_style == 'class':
            b = self.num_classes
        else:
            b = self.batch_size

        if self.deform_variable == 'increased_epoch':
            deform_offset_range = self.step_base_cnt / self.n_tgt * self.deform_variable_range
        elif self.deform_variable == 'rand_uniform':
            deform_offset_range = self.deform_variable_range * torch.rand(1).item()
            # print('deform:{}'.format(deform_offset_range))
        elif self.deform_variable == 'none':
            deform_offset_range = self.deform_variable_range

        rand_size = [b, 2*k_size1*k_size2, x1, x2]
        if self.deform_offset_style == 'normal': # [std:offset_range]
            if self.deform_style == 'class':
                offset[0:b,:,:,:] = torch.normal(mean = 0, std = deform_offset_range * torch.ones(rand_size)).cuda() # normal distribution
            else:
                offset[:,:,:,:] = torch.normal(mean = 0, std = deform_offset_range * torch.ones(rand_size)).cuda() # normal distribution
            if self.deform_offset_clamp:
                offset[offset>1.0] = 1.0 # clamp
                offset[offset<-1.0] = 1.0 # clamp
        elif self.deform_offset_style == 'uniform': # [-offset_range, offset_range]
            if self.deform_style == 'class':
                offset[0:b,:,:,:] = torch.rand(rand_size).cuda()
            else:
                offset[:,:,:,:] = torch.rand(rand_size).cuda()
            offset = deform_offset_range * (offset * 2.0 - 1.0) # for uniform

        if self.encoder_style == 'class' or self.encoder_style == 'instance':
            offset = torch.reshape(offset, (1, offset.shape[0]*offset.shape[1], offset.shape[2], offset.shape[3]))

        max_offset = torch.abs(torch.max(offset)).item()
        min_offset = torch.abs(torch.min(offset)).item()

        # for solve padding problem
        for k in range(min(math.ceil(max([max_offset, min_offset])), offset.size(2))):
            offset[:, :, k, :] = 0.0
            offset[:, :, -k-1, :] = 0.0
            offset[:, :, :, k] = 0.0
            offset[:, :, :, -k-1] = 0.0

        return offset

    # def deform_weight_flag(self, output, init=False):

    #     deform_weight = output.weight.data.cuda()
    #     self.m_pad = torch.nn.ReplicationPad2d(((self.kernelsize-1)//2, (self.kernelsize-1)//2, (self.kernelsize-1)//2, (self.kernelsize-1)//2)).cuda()

    #     return deform_weight

    def deform_offset_flag(self, init=False):
        # offset
        x_size1 = self.im_size[0]
        x_size2 = self.im_size[1]
        deform_offset = self.compute_offset(self.kernelsize, self.kernelsize, x_size1, x_size2)
        # if not init:
        deform_offset = deform_offset.cuda()

        return deform_offset

    def forward_noise(self, x, y, flag_ori = False):

        if self.noise_model == 'grf':

            if (self.noise_first and flag_ori) or self.noise_first == False:
                if self.noise_style == 'batch':
                    n = 1
                elif self.noise_style == 'class':
                    n = self.num_classes
                elif self.noise_style == 'instance':
                    n = self.batch_size
                x_noise_all = []

                for _ in range(n):
                    x_noise_channel = []
                    for _ in range(self.noise_channel):
                        x_noise_channel.append(torch.Tensor(grf.gaussian_random_field(alpha=torch.rand(1).item()*self.noise_alpha, size=x.shape[-1])).cuda())
                    x_noise_channel_all = torch.stack(x_noise_channel)
                    if torch.rand(1) > self.noise_prop:
                        x_noise_channel_all[:] = 0.0
                    x_noise_all.append(x_noise_channel_all)
                x_noise = torch.stack(x_noise_all)/self.noise_scale

                if self.noise_style == 'class':
                    xx = torch.zeros_like(x)
                    for y_local in range(self.num_classes):
                        idx = [j for j, k in enumerate(y) if k == y_local]
                        if len(idx) > 0:
                            for idx_local in idx:
                                xx[idx_local] = x_noise[y_local,:,:,:]

                    x = x + xx.detach()
                else:
                    x = x + x_noise.detach()
        return x

    def forward_init_norm(self, x, y, debug = False, w=[], b=[]):

        for i in range(len(self.init_norm_module)):
            if self.init_norm_model == 'self_in':
                if len(w) == 0:
                    x = self.init_norm_module[i](x, w=self.saved_w, b=self.saved_b)
                else:
                    x = self.init_norm_module[i](x, w=w, b=b)
            else:
                x = self.init_norm_module[i](x, y=y)
            if debug:
                self.debug_name.append('init_norm_'+self.norm_model+str(i))
                self.debug_output.append(x.clone().detach())

        return x

    def forward_encoder(self, x, y, debug = False):

        if self.encoder_model != 'none':
            if not self.encoder_same: # re-initialization

                if self.variable_kernelsize == 'none':
                    if ('multi' in self.encoder_model) and (not self.multiconv_same):
                        kernelsize = self.multiconv_kernelsize
                        try:
                            kernelsize = [int(i) for i in kernelsize.split('_')]
                        except:
                            kernelsize = [int(i[1:]) for i in kernelsize.split('_')]
                        kernelsize = kernelsize[np.random.randint(len(kernelsize))]
                        self.kernelsize = kernelsize

                self.encoder = []
                if self.encoder_num == 1:
                    self.encoder.append(self.conv_flag(self.im_dim, self.norm_dim, self.bias, self.encoder_model, self.encoder_style))
                elif self.encoder_num >= 2:
                    self.encoder.append(self.conv_flag(self.im_dim, self.conv_dim, self.bias, self.encoder_model, self.encoder_style))
                    for _ in range(self.encoder_num - 2):
                        self.encoder.append(self.conv_flag(self.conv_dim, self.conv_dim, self.bias, self.encoder_model, self.encoder_style))
                    self.encoder.append(self.conv_flag(self.conv_dim, self.norm_dim, self.bias, self.encoder_model, self.encoder_style))
                self.encoder_module = nn.Sequential(*self.encoder).cuda()

                # if 'deform' in self.encoder_model:
                #     self.deform_weight = []
                #     for i in range(self.encoder_num):
                #         self.deform_weight.append(self.deform_weight_flag(self.encoder[i], init=False))

        if 'deform' in self.encoder_model:
            if self.change_deform_offset:
                # print('deform_offset [changed]')
                self.deform_offset = []
                for i in range(self.encoder_num):
                    self.deform_offset.append(self.deform_offset_flag(init=False))

        # Encoder
        for i in range(len(self.encoder_module)):
            # enc module
            x_shape = x.shape
            if 'deform' in self.encoder_model:
                x = self.m_pad(x)

            if self.encoder_style == 'old_class': # y
                xx = torch.zeros_like(x)
                for y_local in range(self.num_classes):
                    idx = [j for j, k in enumerate(y) if k == y_local]
                    if len(idx) > 0:
                        if 'deform' in self.encoder_model:
                            xx[idx] = self.encoder_module[i][y_local](x[idx], self.deform_offset[i])
                        else:
                            xx[idx] = self.encoder_module[i][y_local](x[idx])
                x = xx.clone().detach()

            elif self.encoder_style == 'class': # y
                if 'deform' in self.encoder_model:
                    x = torch.reshape(x, (1, x_shape[0]*x_shape[1], x.shape[2], x.shape[3]))
                    # weight extract (assumption: 3, 3, 3, 3, 3 )
                    classwise_weight = []
                    for j in range(self.num_classes):
                        # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                        j_idx = [j*x_shape[1]+z for z in range(x_shape[1])]
                        classwise_weight.append(self.encoder_module[i].weight.data[j_idx].clone().detach())
                    # weight allocation
                    for j in range(self.batch_size):
                        y_local = y[j]
                        # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                        j_idx = [j*x_shape[1]+z for z in range(x_shape[1])]
                        self.encoder_module[i].weight.data[j_idx] = classwise_weight[y_local]

                    if self.deform_style == 'class':
                        deform_offset = torch.zeros_like(self.deform_offset[i][0])
                        deform_offset2 = self.deform_offset[i][0].clone().detach()
                        # weight extract (assumption: 2*kernel*kernel )
                        classwise_offset = []
                        for j in range(self.num_classes):
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            classwise_offset.append(deform_offset2[j_idx])
                        # weight allocation
                        for j in range(self.batch_size):
                            y_local = y[j]
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            deform_offset[j_idx] = classwise_offset[y_local]
                        x = self.encoder_module[i](x, deform_offset.unsqueeze(0))
                    else:
                        x = self.encoder_module[i](x, self.deform_offset[i])

                    x = torch.reshape(x, (x_shape[0], x_shape[1], x.shape[2], x.shape[3]))
                else:
                    x = torch.reshape(x, (1, x_shape[0]*x_shape[1], x.shape[2], x.shape[3]))
                    # weight extract (assumption: 3, 3, 3, 3, 3 )
                    classwise_weight = []
                    for j in range(self.num_classes):
                        # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                        j_idx = [j*x_shape[1]+z for z in range(x_shape[1])]
                        classwise_weight.append(self.encoder_module[i].weight.data[j_idx].clone().detach())
                    # weight allocation
                    for j in range(self.batch_size):
                        y_local = y[j]
                        # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                        j_idx = [j*x_shape[1]+z for z in range(x_shape[1])]
                        self.encoder_module[i].weight.data[j_idx] = classwise_weight[y_local]

                    x = self.encoder_module[i](x)
                    x = torch.reshape(x, (x_shape[0], x_shape[1], x.shape[2], x.shape[3]))
            elif self.encoder_style == 'instance':
                if 'deform' in self.encoder_model:
                    x = torch.reshape(x, (1, x_shape[0]*x_shape[1], x.shape[2], x.shape[3]))
                    if self.deform_style == 'class':
                        deform_offset = torch.zeros_like(self.deform_offset[i][0])
                        deform_offset2 = self.deform_offset[i][0].clone().detach()
                        # weight extract (assumption: 2*kernel*kernel )
                        classwise_offset = []
                        for j in range(self.num_classes):
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            classwise_offset.append(deform_offset2[j_idx])
                        # weight allocation
                        for j in range(self.batch_size):
                            y_local = y[j]
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            deform_offset[j_idx] = classwise_offset[y_local]
                        x = self.encoder_module[i](x, deform_offset.unsqueeze(0))
                    else:
                        x = self.encoder_module[i](x, self.deform_offset[i])
                    x = torch.reshape(x, (x_shape[0], x_shape[1], x.shape[2], x.shape[3]))
                else:
                    x = torch.reshape(x, (1, x_shape[0]*x_shape[1], x.shape[2], x.shape[3]))
                    x = self.encoder_module[i](x)
                    x = torch.reshape(x, (x_shape[0], x_shape[1], x.shape[2], x.shape[3]))
            else:
                if 'deform' in self.encoder_model:
                    if self.deform_style == 'class':

                        # added (due to different size)
                        offset = torch.reshape(self.deform_offset[i], (1, self.deform_offset[i].shape[0]*self.deform_offset[i].shape[1], self.deform_offset[i].shape[2], self.deform_offset[i].shape[3]))
                        deform_offset = torch.zeros_like(offset[0])
                        deform_offset2 = offset[0].clone().detach()
                        # weight extract (assumption: 2*kernel*kernel )
                        classwise_offset = []
                        for j in range(self.num_classes):
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            classwise_offset.append(deform_offset2[j_idx])
                        # weight allocation
                        for j in range(self.batch_size):
                            y_local = y[j]
                            # j_idx = [z*x_shape[0]+j for z in range(x_shape[1])]
                            j_idx = [j*2*self.kernelsize*self.kernelsize+z for z in range(2*self.kernelsize*self.kernelsize)]
                            deform_offset[j_idx] = classwise_offset[y_local]
                        # added (due to different size)
                        deform_offset = torch.reshape(deform_offset, self.deform_offset[i].shape)
                        x = self.encoder_module[i](x, deform_offset)
                    else:
                        x = self.encoder_module[i](x, self.deform_offset[i])
                else:
                    x = self.encoder_module[i](x)

            if debug:
                self.debug_name.append('enc_'+self.encoder_model+str(i))
                self.debug_output.append(x.clone().detach())

            # enc activation
            if len(self.encoder_activation_module) > 0:
                x = self.encoder_activation_module[i](x)
                if debug:
                    self.debug_name.append('enc_'+self.encoder_model+str(i)+'_'+self.encoder_act)
                    self.debug_output.append(x.clone().detach())

        return x

    def forward_norm(self, x, y, debug = False, w=[], b=[], rand=False):

        # Norm
        for i in range(len(self.norm_module)):
            if self.norm_model == 'adain':
                if rand:
                    z = torch.randn(len(x), self.zdim).cuda()
                    x = self.norm_module[i](x, z)
            elif self.norm_model == 'self_in':
                if len(w) == 0:
                    x = self.norm_module[i](x, w=self.saved_w, b=self.saved_b)
                else:
                    x = self.norm_module[i](x, w=w, b=b)
            else:
                x = self.norm_module[i](x, y=y)
            if debug:
                self.debug_name.append('norm_'+self.norm_model+str(i))
                self.debug_output.append(x.clone().detach())

        return x

    def forward_decoder(self, x, y, debug = False):

        # Decoder
        if len(self.decoder_activation_module) > len(self.decoder_module):
            NotImplementedError

        if len(self.decoder_module) > 0:
            for i in range(len(self.decoder_module)):

                # dec module
                x = self.decoder_module[i](x)
                if debug:
                    self.debug_name.append('dec_'+self.decoder_model+str(i))
                    self.debug_output.append(x.clone().detach())

                # dec activation
                if len(self.decoder_activation_module)-1 >= i:
                    x = self.decoder_activation_module[i](x)
                    if debug:
                        self.debug_name.append('enc_'+self.decoder_model+str(i)+'_'+self.decoder_act)
                        self.debug_output.append(x.clone().detach())

        else: # no decoder (check decoder_Activation)
            if len(self.decoder_activation_module) == 1: # only final activation
                x = self.decoder_activation_module[0](x)
            elif len(self.decoder_activation_module) > 1:
                NotImplementedError

        return x

    def forward_activation(self, x):

        if self.final_act != 'none':
            if self.final_act == 'sigmoid' and self.input_normalize:
                x = (x - 0.5) / 0.5
            elif self.final_act == 'tanh' and not self.input_normalize:
                x = x * 0.5 + 0.5

        if self.final_norm == 'clamp':
            x[x>1.0] = 1.0
            if self.input_normalize: # [-1,1]
                x[x<-1.0] = -1.0
            else:
                x[x<0] = 0.0
        elif self.final_norm == 'each_minmax':
            for j in range(len(x)):
                x[j] = (x[j] - x[j].min()) / (x[j].max() - x[j].min() + 0.00000001)
                if self.input_normalize:
                    x[j] = (x[j] - 0.5) / 0.5
        elif self.final_norm == 'batch_minmax':
            x = (x-x.min()) / (x.max()-x.min() + 0.00000001)
            if self.input_normalize:
                x = (x-0.5)/0.5

        return x

    def forward(self, x, rand=False, debug=False, w=[], b=[], y=[], flag_ori=False):
        if len(w) > 0:
            self.saved_w = w
            self.saved_b = b

        x = self.forward_noise(x, y, flag_ori)
        x = self.forward_init_norm(x, y, debug, w, b)
        x = self.forward_encoder(x, y, debug)
        x = self.forward_norm(x, y, debug, w, b, rand)
        x = self.forward_decoder(x, y, debug)
        x = self.forward_activation(x)

        if debug:
            return x, self.debug_output, self.debug_name
        else:
            return x, None, None

class stnGenerator(nn.Module):
    ''' 仿射变换 '''
    def __init__(self, zdim=10, imsize=[32,32], mode=None):
        super().__init__()
        self.mode = mode
        self.zdim = zdim

        self.mapz = nn.Linear(zdim, imsize[0]*imsize[1])
        if imsize == [32,32]:
            self.loc = nn.Sequential(
                    nn.Conv2d( 4,  16, 5), nn.MaxPool2d(2), nn.ReLU(),
                    nn.Conv2d( 16, 32, 5), nn.MaxPool2d(2), nn.ReLU(),)
            self.fc_loc = nn.Sequential(
                    nn.Linear(32*5*5, 32), nn.ReLU(),
                    nn.Linear(32, 6))
        # init the weight
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([1,0,0,0,1,0]))
    def forward(self, x, rand, return_H=False):
        if rand:
            z = torch.randn(len(x), self.zdim).cuda()
        z = self.mapz(z).view(len(x), 1, x.size(2), x.size(3))
        loc = self.loc(torch.cat([x, z], dim=1)) # [N, -1]
        loc = loc.view(len(loc), -1)
        H = self.fc_loc(loc)
        H = H.view(len(H), 2, 3)
        if self.mode == 'translate':
            H[:,0,0] = 1
            H[:,0,1] = 0
            H[:,1,0] = 0
            H[:,1,1] = 1
        grid = F.affine_grid(H, x.size())
        x = F.grid_sample(x, grid)
        if return_H:
            return x, H
        else:
            return x