from . import Utils
from . import DataMisfit
import numpy as np
import matplotlib.pyplot as plt
import warnings
from .PF import Magnetics
from . import Regularization
from . import Mesh
from . import ObjectiveFunction
import json

class InversionDirective:
    """InversionDirective"""

    debug = False  #: Print debugging information
    _regPair = [
        Regularization.BaseComboRegularization,
        Regularization.BaseRegularization,
        ObjectiveFunction.ComboObjectiveFunction,
    ]
    _dmisfitPair = [DataMisfit.BaseDataMisfit, ObjectiveFunction.ComboObjectiveFunction]

    def __init__(self, **kwargs):
        Utils.setKwargs(self, **kwargs)

    @property
    def inversion(self):
        """This is the inversion of the InversionDirective instance."""
        return getattr(self, "_inversion", None)

    @inversion.setter
    def inversion(self, i):
        # if getattr(self, '_inversion', None) is not None:
        #     warnings.warn(
        #         'InversionDirective {0!s} has switched to a new inversion.'
        #         .format(self.__class__.__name__)
        #     )
        self._inversion = i

    @property
    def invProb(self):
        return self.inversion.invProb

    @property
    def opt(self):
        return self.invProb.opt

    @property
    def reg(self):
        if getattr(self, "_reg", None) is None:
            self.reg = self.invProb.reg  # go through the setter
        return self._reg

    @reg.setter
    def reg(self, value):
        assert any(
            [isinstance(value, regtype) for regtype in self._regPair]
        ), "Regularization must be in {}, not {}".format(self._regPair, type(value))

        if isinstance(value, Regularization.BaseComboRegularization):
            value = 1 * value  # turn it into a combo objective function
        self._reg = value

    @property
    def dmisfit(self):
        if getattr(self, "_dmisfit", None) is None:
            self.dmisfit = self.invProb.dmisfit  # go through the setter
        return self._dmisfit

    @dmisfit.setter
    def dmisfit(self, value):

        assert any(
            [isinstance(value, dmisfittype) for dmisfittype in self._dmisfitPair]
        ), "Regularization must be in {}, not {}".format(self._dmisfitPair, type(value))

        if not isinstance(value, ObjectiveFunction.ComboObjectiveFunction):
            value = 1 * value  # turn it into a combo objective function
        self._dmisfit = value

    @property
    def survey(self):
        """
           Assuming that dmisfit is always a ComboObjectiveFunction,
           return a list of surveys for each dmisfit [survey1, survey2, ... ]
        """
        return [objfcts.survey for objfcts in self.dmisfit.objfcts]

    @property
    def prob(self):
        """
           Assuming that dmisfit is always a ComboObjectiveFunction,
           return a list of problems for each dmisfit [prob1, prob2, ...]
        """
        return [objfcts.prob for objfcts in self.dmisfit.objfcts]

    def initialize(self):
        pass

    def endIter(self):
        pass

    def finish(self):
        pass

    def validate(self, directiveList=None):
        return True

class DirectiveList:

    dList = None  #: The list of Directives

    def __init__(self, *directives, **kwargs):
        self.dList = []
        for d in directives:
            assert isinstance(
                d, InversionDirective
            ), "All directives must be InversionDirectives not {}".format(type(d))
            self.dList.append(d)
        Utils.setKwargs(self, **kwargs)

    @property
    def debug(self):
        return getattr(self, "_debug", False)

    @debug.setter
    def debug(self, value):
        for d in self.dList:
            d.debug = value
        self._debug = value

    @property
    def inversion(self):
        """This is the inversion of the InversionDirective instance."""
        return getattr(self, "_inversion", None)

    @inversion.setter
    def inversion(self, i):
        if self.inversion is i:
            return
        if getattr(self, "_inversion", None) is not None:
            warnings.warn(
                f"{self.__class__.__name__!s} has switched to a new inversion."
            )
        for d in self.dList:
            d.inversion = i
        self._inversion = i

    def call(self, ruleType):
        if self.dList is None:
            if self.debug:
                print("DirectiveList is None, no directives to call!")
            return

        directives = ["initialize", "endIter", "finish"]
        assert ruleType in directives, 'Directive type must be in ["{!s}"]'.format(
            '", "'.join(directives)
        )
        for r in self.dList:
            getattr(r, ruleType)()

    def validate(self):
        [directive.validate(self) for directive in self.dList]
        return True

class BetaEstimate_ByEig(InversionDirective):
    """BetaEstimate"""

    beta0 = None  #: The initial Beta (regularization parameter)
    beta0_ratio = 1e2  #: estimateBeta0 is used with this ratio

    def initialize(self):
        """
            The initial beta is calculated by comparing the estimated
            eigenvalues of JtJ and WtW.

            To estimate the eigenvector of **A**, we will use one iteration
            of the *Power Method*:

            .. math::

                \\mathbf{x_1 = A x_0}

            Given this (very course) approximation of the eigenvector, we can
            use the *Rayleigh quotient* to approximate the largest eigenvalue.

            .. math::

                \\lambda_0 = \\frac{\\mathbf{x^\\top A x}}{\\mathbf{x^\\top x}}

            We will approximate the largest eigenvalue for both JtJ and WtW,
            and use some ratio of the quotient to estimate beta0.

            .. math::

                \\beta_0 = \\gamma \\frac{\\mathbf{x^\\top J^\\top J x}}{\\mathbf{x^\\top W^\\top W x}}

            :rtype: float
            :return: beta0
        """

        if self.debug:
            print("Calculating the beta0 parameter.")

        m = self.invProb.model
        f = self.invProb.getFields(m, store=True, deleteWarmstart=False)

        x0 = np.random.rand(m.shape[0])
        t = np.dot(x0, self.dmisfit.deriv2(m, x0, f=f))
        b = np.dot(x0, self.reg.deriv2(m, v=x0))

        self.beta0 = self.beta0_ratio * (t / b)

        self.invProb.beta = self.beta0

class BetaSchedule(InversionDirective):
    """BetaSchedule"""

    coolingFactor = 8.0
    coolingRate = 3

    def endIter(self):
        if self.opt.iter > 0 and self.opt.iter % self.coolingRate == 0:
            if self.debug:
                print(
                    "BetaSchedule is cooling Beta. Iteration: {:d}".format(
                        self.opt.iter
                    )
                )
            self.invProb.beta /= self.coolingFactor

class TargetMisfit(InversionDirective):

    chifact = 1.0
    phi_d_star = None

    @property
    def target(self):
        if getattr(self, "_target", None) is None:
            # the factor of 0.5 is because we do phid = 0.5*|| dpred - dobs||^2
            if self.phi_d_star is None:

                nD = 0
                for survey in self.survey:
                    nD += survey.nD

                self.phi_d_star = 0.5 * nD

            self._target = self.chifact * self.phi_d_star
        return self._target

    @target.setter
    def target(self, val):
        self._target = val

    def endIter(self):
        if self.invProb.phi_d < self.target:
            self.opt.stopNextIteration = True

class SaveEveryIteration(InversionDirective):
    @property
    def name(self):
        if getattr(self, "_name", None) is None:
            self._name = "InversionModel"
        return self._name

    @name.setter
    def name(self, value):
        self._name = value

    @property
    def fileName(self):
        if getattr(self, "_fileName", None) is None:
            from datetime import datetime

            self._fileName = "{!s}-{!s}".format(
                self.name, datetime.now().strftime("%Y-%m-%d-%H-%M")
            )
        return self._fileName

    @fileName.setter
    def fileName(self, value):
        self._fileName = value

class SaveModelEveryIteration(SaveEveryIteration):
    """SaveModelEveryIteration"""

    def initialize(self):
        print(
            "SimPEG.SaveModelEveryIteration will save your models as: '###-{!s}.npy'".format(
                self.fileName
            )
        )

    def endIter(self):
        np.save(f"{self.opt.iter:03d}-{self.fileName!s}", self.opt.xc)

class SaveUBCModelEveryIteration(SaveEveryIteration):
    """SaveModelEveryIteration"""

    replace = True
    saveComp = True
    mapping = None
    vector = False
    mesh = None

    def initialize(self):

        if getattr(self, "mapping", None) is None:
            return self.mapPair()
        print(
            "SimPEG.SaveModelEveryIteration will save your models"
            + f" in UBC format as: '###-{self.fileName!s}.mod'"
        )

    def endIter(self):

        if not self.replace:
            fileName = self.fileName + "Iter" + str(self.opt.iter)
        else:
            fileName = self.fileName

        count = -1
        for prob, survey, reg in zip(self.prob, self.survey, self.reg.objfcts):

            count += 1

            # if getattr(prob, 'mapping', None) is not None:
            #     xc = prob.mapping * self.opt.xc

            # else:
            xc = self.mapping * self.opt.xc

            # # Save predicted data
            # if len(self.prob) > 1:
            #     Magnetics.writeUBCobs(fileName + "Prob" + str(count) + '.pre', survey, survey.dpred(m=self.opt.xc))

            # else:
            #     Magnetics.writeUBCobs(fileName + '.pre', survey, survey.dpred(m=self.opt.xc))

            # Save model
            if not self.vector:

                if isinstance(self.mesh, Mesh.TreeMesh):
                    Mesh.TreeMesh.writeUBC(
                        self.mesh, fileName + ".msh", models={fileName + ".mod": xc}
                    )

                else:
                    Mesh.TensorMesh.writeModelUBC(self.mesh, fileName + ".mod", xc)
            else:

                nC = self.mesh.nC

                if prob.coordinate_system == "spherical":
                    vec_xyz = Utils.matutils.atp2xyz(
                        xc.reshape((int(len(xc) / 3), 3), order="F")
                    )
                    theta = xc[nC : 2 * nC]
                    phi = xc[2 * nC :]
                else:
                    vec_xyz = xc
                    atp = Utils.matutils.xyz2atp(
                        xc.reshape((int(len(xc) / 3), 3), order="F")
                    )
                    theta = atp[nC : 2 * nC]
                    phi = atp[2 * nC :]

                vec_x = vec_xyz[:nC]
                vec_y = vec_xyz[nC : 2 * nC]
                vec_z = vec_xyz[2 * nC :]

                vec = np.c_[vec_x, vec_y, vec_z]

                m_pst = Utils.matutils.xyz2pst(vec, self.survey[0].srcField.param)
                m_ind = m_pst.copy()
                m_ind[:, 1:] = 0.0
                m_ind = Utils.matutils.pst2xyz(m_ind, self.survey[0].srcField.param)

                m_rem = m_pst.copy()
                m_rem[:, 0] = 0.0
                m_rem = Utils.matutils.pst2xyz(m_rem, self.survey[0].srcField.param)

                if self.saveComp:
                    if isinstance(self.mesh, Mesh.TreeMesh):
                        Mesh.TreeMesh.writeUBC(
                            self.mesh,
                            fileName + ".msh",
                            models={
                                fileName + ".dip": (np.rad2deg(theta)),
                                fileName + ".azm": ((450 - np.rad2deg(phi)) % 360),
                                fileName + "_TOT.mod": np.sum(vec ** 2, axis=1) ** 0.5,
                                fileName
                                + "_IND.mod": np.sum(m_ind ** 2, axis=1) ** 0.5,
                                fileName
                                + "_REM.mod": np.sum(m_rem ** 2, axis=1) ** 0.5,
                            },
                        )

                        Utils.io_utils.writeVectorUBC(self.mesh, fileName + ".fld", vec)

                    else:
                        Mesh.TensorMesh.writeModelUBC(
                            self.mesh, fileName + ".dip", (np.rad2deg(theta))
                        )
                        Mesh.TensorMesh.writeModelUBC(
                            self.mesh, fileName + ".azm", (450 - np.rad2deg(phi)) % 360
                        )
                        Mesh.TensorMesh.writeModelUBC(
                            self.mesh,
                            fileName + "_TOT.mod",
                            np.sum(vec ** 2, axis=1) ** 0.5,
                        )
                        Mesh.TensorMesh.writeModelUBC(
                            self.mesh,
                            fileName + "_IND.mod",
                            np.sum(m_ind ** 2, axis=1) ** 0.5,
                        )
                        Mesh.TensorMesh.writeModelUBC(
                            self.mesh,
                            fileName + "_REM.mod",
                            np.sum(m_rem ** 2, axis=1) ** 0.5,
                        )
                        Utils.io_utils.writeVectorUBC(
                            self.mesh, fileName + "_VEC.fld", vec
                        )

class SaveUBCPredictedEveryIteration(SaveEveryIteration):
    """SaveModelEveryIteration"""

    replace = True
    format = "grav"
    fileName = "Predicted"
    residuals = True
    survey = None

    def initialize(self):

        self.invProb.evalFunction(self.invProb.model, return_g=False, return_H=False)

        residuals = self.survey.dobs - self.invProb.dpred

        if self.format == "grav":
            Utils.io_utils.writeUBCgravityObservations(
                self.fileName + "_Initial.pre", self.survey, self.invProb.dpred
            )

            if self.residuals:
                Utils.io_utils.writeUBCgravityObservations(
                    self.fileName + "_Initial_Residual.pre", self.survey, residuals
                )

        elif self.format in ["mag", "mvi", "mvis"]:
            Utils.io_utils.writeUBCmagneticsObservations(
                self.fileName + "_Initial.pre", self.survey, self.invProb.dpred
            )

            if self.residuals:
                Utils.io_utils.writeUBCmagneticsObservations(
                    self.fileName + "_Initial_Residual.pre", self.survey, residuals
                )

        print(
            "SimPEG.SavePredictedEveryIteration will save your predicted data"
            + f" in UBC format as: '###-{self.fileName!s}.mod'"
        )

    def endIter(self):

        if not self.replace:
            fileName = self.fileName + "Iter" + str(self.opt.iter)
        else:
            fileName = self.fileName

        residuals = self.survey.dobs - self.invProb.dpred

        if self.format == "grav":
            Utils.io_utils.writeUBCgravityObservations(
                fileName + ".pre", self.survey, self.invProb.dpred
            )

            if self.residuals:
                Utils.io_utils.writeUBCgravityObservations(
                    fileName + "_Residual.pre", self.survey, residuals
                )

        elif self.format in ["mag", "mvi", "mvis"]:
            Utils.io_utils.writeUBCmagneticsObservations(
                fileName + ".pre", self.survey, self.invProb.dpred
            )

            if self.residuals:
                Utils.io_utils.writeUBCmagneticsObservations(
                    fileName + "_Residual.pre", self.survey, residuals
                )

class SaveOutputEveryIteration(SaveEveryIteration):
    """SaveModelEveryIteration"""

    header = None
    save_txt = True
    beta = None
    phi_d = None
    phi_m = None
    phi_m_small = None
    phi_m_smooth_x = None
    phi_m_smooth_y = None
    phi_m_smooth_z = None
    phi = None

    def initialize(self):
        if self.save_txt is True:
            print(
                "SimPEG.SaveOutputEveryIteration will save your inversion "
                "progress as: '###-{!s}.txt'".format(self.fileName)
            )
            f = open(self.fileName + ".txt", "w")
            self.header = "  #     beta     phi_d     phi_m   phi_m_small     phi_m_smoomth_x     phi_m_smoomth_y     phi_m_smoomth_z      phi\n"
            f.write(self.header)
            f.close()

        self.beta = []
        self.phi_d = []
        self.phi_m = []
        self.phi_m_small = []
        self.phi_m_smooth_x = []
        self.phi_m_smooth_y = []
        self.phi_m_smooth_z = []
        self.phi = []

    def endIter(self):
        phi_s, phi_x, phi_y, phi_z = 0, 0, 0, 0
        for reg in self.reg.objfcts:

            # f_m = reg.objfcts[0].f_m
            # phi_s += np.sum(f_m**2./(f_m**2. + 1e-8)**(1-reg.objfcts[0].norm/2.))

            # f_m = reg.objfcts[1].f_m
            # phi_x += np.sum(f_m**2./(f_m**2. + 1e-8)**(1-reg.objfcts[1].norm/2.))
            phi_s += reg.objfcts[0](self.invProb.model) * reg.alpha_s
            phi_x += reg.objfcts[1](self.invProb.model) * reg.alpha_x

            if reg.regmesh.dim > 1:
                # f_m = reg.objfcts[2].f_m
                # phi_x += np.sum(f_m**2./(f_m**2. + 1e-8)**(1-reg.objfcts[2].norm/2.))
                phi_y += reg.objfcts[2](self.invProb.model) * reg.alpha_y

            if reg.regmesh.dim > 2:
                # f_m = reg.objfcts[3].f_m
                # phi_x += np.sum(f_m**2./(f_m**2. + 1e-8)**(1-reg.objfcts[3].norm/2.))

                phi_z += reg.objfcts[3](self.invProb.model) * reg.alpha_z

            # elif reg.regmesh.dim == 3:
            #     phi_y += (
            #         reg.objfcts[2](self.invProb.model) * reg.alpha_y
            #     )

        self.beta.append(self.invProb.beta)
        self.phi_d.append(self.invProb.phi_d)
        self.phi_m.append(self.invProb.phi_m)
        self.phi_m_small.append(phi_s)
        self.phi_m_smooth_x.append(phi_x)
        self.phi_m_smooth_y.append(phi_y)
        self.phi_m_smooth_z.append(phi_z)
        self.phi.append(self.opt.f)

        if self.save_txt:
            f = open(self.fileName + ".txt", "a")
            f.write(
                " {:3d} {:1.4e} {:1.4e} {:1.4e} {:1.4e} {:1.4e} "
                "{:1.4e}  {:1.4e}  {:1.4e}\n".format(
                    self.opt.iter,
                    self.beta[self.opt.iter - 1],
                    self.phi_d[self.opt.iter - 1],
                    self.phi_m[self.opt.iter - 1][0],
                    self.phi_m_small[self.opt.iter - 1],
                    self.phi_m_smooth_x[self.opt.iter - 1],
                    self.phi_m_smooth_y[self.opt.iter - 1],
                    self.phi_m_smooth_z[self.opt.iter - 1],
                    self.phi[self.opt.iter - 1][0],
                )
            )
            f.close()

    def load_results(self):
        results = np.loadtxt(self.fileName + ".txt", comments="#")
        self.beta = results[:, 1]
        self.phi_d = results[:, 2]
        self.phi_m = results[:, 3]
        self.phi_m_small = results[:, 4]
        self.phi_m_smooth_x = results[:, 5]
        self.phi_m_smooth_y = results[:, 6]
        self.phi_m_smooth_z = results[:, 7]

        if self.reg.regmesh.dim == 1:
            self.phi_m_smooth = self.phi_m_smooth_x.copy()
        elif self.reg.regmesh.dim == 2:
            self.phi_m_smooth = self.phi_m_smooth_x + self.phi_m_smooth_y
        elif self.reg.regmesh.dim == 3:
            self.phi_m_smooth = (
                self.phi_m_smooth_x + self.phi_m_smooth_y + self.phi_m_smooth_z
            )

        self.f = results[:, 7]

        self.target_misfit = self.invProb.dmisfit.prob.survey.nD / 2.0
        self.i_target = None

        if self.invProb.phi_d < self.target_misfit:
            i_target = 0
            while self.phi_d[i_target] > self.target_misfit:
                i_target += 1
            self.i_target = i_target

    def plot_misfit_curves(self, fname=None, plot_small_smooth=False):

        self.target_misfit = self.invProb.dmisfit.prob.survey.nD / 2.0
        self.i_target = None

        if self.invProb.phi_d < self.target_misfit:
            i_target = 0
            while self.phi_d[i_target] > self.target_misfit:
                i_target += 1
            self.i_target = i_target

        fig = plt.figure(figsize=(5, 2))
        ax = plt.subplot(111)
        ax_1 = ax.twinx()
        ax.semilogy(np.arange(len(self.phi_d)), self.phi_d, "k-", lw=2)
        ax_1.semilogy(np.arange(len(self.phi_d)), self.phi_m, "r", lw=2)
        if plot_small_smooth:
            ax_1.semilogy(np.arange(len(self.phi_d)), self.phi_m_small, "ro")
            ax_1.semilogy(np.arange(len(self.phi_d)), self.phi_m_smooth, "rx")
            ax_1.legend((r"$\phi_m$", "small", "smooth"), bbox_to_anchor=(1.5, 1.0))

        ax.plot(
            np.r_[ax.get_xlim()[0], ax.get_xlim()[1]],
            np.ones(2) * self.target_misfit,
            "k:",
        )
        ax.set_xlabel("Iteration")
        ax.set_ylabel(r"$\phi_d$")
        ax_1.set_ylabel(r"$\phi_m$", color="r")
        for tl in ax_1.get_yticklabels():
            tl.set_color("r")
        plt.show()

    def plot_tikhonov_curves(self, fname=None, dpi=200):

        self.target_misfit = self.invProb.dmisfit.prob.survey.nD / 2.0
        self.i_target = None

        if self.invProb.phi_d < self.target_misfit:
            i_target = 0
            while self.phi_d[i_target] > self.target_misfit:
                i_target += 1
            self.i_target = i_target

        fig = plt.figure(figsize=(5, 8))
        ax1 = plt.subplot(311)
        ax2 = plt.subplot(312)
        ax3 = plt.subplot(313)

        ax1.plot(self.beta, self.phi_d, "k-", lw=2, ms=4)
        ax1.set_xlim(np.hstack(self.beta).min(), np.hstack(self.beta).max())
        ax1.set_xlabel("$\\beta$", fontsize=14)
        ax1.set_ylabel(r"$\phi_d$", fontsize=14)

        ax2.plot(self.beta, self.phi_m, "k-", lw=2)
        ax2.set_xlim(np.hstack(self.beta).min(), np.hstack(self.beta).max())
        ax2.set_xlabel("$\\beta$", fontsize=14)
        ax2.set_ylabel(r"$\phi_m$", fontsize=14)

        ax3.plot(self.phi_m, self.phi_d, "k-", lw=2)
        ax3.set_xlim(np.hstack(self.phi_m).min(), np.hstack(self.phi_m).max())
        ax3.set_xlabel(r"$\phi_m$", fontsize=14)
        ax3.set_ylabel(r"$\phi_d$", fontsize=14)

        if self.i_target is not None:
            ax1.plot(self.beta[self.i_target], self.phi_d[self.i_target], "k*", ms=10)
            ax2.plot(self.beta[self.i_target], self.phi_m[self.i_target], "k*", ms=10)
            ax3.plot(self.phi_m[self.i_target], self.phi_d[self.i_target], "k*", ms=10)

        for ax in [ax1, ax2, ax3]:
            ax.set_xscale("log")
            ax.set_yscale("log")
        plt.tight_layout()
        plt.show()
        if fname is not None:
            fig.savefig(fname, dpi=dpi)

class SaveOutputDictEveryIteration(SaveEveryIteration):
    """
        Saves inversion parameters at every iteraion.
    """

    # Initialize the output dict
    outDict = None
    outDict = {}

    def initialize(self):
        print(
            "SimPEG.SaveOutputDictEveryIteration will save your inversion progress as dictionary: '###-{!s}.npz'".format(
                self.fileName
            )
        )

    def endIter(self):

        regCombo = ["phi_ms", "phi_msx"]

        if self.prob[0].mesh.dim >= 2:
            regCombo += ["phi_msy"]

        if self.prob[0].mesh.dim == 3:
            regCombo += ["phi_msz"]

        # Initialize the output dict
        iterDict = {}

        # Save the data.
        iterDict["iter"] = self.opt.iter
        iterDict["beta"] = self.invProb.beta
        iterDict["phi_d"] = self.invProb.phi_d
        iterDict["phi_m"] = self.invProb.phi_m

        for label, fcts in zip(regCombo, self.reg.objfcts[0].objfcts):
            iterDict[label] = fcts(self.invProb.model)

        iterDict["f"] = self.opt.f
        iterDict["m"] = self.invProb.model
        iterDict["dpred"] = self.invProb.dpred

        if hasattr(self.dmisfit.objfcts[0].prob, "coordinate_system") is True:
            iterDict["coordinate_system"] = self.dmisfit.objfcts[
                0
            ].prob.coordinate_system
        else:
            iterDict["coordinate_system"] = False

        if hasattr(self.reg.objfcts[0], "eps_p") is True:
            iterDict["eps_p"] = self.reg.objfcts[0].eps_p
            iterDict["eps_q"] = self.reg.objfcts[0].eps_q

        iterDict["IRLSiterStart"] = None
        for direct in self.inversion.directiveList.dList:
            if isinstance(direct, Update_IRLS) and hasattr(direct, "iterStart") is True:
                if direct.IRLSiter > 0:
                    iterDict["IRLSiterStart"] = direct.iterStart

        if hasattr(self.reg.objfcts[0], "norms") is True:
            for objfct in self.reg.objfcts[0].objfcts:
                objfct.stashedR = None

            iterDict["lps"] = self.reg.objfcts[0].norms[0][0]
            iterDict["lpx"] = self.reg.objfcts[0].norms[0][1]

        iterDict["dphisdm"] = self.reg.objfcts[0].alpha_s * self.reg.objfcts[0].objfcts[
            0
        ].deriv(self.invProb.model)
        iterDict["dphixdm"] = self.reg.objfcts[0].alpha_x * self.reg.objfcts[0].objfcts[
            1
        ].deriv(self.invProb.model)

        # Save the file as a npz
        self.outDict[self.opt.iter] = iterDict

class SaveIterationsGeoH5(InversionDirective):
    """
        Saves inversion results to a geoh5 file
    """

    # Initialize the output dict
    h5_object = None
    channels = ["model"]
    attribute = "model"
    association = "VERTEX"
    sorting = None
    mapping = None
    save_objective_function = False
    data_type = {}
    replace_values = False
    no_data_value = None

    def initialize(self):

        if self.attribute == "predicted":
            if getattr(self.dmisfit, "objfcts", None) is not None:
                dpred = []
                for local_misfit in self.dmisfit.objfcts:
                    dpred.append(
                        np.asarray(local_misfit.survey.dpred(self.invProb.model))
                    )
                prop = np.hstack(dpred)
            else:
                prop = self.dmisfit.survey.dpred(self.invProb.model)
        else:
            prop = self.invProb.model

        if self.mapping is not None:
            prop = self.mapping * prop

        prop = self.check_mvi_format(prop)

        for ii, channel in enumerate(self.channels):

            attr = prop[ii :: len(self.channels)]

            if self.sorting is not None:
                attr = attr[self.sorting]

            data = self.h5_object.add_data(
                {
                    f"Iteration_0_"
                    + channel: {"association": self.association, "values": attr}
                }
            )

            data.entity_type.name = channel
            self.data_type[channel] = data.entity_type

        if self.save_objective_function:
            regCombo = ["phi_ms", "phi_msx"]

            if self.prob[0].mesh.dim >= 2:
                regCombo += ["phi_msy"]

            if self.prob[0].mesh.dim == 3:
                regCombo += ["phi_msz"]

            # Save the data.
            iterDict = {"beta": f"{self.invProb.beta:.3e}"}
            iterDict["phi_d"] = f"{self.invProb.phi_d:.3e}"
            iterDict["phi_m"] = f"{self.invProb.phi_m:.3e}"

            for label, fcts in zip(regCombo, self.reg.objfcts[0].objfcts):
                iterDict[label] = f"{fcts(self.invProb.model):.3e}"

            self.h5_object.parent.add_comment(
                json.dumps(iterDict), author=f"Iteration_{0}"
            )

        self.h5_object.workspace.finalize()

    def endIter(self):
        if self.attribute == "predicted":
            if getattr(self.dmisfit, "objfcts", None) is not None:
                dpred = []
                for local_misfit in self.dmisfit.objfcts:
                    dpred.append(
                        np.asarray(local_misfit.survey.dpred(self.invProb.model))
                    )
                prop = np.hstack(dpred)
            else:
                prop = self.dmisfit.survey.dpred(self.invProb.model)
        else:
            prop = self.invProb.model

        if self.mapping is not None:
            prop = self.mapping * prop

        prop = self.check_mvi_format(prop)

        for ii, channel in enumerate(self.channels):
            attr = prop[ii :: len(self.channels)]

            if self.sorting is not None:
                attr = attr[self.sorting]

            if self.replace_values and self.h5_object.get_data(
                f"Iteration_{self.opt.iter-1}_" + channel
            ):
                data = self.h5_object.get_data(
                    f"Iteration_{self.opt.iter-1}_" + channel
                )[0]
                data.name = f"Iteration_{self.opt.iter}_" + channel
                data.values = attr
            else:
                self.h5_object.add_data(
                    {
                        f"Iteration_{self.opt.iter}_"
                        + channel: {
                            "values": attr,
                            "association": self.association,
                            "entity_type": self.data_type[channel],
                        }
                    }
                )

        if self.save_objective_function:
            regCombo = ["phi_ms", "phi_msx"]

            if self.prob[0].mesh.dim >= 2:
                regCombo += ["phi_msy"]

            if self.prob[0].mesh.dim == 3:
                regCombo += ["phi_msz"]

            # Save objective function.
            if isinstance(self.invProb.beta, float):
                beta = self.invProb.beta
            else:
                beta = self.invProb.beta[0]

            iterDict = {"beta": f"{beta:.3e}"}

            if isinstance(self.invProb.phi_d, float):
                phi_d = self.invProb.phi_d
            else:
                phi_d = self.invProb.phi_d[0]

            if isinstance(self.invProb.phi_m, float):
                phi_m = self.invProb.phi_m
            else:
                phi_m = self.invProb.phi_m[0]

            iterDict["phi_d"] = f"{phi_d:.3e}"
            iterDict["phi_m"] = f"{phi_m:.3e}"

            for label, fcts in zip(regCombo, self.reg.objfcts[0].objfcts):
                iterDict[label] = f"{fcts(self.invProb.model):.3e}"

            self.h5_object.parent.add_comment(
                json.dumps(iterDict), author=f"Iteration_{self.opt.iter}"
            )

        self.h5_object.workspace.finalize()

    def check_mvi_format(self, values):
        if "mvi" in self.attribute:
            values = values.reshape((-1, 3), order="F")
            if self.no_data_value is not None:
                ndv_ind = values[:, 0] == self.no_data_value
                values[ndv_ind, :] = 0
            else:
                ndv_ind = np.zeros(values.shape[0], dtype="bool")

            if self.attribute == "mvi_model":
                values = np.linalg.norm(values, axis=1)
            elif self.attribute == "mvi_model_s":
                values = values[:, 0]
            elif self.attribute == "mvi_angles":
                atp = Utils.matutils.xyz2atp(values)
                values = atp.reshape((-1, 3), order="F")

            if "model" in self.attribute:
                values[ndv_ind] = self.no_data_value
            elif "angles" in self.attribute:
                values = np.rad2deg(values[:, 1:])
                values[ndv_ind, :] = self.no_data_value
                values = values.ravel()

        return values

class VectorInversion(InversionDirective):
    """
    Control a vector inversion from Cartesian to spherical coordinates
    """

    chifact_target = 1.0
    mref = None
    mode = "cartesian"
    inversion_type = "mvis"
    norms = []
    alphas = []

    @property
    def target(self):
        if getattr(self, "_target", None) is None:
            nD = 0
            for survey in self.survey:
                nD += survey.nD

            self._target = nD * 0.5 * self.chifact_target

        return self._target

    @target.setter
    def target(self, val):
        self._target = val

    def initialize(self):

        for reg in self.reg.objfcts:
            reg.model = self.invProb.model

        self.mref = reg.mref

        for prob in self.prob:
            if getattr(prob, "coordinate_system", None) is not None:
                prob.coordinate_system = self.mode

    def endIter(self):
        if (
            self.invProb.phi_d < self.target
        ) and self.mode == "cartesian":  # and self.inversion_type == 'mvis':
            print("Switching MVI to spherical coordinates")
            self.mode = "spherical"

            mstart = Utils.matutils.xyz2atp(
                self.invProb.model.reshape((-1, 3), order="F")
            )
            mref = Utils.matutils.xyz2atp(self.mref.reshape((-1, 3), order="F"))

            self.invProb.model = mstart
            self.invProb.beta *= 2
            self.opt.xc = mstart

            nC = mstart.reshape((-1, 3)).shape[0]
            self.opt.lower = np.kron(np.asarray([0, -np.inf, -np.inf]), np.ones(nC))
            self.opt.upper[nC:] = np.inf

            self.reg.mref = mref
            self.reg.model = mstart

            for prob in self.prob:
                if getattr(prob, "coordinate_system", None) is not None:
                    prob.coordinate_system = self.mode
                    prob.model = mstart

            for ind, reg_fun in enumerate(self.reg.objfcts):

                reg_fun.mref = mref
                reg_fun.model = mstart

                if ind > 0:
                    reg_fun.alpha_s = 0
                    reg_fun.eps_q = np.pi
                    for reg in reg_fun.objfcts:
                        reg.space = "spherical"

            # Add directives
            directiveList = []
            update_Jacobi = []
            IRLS = []
            for directive in self.inversion.directiveList.dList:
                if isinstance(directive, SaveIterationsGeoH5):
                    channels = []
                    for channel in directive.channels:
                        channels.append(channel + "_s")
                        directive.data_type[channel + "_s"] = directive.data_type[
                            channel
                        ]

                    directive.channels = channels

                    if directive.attribute == "mvi_model":
                        directive.attribute = "mvi_model_s"
                    elif directive.attribute == "mvi_angles":
                        directive.attribute = "mvi_angles_s"

                    directiveList.append(directive)

                elif isinstance(directive, SaveUBCModelEveryIteration):
                    directive.fileName = directive.fileName + "_S"
                    directiveList.append(directive)

                elif isinstance(directive, SaveUBCPredictedEveryIteration):
                    directive.fileName = directive.fileName + "_S"
                    directiveList.append(directive)

                elif isinstance(directive, Update_IRLS):
                    directive.sphericalDomain = True
                    directive.model = mstart
                    directive.coolingFactor = 1.5
                    IRLS = directive

                elif isinstance(directive, UpdatePreconditioner):
                    update_Jacobi = directive
                else:
                    directiveList.append(directive)

            directiveList = [
                ProjSpherical(),
                IRLS,
                UpdateSensitivityWeights(),
                update_Jacobi,
            ] + directiveList

            self.inversion.directiveList = directiveList
            directiveList[1].endIter()
            directiveList[2].endIter()
            directiveList[3].endIter()
        elif (self.invProb.phi_d < self.target) and self.mode == "spherical":

            for directive in self.inversion.directiveList.dList:
                if isinstance(directive, Update_IRLS) and directive.mode != 1:
                    directive.coolingFactor = 2

class Update_IRLS(InversionDirective):

    updateGamma = False
    f_old = np.inf
    f_min_change = 1e-2
    beta_tol = 1e-1
    beta_ratio_l2 = None
    prctile = 90
    chifact_target = 1.0

    # Solving parameter for IRLS (mode:2)
    IRLSiter = 0
    minGNiter = 1
    maxIRLSiter = 20
    iterStart = 0
    sphericalDomain = False

    # Beta schedule
    updateBeta = True
    betaSearch = True
    coolingFactor = 2.0
    coolingRate = 1
    ComboObjFun = False
    mode = 1
    coolEpsOptimized = True
    coolEps_p = True
    coolEps_q = True
    floorEps_p = [1e-8, 1e-8, 1e-8]
    floorEps_q = [1e-8, 1e-8, 1e-8]
    floorEpsEnforced = True
    coolEpsFact = 1.2
    silent = False
    fix_Jmatrix = False

    @property
    def target(self):
        if getattr(self, "_target", None) is None:
            nD = 0
            for survey in self.survey:
                nD += survey.nD

            self._target = nD * 0.5 * self.chifact_target

        return self._target

    @target.setter
    def target(self, val):
        self._target = val

    def initialize(self):

        if self.mode == 1:

            self.norms = []
            for reg in self.reg.objfcts:
                self.norms.append(reg.norms)
                reg.norms = np.c_[2.0, 2.0, 2.0, 2.0]
                reg.model = self.invProb.model

        for reg in self.reg.objfcts:
            reg.model = self.invProb.model
            for comp in reg.objfcts:
                self.f_old += comp(reg.model)

        self.phi_dm = []
        self.phi_dmx = []
        # Look for cases where the block models in to be scaled
        for prob in self.prob:

            if getattr(prob, "coordinate_system", None) is not None:
                if prob.coordinate_system == "spherical":
                    self.sphericalDomain = True

        if self.sphericalDomain:
            self.angleScale()

    def endIter(self):

        if self.sphericalDomain:
            self.angleScale()

        # Check if misfit is within the tolerance, otherwise scale beta
        if np.all(
            [
                np.abs(1.0 - float(self.invProb.phi_d) / self.target) > self.beta_tol,
                self.updateBeta,
                self.mode != 1,
            ]
        ):

            ratio = self.target / self.invProb.phi_d

            if ratio > 1:
                ratio = np.min([2.0, ratio])

            else:
                ratio = np.max([0.75, ratio])

            self.invProb.beta = self.invProb.beta * ratio

            if np.all([self.mode != 1, self.betaSearch]):
                # Re-use previous model and continue with new beta
                self.invProb.model = self.reg.objfcts[0].model
                self.opt.xc = self.reg.objfcts[0].model
                self.opt.iter -= 1
                return

        elif np.all([self.mode == 1, (self.opt.iter % self.coolingRate) == 0]):

            self.invProb.beta = self.invProb.beta / self.coolingFactor

        phim_new = 0
        for reg in self.reg.objfcts:
            reg.model = self.invProb.model
            for comp in reg.objfcts:
                phim_new += comp(reg.model)

        # Update the model used by the regularization
        phi_m_last = []
        for reg in self.reg.objfcts:
            reg.model = self.invProb.model
            phi_m_last += [reg(self.invProb.model)]

        # After reaching target misfit with l2-norm, switch to IRLS (mode:2)
        if np.all([float(self.invProb.phi_d) < self.target, self.mode == 1]):
            self.startIRLS()
            self.f_old = np.sum(phi_m_last)

        # Only update after GN iterations
        if np.all(
            [
                (float(self.opt.iter) - self.iterStart) % self.minGNiter == 0,
                self.mode != 1,
            ]
        ):
            if self.fix_Jmatrix:
                self.invProb.dmisfit.prob.fix_Jmatrix = True

            # Check for maximum number of IRLS cycles
            if self.IRLSiter == self.maxIRLSiter:
                if not self.silent:
                    print(
                        "Reach maximum number of IRLS cycles:"
                        + f" {int(self.maxIRLSiter)}"
                    )

                self.opt.stopNextIteration = True
                return

            # Print to screen
            for ii, reg in enumerate(self.reg.objfcts):

                if reg.eps_p > self.floorEps_p[ii] and self.coolEps_p:
                    reg.eps_p /= self.coolEpsFact

                elif self.floorEpsEnforced:
                    reg.eps_p = self.floorEps_p[ii]

                if reg.eps_q > self.floorEps_q[ii] and self.coolEps_q:
                    reg.eps_q /= self.coolEpsFact

                elif self.floorEpsEnforced:
                    reg.eps_q = self.floorEps_q[ii]

            self.IRLSiter += 1

            # Reset the regularization matrices so that it is
            # recalculated for current model. Do it to all levels of comboObj
            for reg in self.reg.objfcts:
                # If comboObj, go down one more level
                for comp in reg.objfcts:
                    comp.stashedR = None

            for dmis in self.dmisfit.objfcts:
                if getattr(dmis, "stashedR", None) is not None:
                    dmis.stashedR = None

            # Compute new model objective function value
            phi_m_new = []
            for reg in self.reg.objfcts:
                phi_m_new += [reg(self.invProb.model)]

            self.f_change = np.abs(self.f_old - phim_new) / self.f_old

            # Check if the function has changed enough
            if np.all(
                [
                    self.f_change < self.f_min_change,
                    self.IRLSiter > 1,
                    np.abs(1.0 - float(self.invProb.phi_d) / self.target)
                    < self.beta_tol,
                ]
            ):

                print("Minimum decrease in regularization." + "End of IRLS")
                self.opt.stopNextIteration = True
                return

            self.f_old = phim_new

            self.updateBeta = True
            self.invProb.phi_m_last = self.reg(self.invProb.model)

    def startIRLS(self):
        if not self.silent:
            print(
                "Reached starting chifact with l2-norm regularization:"
                + " Start IRLS steps..."
            )

        self.mode = 2

        if getattr(self.opt, "iter", None) is None:
            self.iterStart = 0
        else:
            self.iterStart = self.opt.iter

        self.invProb.phi_m_last = self.reg(self.invProb.model)

        # Either use the supplied epsilon, or fix base on distribution of
        # model values
        for reg in self.reg.objfcts:

            if getattr(reg, "eps_p", None) is None:

                reg.eps_p = np.percentile(
                    np.abs(reg.mapping * reg._delta_m(self.invProb.model)), self.prctile
                )
            if getattr(reg, "eps_q", None) is None:
                reg.eps_q = reg.eps_p

        # Re-assign the norms supplied by user l2 -> lp
        for reg, norms in zip(self.reg.objfcts, self.norms):
            reg.norms = norms

        # Save l2-model
        self.invProb.l2model = self.invProb.model.copy()

    def angleScale(self):
        """
            Update the scales used by regularization for the
            different block of models
        """
        # Currently implemented for MVI-S only
        max_p = []
        for reg in self.reg.objfcts[0].objfcts:
            eps_p = reg.epsilon
            norm_p = 2  # self.reg.objfcts[0].norms[0]
            f_m = abs(reg.f_m)
            max_p += [np.max(f_m)]

        max_p = np.asarray(max_p).max() * 2.0

        scales = [max_p / np.pi, max_p / np.pi]
        for obj, scale in zip(self.reg.objfcts[1:3], scales):
            obj.scales = np.ones(obj.scales.shape) * scale
        # self.reg.objfcts[0].scales = np.ones(self.reg.objfcts[0].scales.shape)
        # Probably doing rotated obj fun
        if len(self.reg) > 3:

            for obj, scale in zip(self.reg.objfcts[4:], scales):
                obj.scales = np.ones(obj.scales.shape) * scale

    def validate(self, directiveList):
        # check if a linear preconditioner is in the list, if not warn else
        # assert that it is listed after the IRLS directive
        dList = directiveList.dList
        self_ind = dList.index(self)
        lin_precond_ind = [isinstance(d, UpdatePreconditioner) for d in dList]

        if any(lin_precond_ind):
            assert lin_precond_ind.index(True) > self_ind, (
                "The directive 'UpdatePreconditioner' must be after Update_IRLS "
                "in the directiveList"
            )
        else:
            warnings.warn(
                "Without a Linear preconditioner, convergence may be slow. "
                "Consider adding `Directives.UpdatePreconditioner` to your "
                "directives list"
            )
        return True

class UpdatePreconditioner(InversionDirective):
    """
    Create a Jacobi preconditioner for the linear problem
    """

    onlyOnStart = False
    mapping = None
    misfitDiag = None
    epsilon = 1e-8

    def initialize(self):

        m = self.invProb.model
        # Create the pre-conditioner
        regDiag = np.zeros_like(self.invProb.model)

        for reg in self.reg.objfcts:
            # # Check if regularization has a projection
            # if getattr(reg.mapping, 'P', None) is None:
            #     regDiag += (reg.W.T*reg.W).diagonal()
            # else:
            #     P = reg.mapping.P
            #     regDiag += (P.T * (reg.W.T * (reg.W * P))).diagonal()
            regDiag += reg.deriv2(m).diagonal()

        # Deal with the linear case
        if getattr(self.opt, "JtJdiag", None) is None:
            JtJdiag = np.zeros_like(self.invProb.model)
            m = self.invProb.model
            for prob, dmisfit in zip(self.prob, self.dmisfit.objfcts):

                if getattr(prob, "getJtJdiag", None) is None:
                    assert getattr(prob, "getJ", None) is not None, (
                        "Problem does not have a getJ attribute."
                        + "Cannot form the sensitivity explicitely"
                    )
                    JtJdiag += np.sum(np.power((dmisfit.W * prob.getJ(m)), 2), axis=0)
                else:
                    JtJdiag += prob.getJtJdiag(m, W=dmisfit.W)
            self.opt.JtJdiag = JtJdiag

        diagA = self.opt.JtJdiag + self.invProb.beta * regDiag

        diagA[diagA != 0] = diagA[diagA != 0] ** -1.0
        PC = Utils.sdiag(diagA)

        self.opt.approxHinv = PC

    def endIter(self):

        # Cool the threshold parameter
        if self.onlyOnStart is True:
            return

        m = self.invProb.model
        # Create the pre-conditioner
        regDiag = np.zeros_like(self.invProb.model)

        for reg in self.reg.objfcts:
            regDiag += reg.deriv2(m).diagonal()

        # Assumes that opt.JtJdiag has been updated or static
        diagA = self.opt.JtJdiag + self.invProb.beta * regDiag

        diagA[diagA != 0] = diagA[diagA != 0] ** -1.0
        PC = Utils.sdiag(diagA)
        self.opt.approxHinv = PC

class UpdateSensitivityWeights(InversionDirective):
    """
    Directive to take care of re-weighting
    the non-linear magnetic problems.
    """

    mapping = None
    JtJdiag = None
    everyIter = True
    threshold = 1e-12
    switch = True

    def initialize(self):

        # Calculate and update sensitivity
        # for optimization and regularization
        self.update()

    def endIter(self):

        if self.everyIter:
            # Update inverse problem
            self.update()

    def update(self):
        # Get sum square of columns of J
        self.getJtJdiag()

        # Compute normalized weights
        self.wr = self.getWr()

        # Send a copy of JtJdiag for the preconditioner
        self.updateOpt()

        # Update the regularization
        self.updateReg()

    def getJtJdiag(self):
        """
            Compute explicitely the main diagonal of JtJ
            Good for any problem where J is formed explicitely
        """
        self.JtJdiag = []

        for prob, dmisfit in zip(self.prob, self.dmisfit.objfcts):
            m = self.invProb.model

            if getattr(prob, "getJtJdiag", None) is None:
                assert getattr(prob, "getJ", None) is not None, (
                    "Problem does not have a getJ attribute."
                    + "Cannot form the sensitivity explicitely"
                )

                self.JtJdiag += [
                    Utils.mkvc(np.sum((dmisfit.W * prob.getJ(m)) ** (2.0), axis=0))
                ]
            else:
                self.JtJdiag += [prob.getJtJdiag(m, W=dmisfit.W)]

        return self.JtJdiag

    def getWr(self):
        """
            Take the diagonal of JtJ and return
            a normalized sensitivty weighting vector
        """

        wr = np.zeros_like(self.invProb.model)
        if self.switch:
            for prob_JtJ, prob, dmisfit in zip(
                self.JtJdiag, self.prob, self.dmisfit.objfcts
            ):

                wr += prob_JtJ + self.threshold

            wr = wr ** 0.5
            wr /= wr.max()
        else:
            wr += 1.0

        return wr

    def updateReg(self):
        """
            Update the cell weights with the approximated sensitivity
        """

        for reg in self.reg.objfcts:
            reg.cell_weights = reg.mapping * (self.wr)

    def updateOpt(self):
        """
            Update a copy of JtJdiag to optimization for preconditioner
        """
        # if self.ComboMisfitFun:
        JtJdiag = np.zeros_like(self.invProb.model)
        for prob, JtJ, dmisfit in zip(self.prob, self.JtJdiag, self.dmisfit.objfcts):

            JtJdiag += JtJ

        self.opt.JtJdiag = JtJdiag

class ProjSpherical(InversionDirective):
    """
        Trick for spherical coordinate system.
        Project \theta and \\phi angles back to [-\\pi,\\pi] using
        back and forth conversion.
        spherical->cartesian->spherical
    """

    def initialize(self):

        x = self.invProb.model
        # Convert to cartesian than back to avoid over rotation
        nC = int(len(x) / 3)

        xyz = Utils.matutils.atp2xyz(x.reshape((nC, 3), order="F"))
        m = Utils.matutils.xyz2atp(xyz.reshape((nC, 3), order="F"))

        self.invProb.model = m

        for prob in self.prob:
            prob.model = m

        self.opt.xc = m

    def endIter(self):
        x = self.invProb.model
        nC = int(len(x) / 3)

        # Convert to cartesian than back to avoid over rotation
        xyz = Utils.matutils.atp2xyz(x.reshape((nC, 3), order="F"))
        m = Utils.matutils.xyz2atp(xyz.reshape((nC, 3), order="F"))

        self.invProb.model = m

        phi_m_last = []
        for reg in self.reg.objfcts:
            reg.model = self.invProb.model
            phi_m_last += [reg(self.invProb.model)]

        self.invProb.phi_m_last = phi_m_last

        for prob in self.prob:
            prob.model = m

        self.opt.xc = m

class JointAmpMVI(InversionDirective):
    """
        Directive controlling the joint inversion of
        magnetic amplitude data and MVI. Use the vector
        magnetization model (M) to update the linear amplitude
        operator.

    """

    amp = None
    minGNiter = 1
    jointMVIS = False
    updateM = False

    def initialize(self):

        # Get current MVI model and update MAI sensitivity
        # if isinstance(self.prob, list):

        m = self.invProb.model.copy()
        for prob in self.prob:

            if isinstance(prob, Magnetics.MagneticVector):
                if prob.coordinate_system == "spherical":
                    xyz = Magnetics.atp2xyz(
                        (prob.chiMap * m).reshape((int(len(m) / 3), 3), order="F")
                    )
                    self.jointMVIS = True
                elif prob.coordinate_system == "cartesian":
                    xyz = prob.chiMap * m

            if isinstance(prob, Magnetics.MagneticAmplitude):
                self.amp = prob.chiMap * m

        for prob in self.prob:
            if isinstance(prob, Magnetics.MagneticAmplitude):
                if self.jointMVIS:
                    prob.jointMVIS = True

                nC = prob.mesh.nC

                mcol = xyz.reshape((nC, 3), order="F")
                amp = np.sum(mcol ** 2.0, axis=1) ** 0.5
                M = Utils.sdiag(1.0 / amp) * mcol

        else:
            assert "This directive needs to used on a ComboObjective"

    def endIter(self):

        # Get current MVI model and update magnetization model for MAI
        m = self.invProb.model.copy()
        for prob in self.prob:

            if isinstance(prob, Magnetics.MagneticVector):
                if prob.coordinate_system == "spherical":
                    xyz = Magnetics.atp2xyz(
                        (prob.chiMap * m).reshape((int(len(m) / 3), 3), order="F")
                    )

                elif prob.coordinate_system == "cartesian":
                    xyz = prob.chiMap * m

            if isinstance(prob, Magnetics.MagneticAmplitude):
                if prob.chiMap.shape[0] == 3 * prob.mesh.nC:

                    nC = prob.mesh.nC

                    mcol = (prob.chiMap * m).reshape((nC, 3), order="F")
                    self.amp = np.sum(mcol ** 2.0, axis=1) ** 0.5
                else:

                    self.amp = prob.chiMap * m

        for prob in self.prob:
            if np.all([isinstance(prob, Magnetics.MagneticAmplitude), self.updateM]):

                nC = prob.mesh.nC

                mcol = xyz.reshape((nC, 3), order="F")
                amp = np.sum(mcol ** 2.0, axis=1) ** 0.5

                M = Utils.sdiag(1.0 / amp) * mcol

                prob.M = M
                prob._Mxyz = None
                prob.Mxyz
                # if prob.model is None:
                #     prob.model = prob.chiMap * m

                # ampW = (amp/amp.max() + 1e-2)**-1.
                # prob.W = ampW

            if isinstance(prob, Magnetics.MagneticVector):

                if (prob.coordinate_system == "cartesian") and (self.amp is not None):

                    ampW = (self.amp / self.amp.max() + 1e-2) ** -1.0
                    ampW = np.r_[ampW, ampW, ampW]

                    # Scale max values
                    prob.W = ampW

class UpdateApproxJtJ(InversionDirective):
    """
        Create approx-sensitivity base weighting using the probing method
    """

    k = None  # Number of probing cycles
    itr = None  # Iteration number to update Wj, or always update if None

    def endIter(self):

        if self.itr is None or self.itr == self.opt.iter:

            m = self.invProb.model
            if self.k is None:

                nD = 0
                for survey in self.survey:
                    nD += survey.nD

                self.k = int(nD / 10)

            def JtJv(v):

                Jv = self.prob.Jvec(m, v)

                return self.prob.Jtvec(m, Jv)

            JtJdiag = Utils.diagEst(JtJv, len(m), k=self.k)
            JtJdiag = JtJdiag / max(JtJdiag)

            self.reg.wght = JtJdiag

class ScaleComboReg(InversionDirective):
    """
    Directive to take care of re-weighting
    the non-linear magnetic problems.

    """

    # coordinate_system = 'Amp'
    # test = False
    mapping = None
    ComboRegFun = False
    ComboMisfitFun = False
    JtJdiag = None
    everyIter = True

    def initialize(self):

        # for reg in self.reg.objfcts:
        m = self.invProb.model

        scale = (
            np.abs(self.reg.objfcts[0](m)).max() / np.abs(self.reg.objfcts[1](m)).max()
        )
        self.reg.objfcts[1].scale = scale

    def endIter(self):

        m = self.invProb.model

        scale = (
            np.abs(self.reg.objfcts[0].deriv(m)).max()
            / np.abs(self.reg.objfcts[1].deriv(m)).max()
        )
        self.reg.objfcts[1].scale = scale


#!/usr/bin/env python

import os
import time
import warnings

import numpy

from cogent3.maths.optimisers import ParameterOutOfBoundsError, maximise
from cogent3.maths.solve import find_root

Float = numpy.core.numerictypes.sctype2char(float)

TRACE_DEFAULT = "COGENT3_TRACE" in os.environ
TRACE_SCALE = 100000

__author__ = "Peter Maxwell"
__copyright__ = "Copyright 2007-2020, The Cogent Project"
__credits__ = ["Peter Maxwell", "Gavin Huttley", "Daniel McDonald"]
__license__ = "BSD-3"
__version__ = "2020.2.7a"
__maintainer__ = "Peter Maxwell"
__email__ = "pm67nz@gmail.com"
__status__ = "Production"

# This is the 'live' layer of the recalculation system
# Cells and OptPars are held by a Calculator
# For docstring see definitions.py

class CalculationInterupted(Exception):
    pass

class OptPar(object):
    """One parameter, as seen by the optimiser, eg: length of one edge.
    An OptPar reports changes to the ParameterValueSet for its parameter.
    """

    is_constant = False
    recycled = False
    args = ()
    # Use of __slots__ here and in Cell gives 8% speedup on small calculators.
    __slots__ = [
        "clients",
        "client_ranks",
        "name",
        "lower",
        "default_value",
        "upper",
        "scope",
        "order",
        "label",
        "consequences",
        "rank",
    ]

    def __init__(self, name, scope, bounds):
        self.clients = []
        self.client_ranks = []
        self.name = name
        for (attr, v) in zip(["lower", "default_value", "upper"], bounds):
            setattr(self, attr, float(v))

        # controls order in optimiser - group for LF
        self.scope = scope
        self.order = (len(scope), scope and min(scope), name)
        self.label = self.name

    def add_client(self, client):
        self.clients.append(client)

    def __lt__(self, other):
        # optimisation is more efficient if params for one edge are neighbours
        return self.order < other.order

    def __eq__(self, other):
        # optimisation is more efficient if params for one edge are neighbours
        return self.order == other.order

    def __ne__(self, other):
        # optimisation is more efficient if params for one edge are neighbours
        return self.order != other.order

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, self.label)

    def get_optimiser_bounds(self):
        lower = self.transform_to_optimiser(self.lower)
        upper = self.transform_to_optimiser(self.upper)
        return (lower, upper)

    def transform_from_optimiser(self, value):
        return value

    def transform_to_optimiser(self, value):
        return value

class LogOptPar(OptPar):
    # For ratios, optimiser sees log(param value).  Conversions to/from
    # optimiser representation are only done by Calculator.change(),
    # .get_value_array() and .getBoundsArrrays().

    def transform_from_optimiser(self, value):
        return numpy.exp(value)

    def transform_to_optimiser(self, value):
        try:
            return numpy.log(value)
        except OverflowError:
            raise OverflowError("log(%s)" % value)

class EvaluatedCell(object):
    __slots__ = [
        "client_ranks",
        "rank",
        "calc",
        "args",
        "is_constant",
        "clients",
        "failure_count",
        "name",
        "arg_ranks",
        "consequences",
        "recycled",
        "default",
    ]

    def __init__(self, name, calc, args, recycling=None, default=None):
        self.name = name
        self.rank = None
        self.calc = calc
        self.default = default
        self.args = tuple(args)

        self.recycled = recycling
        if recycling:
            self.args = (self,) + self.args

        self.is_constant = True
        for arg in args:
            arg.add_client(self)
            if not arg.is_constant:
                self.is_constant = False

        self.clients = []
        self.client_ranks = []
        self.failure_count = 0

    def add_client(self, client):
        self.clients.append(client)

    def update(self, data):
        data[self.rank] = self.calc(*[data[arg_rank] for arg_rank in self.arg_ranks])

    def prime(self, data_sets):
        if self.is_constant:
            # Just calc once
            self.update(data_sets[0])
            for data in data_sets[1:]:
                data[self.rank] = data_sets[0][self.rank]
        else:
            for data in data_sets:
                self.update(data)

    def report_error(self, detail, data):
        self.failure_count += 1
        if self.failure_count <= 5:
            print(("%s in calculating %s:", detail.__class__.__name__, self.name))
        if self.failure_count == 5:
            print("Additional failures of this type will not be reported.")
        if self.failure_count < 2:
            print("%s inputs were:", len(self.arg_ranks))
            for (i, arg) in enumerate(self.arg_ranks):
                print("%s: " % i + repr(data[arg]))

class ConstCell(object):
    __slots__ = ["name", "scope", "value", "rank", "consequences", "clients"]

    recycled = False
    is_constant = True
    args = ()

    def __init__(self, name, value):
        self.name = name
        self.clients = []
        self.value = value

    def add_client(self, client):
        self.clients.append(client)

class Calculator(object):
    """A complete hierarchical function with N evaluation steps to call
    for each change of inputs.  Made by a ParameterController."""

    def __init__(self, cells, defns, trace=None, with_undo=True):
        if trace is None:
            trace = TRACE_DEFAULT
        self.with_undo = with_undo
        self.results_by_id = defns
        self.opt_pars = []
        other_cells = []
        for cell in cells:
            if isinstance(cell, OptPar):
                self.opt_pars.append(cell)
            else:
                other_cells.append(cell)
        self._cells = self.opt_pars + other_cells
        data_sets = [[0], [0, 1]][self.with_undo]
        self.cell_values = [[None] * len(self._cells) for switch in data_sets]
        self.arg_ranks = [[] for cell in self._cells]
        for (i, cell) in enumerate(self._cells):
            cell.rank = i
            cell.consequences = {}
            if isinstance(cell, OptPar):
                for switch in data_sets:
                    self.cell_values[switch][i] = cell.default_value
            elif isinstance(cell, ConstCell):
                for switch in data_sets:
                    self.cell_values[switch][i] = cell.value
            elif isinstance(cell, EvaluatedCell):
                cell.arg_ranks = []
                for arg in cell.args:
                    if hasattr(arg, "client_ranks"):
                        arg.client_ranks.append(i)
                    self.arg_ranks[i].append(arg.rank)
                    cell.arg_ranks.append(arg.rank)

                try:
                    cell.prime(self.cell_values)
                except KeyboardInterrupt:
                    raise
                except Exception as detail:
                    print(("Failed initial calculation of %s" % cell.name))
                    raise
            else:
                raise RuntimeError("Unexpected Cell type %s" % type(cell))

        self._switch = 0
        self.recycled_cells = [cell.rank for cell in self._cells if cell.recycled]
        self.spare = [None] * len(self._cells)

        for cell in self._cells[::-1]:
            for arg in cell.args:
                arg.consequences[cell.rank] = True
                arg.consequences.update(cell.consequences)

        self._programs = {}
        # Just for timings pre-calc these
        for opt_par in self.opt_pars:
            self.cells_changed_by([(opt_par.rank, None)])

        self.last_values = self.get_value_array()
        self.last_undo = []
        self.elapsed_time = 0.0
        self.evaluations = 0
        self.set_tracing(trace)
        self.optimised = False

    def graphviz(self):
        """Returns a string in the 'dot' graph description language used by the
        program 'Graphviz'.  One box per cell, grouped by Defn."""

        lines = ["digraph G {\n rankdir = LR\n ranksep = 1\n"]
        evs = []
        for cell in self._cells:
            if cell.name not in evs:
                evs.append(cell.name)
        nodes = dict([(name, []) for name in evs])
        edges = []
        for cell in self._cells:
            if hasattr(cell, "name"):
                nodes[cell.name].append(cell)
                for arg in cell.args:
                    if arg is not cell:
                        edges.append(
                            '"%s":%s -> "%s":%s'
                            % (arg.name, arg.rank, cell.name, cell.rank)
                        )
        for name in evs:
            all_const = True
            some_const = False
            enodes = [name.replace("edge", "QQQ")]
            for cell in nodes[name]:
                value = self._get_current_cell_value(cell)
                if isinstance(value, float):
                    label = "%5.2e" % value
                else:
                    label = "[]"
                label = "<%s> %s" % (cell.rank, label)
                enodes.append(label)
                all_const = all_const and cell.is_constant
                some_const = some_const or cell.is_constant
            enodes = "|".join(enodes)
            colour = ["", " fillcolor=gray90, style=filled,"][some_const]
            colour = [colour, " fillcolor=gray, style=filled,"][all_const]
            lines.append(
                '"%s" [shape = "record",%s label="%s"];' % (name, colour, enodes)
            )
        lines.extend(edges)
        lines.append("}")
        return "\n".join(lines).replace("edge", "egde").replace("QQQ", "edge")

    def optimise(self, **kw):
        x = self.get_value_array()
        bounds = self.get_bounds_vectors()
        maximise(self, x, bounds, **kw)
        self.optimised = True

    def set_tracing(self, trace=False):
        """With 'trace' true every evaluated is printed.  Useful for profiling
        and debugging."""

        self.trace = trace
        if trace:
            print()
            n_opars = len(self.opt_pars)
            n_cells = len([c for c in self._cells if not c.is_constant])
            print(n_opars, "OptPars and", n_cells - n_opars, "derived values")
            print("OptPars: ", ", ".join([par.name for par in self.opt_pars]))
            print("Times in 1/%sths of a second" % TRACE_SCALE)

            groups = []
            groupd = {}
            for cell in self._cells:
                if cell.is_constant or not isinstance(cell, EvaluatedCell):
                    continue
                if cell.name not in groupd:
                    group = []
                    groups.append((cell.name, group))
                    groupd[cell.name] = group
                groupd[cell.name].append(cell)

            widths = []
            for (name, cells) in groups:
                width = 4 + len(cells)
                widths.append(min(15, width))
            self._cellsGroupedForDisplay = list(zip(groups, widths))
            for ((name, cells), width) in self._cellsGroupedForDisplay:
                print(name[:width].ljust(width), "|", end=" ")
            print()
            for width in widths:
                print("-" * width, "|", end=" ")
            print()

    def get_value_array(self):
        """This being a caching function, you can ask it for its current
        input!  Handy for initialising the optimiser."""
        values = [
            p.transform_to_optimiser(self._get_current_cell_value(p))
            for p in self.opt_pars
        ]
        return values

    # get_bounds_vectors and testoptparvector make up the old LikelihoodFunction
    # interface expected by the optimiser.

    def get_bounds_vectors(self):
        """2 arrays: minimums, maximums"""
        lower = numpy.zeros([len(self.opt_pars)], Float)
        upper = numpy.zeros([len(self.opt_pars)], Float)
        for (i, opt_par) in enumerate(self.opt_pars):
            (lb, ub) = opt_par.get_optimiser_bounds()
            lower[i] = lb
            upper[i] = ub
        return (lower, upper)

    def fuzz(self, random_series=None, seed=None):
        # Slight randomisation suitable for removing right-on-the-
        # ridge starting points before local optimisation.
        if random_series is None:
            import random

            random_series = random.Random()
        if seed is not None:
            random_series.seed(seed)
        X = self.get_value_array()
        for (i, (l, u)) in enumerate(zip(*self.get_bounds_vectors())):
            sign = random_series.choice([-1, +1])
            step = random_series.uniform(+0.05, +0.025)
            X[i] = max(l, min(u, (1.0 + sign * step * X[i])))
        self.testoptparvector(X)
        self.optimised = False

    def testoptparvector(self, values):
        """AKA self().  Called by optimisers.  Returns the output value
        after doing any recalculation required for the new input 'values'
        array"""

        assert len(values) == len(self.opt_pars)
        changes = [
            (i, new)
            for (i, (old, new)) in enumerate(zip(self.last_values, values))
            if old != new
        ]
        return self.change(changes)

    __call__ = testoptparvector

    def testfunction(self):
        """Return the current output value without changing any inputs"""
        return self._get_current_cell_value(self._cells[-1])

    def change(self, changes):
        """Returns the output value after applying 'changes', a list of
        (optimisable_parameter_ordinal, new_value) tuples."""

        t0 = time.time()
        self.evaluations += 1

        # If ALL of the changes made in the last step are reversed in this step
        # then it is safe to undo them first, taking advantage of the 1-deep
        # cache.
        if self.with_undo and self.last_undo:
            for (i, v) in self.last_undo:
                if (i, v) not in changes:
                    break
            else:
                changes = [ch for ch in changes if ch not in self.last_undo]
                self._switch = not self._switch
                for (i, v) in self.last_undo:
                    self.last_values[i] = v

        self.last_undo = []
        program = self.cells_changed_by(changes)

        if self.with_undo:
            self._switch = not self._switch
            data = self.cell_values[self._switch]
            base = self.cell_values[not self._switch]

            # recycle and undo interact in bad ways
            for rank in self.recycled_cells:
                if data[rank] is not base[rank]:
                    self.spare[rank] = data[rank]
            data[:] = base[:]
            for cell in program:
                if cell.recycled:
                    if data[cell.rank] is base[cell.rank]:
                        data[cell.rank] = self.spare[cell.rank]
                        assert data[cell.rank] is not base[cell.rank]
        else:
            data = self.cell_values[self._switch]

        # Set new OptPar values
        changed_optpars = []
        for (i, v) in changes:
            if i < len(self.opt_pars):
                assert isinstance(v * 1.0, float), v
                changed_optpars.append((i, self.last_values[i]))
                self.last_values[i] = v
                data[i] = self.opt_pars[i].transform_from_optimiser(v)
            else:
                data[i] = v

        try:
            if self.trace:
                self.tracing_update(changes, program, data)
            else:
                self.plain_update(program, data)

            # if non-optimiser parameter was set then undo is invalid
            if self.last_undo and max(self.last_undo)[0] >= len(self.opt_pars):
                self.last_undo = []
            else:
                self.last_undo = changed_optpars

        except CalculationInterupted as detail:
            if self.with_undo:
                self._switch = not self._switch
            for (i, v) in changed_optpars:
                self.last_values[i] = v
            self.last_undo = []
            (cell, exception) = detail.args
            raise exception

        finally:
            self.elapsed_time += time.time() - t0

        return self.cell_values[self._switch][-1]

    def cells_changed_by(self, changes):
        # What OptPars have been changed determines cells to update
        change_key = list(dict(changes).keys())
        change_key.sort()
        change_key = tuple(change_key)
        if change_key in self._programs:
            program = self._programs[change_key]
        else:
            # Make a list of the cells to update and cache it.
            consequences = {}
            for i in change_key:
                consequences.update(self._cells[i].consequences)
            self._programs[change_key] = program = [
                cell for cell in self._cells if cell.rank in consequences
            ]
        return program

    def plain_update(self, program, data):
        try:
            for cell in program:
                data[cell.rank] = cell.calc(*[data[a] for a in cell.arg_ranks])
        except ParameterOutOfBoundsError as detail:
            # Non-fatal error, just cancel this calculation.
            raise CalculationInterupted(cell, detail)
        except ArithmeticError as detail:
            # Non-fatal but unexpected error. Warn and cancel this calculation.
            cell.report_error(detail, data)
            raise CalculationInterupted(cell, detail)

    def tracing_update(self, changes, program, data):
        # Does the same thing as plain_update, but also produces lots of
        # output showing how long each step of the calculation takes.
        # One line per call, '-' for undo, '+' for calculation

        exception = None
        elapsed = {}
        for cell in program:
            try:
                t0 = time.time()
                data[cell.rank] = cell.calc(*[data[a] for a in cell.arg_ranks])
                t1 = time.time()
            except (ParameterOutOfBoundsError, ArithmeticError) as exception:
                error_cell = cell
                break
            elapsed[cell.rank] = t1 - t0

        tds = []
        for ((name, cells), width) in self._cellsGroupedForDisplay:
            text = "".join([" +"[cell.rank in elapsed] for cell in cells])
            elap = sum([elapsed.get(cell.rank, 0) for cell in cells])
            if len(text) > width - 4:
                edge_width = min(len(text), (width - 4 - 3)) // 2
                elipsis = ["   ", "..."][not not text.strip()]
                text = text[:edge_width] + elipsis + text[-edge_width:]
            tds.append("%s%4s" % (text, int(TRACE_SCALE * elap + 0.5) or ""))

        par_descs = []
        for (i, v) in changes:
            cell = self._cells[i]
            if isinstance(cell, OptPar):
                par_descs.append("%s=%8.6f" % (cell.name, v))
            else:
                par_descs.append("%s=?" % cell.name)
        par_descs = ", ".join(par_descs)[:22].ljust(22)
        print(" | ".join(tds + [""]), end=" ")
        if exception:
            print("%15s | %s" % ("", par_descs))
            error_cell.report_error(exception, data)
            raise CalculationInterupted(cell, exception)
        else:
            print("%-15s | %s" % (repr(data[-1])[:15], par_descs))

    def measure_evals_per_second(self, time_limit=1.0, wall=True, sa=False):
        # Returns an estimate of the number of evaluations per second
        # an each-optpar-in-turn simulated annealing type optimiser
        # can achive, spending not much more than 'time_limit' doing
        # so.  'wall'=False causes process time to be used instead of
        # wall time.
        # 'sa' makes it simulated-annealing-like, with frequent backtracks
        if wall:
            now = time.time
        else:
            now = time.clock
        x = self.get_value_array()
        samples = []
        elapsed = 0.0
        rounds_per_sample = 2
        while elapsed < time_limit and len(samples) < 5:
            time.sleep(0.01)
            t0 = now()
            last = []
            for j in range(rounds_per_sample):
                for (i, v) in enumerate(x):
                    # Not a real change, but works like one.
                    self.change(last + [(i, v)])
                    if sa and (i + j) % 2:
                        last = [(i, v)]
                    else:
                        last = []
            # Use one agreed on delta otherwise different cpus will finish the
            # loop at different times causing chaos.
            delta = now() - t0
            if delta < 0.1:
                # time.clock is low res, so need to ensure each sample
                # is long enough to take SOME time.
                rounds_per_sample *= 2
                continue
            else:
                rate = rounds_per_sample * len(x) / delta
                samples.append(rate)
                elapsed += delta

        if wall:
            samples.sort()
            return samples[len(samples) // 2]
        else:
            return sum(samples) / len(samples)

    def _get_current_cell_value(self, cell):
        return self.cell_values[self._switch][cell.rank]

    def get_current_cell_values_for_defn(self, defn):
        cells = self.results_by_id[id(defn)]
        return [self.cell_values[self._switch][cell.rank] for cell in cells]

    def __get_bounded_root(self, func, origX, direction, bound, xtol):
        return find_root(
            func,
            origX,
            direction,
            bound,
            xtol=xtol,
            expected_exception=(ParameterOutOfBoundsError, ArithmeticError),
        )

    def _get_current_cell_interval(self, opt_par, dropoff, xtol=None):
        # (min, opt, max) tuples for each parameter where f(min) ==
        # f(max) == f(opt)-dropoff.  Uses None when a bound is hit.
        # assert self.optimised, "Call optimise() first"
        origY = self.testfunction()
        (lower, upper) = opt_par.get_optimiser_bounds()
        opt_value = self._get_current_cell_value(opt_par)
        origX = opt_par.transform_to_optimiser(opt_value)

        def func(x):
            Y = self.change([(opt_par.rank, x)])
            return Y - (origY - dropoff)

        try:
            lowX = self.__get_bounded_root(func, origX, -1, lower, xtol)
            highX = self.__get_bounded_root(func, origX, +1, upper, xtol)
        finally:
            func(origX)

        triple = []
        for x in [lowX, origX, highX]:
            if x is not None:
                x = opt_par.transform_from_optimiser(x)
            triple.append(x)
        return tuple(triple)


import os
import copy
import numpy as np
import torch
from torch import nn
from torch.autograd import grad
from procedures import get_ntk_n

reward_type2index = {
    'ntk': 0,
    'exp': 1,
    'constraint': 3
}
index2reward_type = {
    0: 'ntk',
    1: 'exp',
    3: 'constraint'
}

def kaiming_normal_fanin_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
        if hasattr(m, 'bias') and m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight.data)
        nn.init.constant_(m.bias.data, 0.0)

def kaiming_normal_fanout_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if hasattr(m, 'bias') and m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.LayerNorm):
        nn.init.ones_(m.weight.data)
        nn.init.constant_(m.bias.data, 0.0)

def init_model(model, method='kaiming_norm_fanin'):
    if method == 'kaiming_norm_fanin':
        model.apply(kaiming_normal_fanin_init)
    elif method == 'kaiming_norm_fanout':
        model.apply(kaiming_normal_fanout_init)
    return model

class TEG(object):
    def __init__(self, loader, loader_val, class_num=1000, repeat=3, size_curve=(500, 3, 16, 16), batch_curve=6, reward_types=["ntk", "exp"], buffer_size=10, constraint_weight=0.):
        # self.__super__()
        self.repeat = repeat
        self.constraint_weight = constraint_weight # e.g. FLOPs constraint
        self.batch_size_curve = batch_curve

        self.reward_type2index = reward_type2index
        self.index2reward_type = index2reward_type
        self._reward_types = reward_types
        self._reward_sign = {"ntk": -1, "exp": 1, "constraint": -1} # ntk: lower the better; exp: higher the better
        self._buffers = {key: [] for key in self._reward_types}
        self._buffers['constraint'] = []
        self._buffers_bad = [] # indicator of bad architectures
        self._buffers_change = {key: [] for key in self._reward_types}
        self._buffers_change['constraint'] = []
        self._buffer_length = buffer_size
        self._class_num = class_num
        # build fixed data samples
        self._ntk_input_data = []
        for i, (inputs, targets) in enumerate(loader):
            if i >= self.repeat: break
            self._ntk_input_data.append((inputs, targets))
            self.batch_size = len(inputs)
        self._ntk_target_data = [] # for NTK kernel regression
        for i, (inputs, targets) in enumerate(loader_val):
            if i >= self.repeat: break
            self._ntk_target_data.append((inputs, targets))
        # Curve complexity
        n_interp, C, H, W = size_curve
        self.theta = []; self.curve_input = []
        for _ in range(self.repeat):
            self.theta.append(torch.linspace(0, 2 * np.pi, n_interp))
            self.theta[-1].requires_grad_(True)
            self.curve_input.append(torch.matmul(torch.svd(torch.randn(H*W*C, 2))[0], torch.stack([torch.cos(self.theta[-1]), torch.sin(self.theta[-1])])).T.reshape((n_interp, C, H, W)).cuda(non_blocking=True))
            self.curve_input[-1].requires_grad_(True)

    def reset(self, constraint_weight=0):
        self.constraint_weight = constraint_weight
        # self._reward_types = reward_types
        self._buffers = {key: [] for key in self._reward_types}; self._buffers['constraint'] = []
        self._buffers_bad = [] # indicator of bad architectures
        self._buffers_change = {key: [] for key in self._reward_types}; self._buffers_change['constraint'] = []

    def set_network(self, network):
        if hasattr(self, "_networks"):
            del self._networks
        self._networks = []
        for _ in range(self.repeat):
            net = copy.deepcopy(network)
            net.apply(net._init_weights)
            self._networks.append(net.cuda())

    def get_ntk(self):
        for net in self._networks:
            net.switch_norm('ln')
        ntks = get_ntk_n(self._ntk_input_data, self._networks, criterion=torch.nn.CrossEntropyLoss(), train_mode=True, num_batch=1, num_classes=self._class_num)
        for network in self._networks:
            network.zero_grad()
        torch.cuda.empty_cache()
        return np.mean(ntks)

    def get_curve_complexity(self):
        for net in self._networks:
            net.switch_norm('id')
        LE = [0 for _ in range(len(self._networks))]
        for net_idx, network in enumerate(self._networks):
            network = network.cuda()
            network.train()
            network.zero_grad()
            _idx = 0
            while _idx < len(self.curve_input[net_idx]):
                output = network.forward_features(self.curve_input[net_idx][_idx:_idx+self.batch_size_curve])[1]
                _idx += self.batch_size_curve
                output = output.reshape(output.size(0), -1)
                n, c = output.size()
                jacobs = []
                for coord in range(c):
                    output[:, coord].backward(torch.ones_like(output[:, coord]), retain_graph=True)
                    # actually only "batch_size" number of thetas have grad, but it is ok, since zeros won't contribute to gE.sum()
                    jacobs.append(self.theta[net_idx].grad.detach().clone())
                    self.theta[net_idx].grad.zero_()
                jacobs = torch.stack(jacobs, 0)
                jacobs = jacobs.permute(1, 0) # num_theta x c
                gE = torch.einsum('nd,nd->n', jacobs, jacobs).sqrt()
                LE[net_idx] += gE.sum().item()
                torch.cuda.empty_cache()
            network = network.cpu()
            torch.cuda.empty_cache()
        for net in self._networks:
            net.switch_norm('ln')
        return np.mean(LE)

    def get_curve_complexity_gauss(self):
        for net in self._networks:
            net.switch_norm('id')
        LG = [0 for _ in range(len(self._networks))]
        for net_idx, network in enumerate(self._networks):
            network = network.cuda()
            network.train()
            network.zero_grad()
            _idx = 0
            v_s = [] # 1st derivative
            while _idx < len(self.curve_input[net_idx]):
                output = network.forward_features(self.curve_input[net_idx][_idx:_idx+self.batch_size_curve])[1]
                output = output.reshape(output.size(0), -1)
                n, c = output.size()
                _v_s = [] # 1st derivative
                for coord in range(c):
                    v = grad(output[:, coord].sum(), self.theta[net_idx], create_graph=True, retain_graph=True)[0][_idx:_idx+self.batch_size_curve] # batch size (of thetas)
                    _v_s.append(v.detach().clone())
                v_s.append(torch.stack(_v_s, 0).permute(1, 0)) # bach_size x c
                _idx += self.batch_size_curve
            v_s = torch.cat(v_s, 0) # num_thetas x c
            v_s_norm = v_s.norm(2, dim=1, keepdim=True) # norm over c of all thetas
            _idx = 0
            while _idx < len(self.curve_input[net_idx]):
                output = network.forward_features(self.curve_input[net_idx][_idx:_idx+self.batch_size_curve])[1]
                output = output.reshape(output.size(0), -1)
                n, c = output.size()
                d_v_hat_s = [] # 2nd derivative
                for coord in range(c):
                    v = grad(output[:, coord].sum(), self.theta[net_idx], create_graph=True, retain_graph=True)[0][_idx:_idx+self.batch_size_curve] # batch size (of thetas)
                    d_v_hat = grad((v / v_s_norm[_idx:_idx+self.batch_size_curve]).sum(), self.theta[net_idx], create_graph=True, retain_graph=True)[0][_idx:_idx+self.batch_size_curve] # batch size (of thetas)
                    d_v_hat_s.append(d_v_hat.detach().clone())
                    del v
                d_v_hat_s = torch.stack(d_v_hat_s, 0).permute(1, 0) # batch_size_curve x c
                gG = torch.einsum('nd,nd->n', d_v_hat_s, d_v_hat_s).sqrt()
                LG[net_idx] += gG.sum().item()
                torch.cuda.empty_cache()
                _idx += self.batch_size_curve
            network = network.cpu()
            torch.cuda.empty_cache()
        return np.mean(LG)

    def get_extrinsic_curvature(self):
        for net in self._networks:
            net.switch_norm('id')
        kappa = [0 for _ in range(len(self._networks))]
        for net_idx, network in enumerate(self._networks):
            network = network.cuda()
            network.train()
            network.zero_grad()
            _idx = 0
            while _idx < len(self.curve_input[net_idx]):
                output = network.forward_features(self.curve_input[net_idx][_idx:_idx+self.batch_size_curve])[1]
                output = output.reshape(output.size(0), -1)
                n, c = output.size()
                v_s = [] # 1st derivative
                a_s = [] # 2nd derivative
                for coord in range(c):
                    v = grad(output[:, coord].sum(), self.theta[net_idx], create_graph=True, retain_graph=True)[0][_idx:_idx+self.batch_size_curve] # batch size (of thetas)
                    a = grad(v.sum(), self.theta[net_idx], create_graph=True, retain_graph=True)[0][_idx:_idx+self.batch_size_curve] # batch size (of thetas)
                    v_s.append(v.detach().clone())
                    a_s.append(a.detach().clone())
                v_s = torch.stack(v_s, 0).permute(1, 0) # batch_size_curve x c
                a_s = torch.stack(a_s, 0).permute(1, 0) # batch_size_curve x c
                vv = torch.einsum('nd,nd->n', v_s, v_s)
                aa = torch.einsum('nd,nd->n', a_s, a_s)
                va = torch.einsum('nd,nd->n', v_s, a_s)
                kappa[net_idx] += (vv**(-3/2) * (vv * aa - va ** 2).sqrt()).sum().item()
                torch.cuda.empty_cache()
                _idx += self.batch_size_curve
            network = network.cpu()
            torch.cuda.empty_cache()
        return np.mean(kappa)

    def _update_bad_cases(self, reward_type, reward):
        # re-set "reward_type" of bad architectures to "reward"
        for _type in self._reward_types:
            for _idx, isbad in enumerate(self._buffers_bad):
                if isbad:
                    self._buffers[_type][_idx] = reward
            for _idx, isbad in enumerate(self._buffers_bad):
                if isbad:
                    self._buffers_change[_type][_idx] = (self._buffers[_type][_idx] - self._buffers[_type][_idx-1]) / (max(self._buffers[_type][max(0, _idx+1-self._buffer_length):_idx+1]) - min(self._buffers[_type][max(0, _idx+1-self._buffer_length):_idx+1]) + 1e-6)
                    if _idx + 1 < len(self._buffers_bad):
                        self._buffers_change[_type][_idx+1] = (self._buffers[_type][_idx+1] - self._buffers[_type][_idx]) / (max(self._buffers[_type][max(0, _idx+2-self._buffer_length):_idx+2]) - min(self._buffers[_type][max(0, _idx+2-self._buffer_length):_idx+2]) + 1e-6)

    def get_reward(self):
        #  changing range comparison ######
        _reward = _type = 0
        if len(self._buffers[self._reward_types[0]]) <= 1:
            # dummy reward for step 0
            return 0
        type_reward = [] # tuples of (type, reward)
        for _type in self._reward_types:
            var = self._buffers_change[_type][-1]
            type_reward.append((self.reward_type2index[_type], self._reward_sign[_type] * var))
        if 'constraint' in self._buffers and len(self._buffers['constraint']) > 0:
            var = self._buffers_change['constraint'][-1]
            type_reward.append((self.reward_type2index['constraint'], self._reward_sign['constraint'] * var * self.constraint_weight))
        if len(type_reward) > 0:
            _reward = sum([_r for _t, _r in type_reward])
        return _reward

    def _buffer_insert(self, results):
        if len(self._buffers[self._reward_types[0]]) == 0:
            self._buffers_bad.append(results['bad'])
            for _type in self._reward_types:
                self._buffers_change[_type].append(0)
                self._buffers[_type].append(results[_type])
            if 'constraint' in results:
                self._buffers_change['constraint'].append(0)
                self._buffers['constraint'].append(results['constraint'])
        else:
            if results['bad']:
                # set ntk of bad architecture as worst case in current buffer
                if 'ntk' in self._reward_types: results['ntk'] = max(self._buffers['ntk'])
            else:
                if 'ntk' in self._reward_types and results['ntk'] > max(self._buffers['ntk']):
                    self._update_bad_cases('ntk', results['ntk'])
            self._buffers_bad.append(results['bad'])
            for _type in self._reward_types:
                self._buffers[_type].append(results[_type])
                var = (self._buffers[_type][-1] - self._buffers[_type][-2]) / (max(self._buffers[_type][-self._buffer_length:]) - min(self._buffers[_type][-self._buffer_length:]) + 1e-6)
                self._buffers_change[_type].append(var)
            if 'constraint' in results:
                self._buffers['constraint'].append(results['constraint'])
                var = (self._buffers['constraint'][-1] - self._buffers['constraint'][-2]) / (max(self._buffers['constraint'][-self._buffer_length:]) - min(self._buffers['constraint'][-self._buffer_length:]) + 1e-6)
                self._buffers_change['constraint'].append(var)

    def get_ntk_exp(self):
        results = {}
        if 'ntk' in self._reward_types:
            ntk = self.get_ntk()
            results['ntk'] = ntk
            results['bad'] = ntk==-1 # networks of bad gradients
        if 'exp' in self._reward_types:
            exp = self.get_curve_complexity()
            results['exp'] = exp
            results['bad'] = False # networks of bad gradients
        torch.cuda.empty_cache()
        return results

    def step(self, network, constraint=None, verbose=False):
        self.set_network(network)
        results = self.get_ntk_exp()
        if constraint is not None:
            results['constraint'] = constraint
        self._buffer_insert(results)
        if verbose:
            print("NTK buffer:", self._buffers['ntk'][-self._buffer_length:])
            print("NTK change buffer:", self._buffers_change['ntk'][-self._buffer_length:])
            print("Exp buffer:", self._buffers['exp'][-self._buffer_length:])
            print("Exp change buffer:", self._buffers_change['exp'][-self._buffer_length:])
            if constraint is not None:
                print("Constraint buffer:", self._buffers['constraint'][-self._buffer_length:])
                print("Constraint change buffer:", self._buffers_change['constraint'][-self._buffer_length:])
        reward = self.get_reward()
        # reward larger the better
        return reward

    def _buffer_rank_best(self):
        # return the index of the best based on rankings over three buffers
        rankings = {}
        buffers_sorted = {}
        rankings_all = []
        for _type in self._reward_types:
            buffers_sorted[_type] = sorted(self._buffers[_type], reverse=self._reward_sign[_type]==1) # by default ascending
            num_samples = len(buffers_sorted[_type])
            rankings[_type] = [ buffers_sorted[_type].index(value) for value in self._buffers[_type] ]
        for _idx in range(num_samples):
            rankings_all.append(sum([ rankings[_type][_idx] for _type in rankings.keys() ]))
        return np.argmin(rankings_all)