import os
import json
import random
import pickle
import numpy as np
from copy import deepcopy

import cma
from svreg.optimizers import GAWrapper#, SofomoreWrapper
from svreg.nodes import SVNode
from svreg.tree import SVTree
from svreg.tree import MultiComponentTree as MCTree

import dask

class SVRegressor:
    """
    A class for running a genetic algorithm to perform symbolic
    regression using structure vectors.

    Attributes:

        settings (Settings):
            The settings used during regression. Use
            Settings.printValidSettings() to see valid settings options and
            suggested values.

        svNodePool (list):
            A list of SVNode objects representing the different types of SVs
            to be used during the regression.

        optimizer (object):
            The constructor for an optimizer object that will be used to
            optimize tree parameters. It is assumed that `optimizer` is a
            population-based optimizer and that it implements an ask() and
            tell() interface. An optimizer instance will be generated by calling
            optimizer(tree.populate(N=1)[0], **optimizerArgs) for tree in
            self.trees.

        optimizerArgs (dict):
            A dictionary of named arguments to be passed to the tree optimizers.

        trees (list):
            A list of SVTree objects being optimized.
    """

    def __init__(
            self,
            settings,
            database,
        ):

        # Note: it is assumed that Settings has already performed all validation
        # checks on the provided settings.

        self.settings = settings
        self.svNodePool = buildSVNodePool(database)

        self.structNames    = database.attrs['structNames']
        self.svNames        = database.attrs['svNames']
        self.elements       = database.attrs['elements']

        if settings['optimizer'] == 'CMA':
            self.optimizer = cma.CMAEvolutionStrategy
            self.optimizerArgs = [
                1.0,  # defaulted sigma0 value
                {
                    'verb_disp': 0,
                    'popsize': settings['optimizerPopSize'],
                    'maxiter': settings['maxNumOptimizerSteps'],
                    'tolx': 1e-8,  # changes in x-values
                    'tolfunhist': 1e-8,
                    'tolfun': 1e-8,
                    'tolfunrel': 1e-8,
                }
            ]
        elif settings['optimizer'] == 'GA':
            self.optimizer = GAWrapper
            self.optimizerArgs = [
                {
                    'verb_disp': 0,
                    'popsize': settings['optimizerPopSize'],
                    'maxiter': settings['maxNumOptimizerSteps'],
                    'tolx': 1e-8,  # changes in x-values
                    'tolfunhist': 1e-8,
                    'tolfun': 1e-8,
                    'tolfunrel': 1e-8,
                }
            ]
        # elif settings['optimizer'] == 'Sofomore':
        #     self.optimizer = SofomoreWrapper
        #     self.optimizerArgs = {
        #         'numStructs': numStructs,
        #         'paretoDimensionality': 2,
        #         'CMApopSize': settings['optimizerPopSize'],
        #         'SofomorePopSize': settings['numberOfTrees'],  # Placeholder
        #         'threads_per_node': settings['PROCS_PER_PHYS_NODE'],
        #         'threads_per_node': None,
        #     }
        else:
            raise NotImplementedError(
                'Must be one of `GA`, `CMA`, or `Sofomore`.'
            )

        self.trees = []
        self.populationDict = None

    def initializeTrees(self, elements):
        """Populates the GA with randomly-generated equation trees."""

        numElements = len(elements)

        if numElements < 1:
            raise RuntimeError("numElements must be >= 1 in initializeTrees()")

        uniqueTreeNames = []

        for ii, tree in enumerate(self.trees):
            tn = str(tree)
            if tn not in uniqueTreeNames:
                uniqueTreeNames.append(tn)
            else:
                del self.trees[ii]
                print("Removed duplicate tree: {}".format(tn))

        treesToAdd = []
        while len(uniqueTreeNames) < self.settings['numberOfTrees']:
            randTree = MCTree.random(
                svNodePool=self.svNodePool,
                maxDepth=random.randint(0, self.settings['maxTreeDepth']),
                elements=elements,
                allSums=self.settings['allSums']
            )

            if str(randTree) not in uniqueTreeNames:
                uniqueTreeNames.append(str(randTree))
                treesToAdd.append(randTree)

        self.trees += treesToAdd

    def evaluateTrees(self, svEng, svFcs, P, trueValues, useDask=True):
        """
        Updates the SVNode objects in the trees with the given values, then
        evaluate the trees

        Args:
            svEng (dict):
                svEng[structName][svName][elem] = computed values for given node

            svFcs (dict):
                svFcs[structName][svName][elem] = computed values for given node

            P (int):
                The number of parameter sets for each node. Used for splitting
                the population of results.

        Return:
            energies, forces (dict):
                {structName: [tree.eval() for tree in self.trees]}
        """

        structNames = list(svEng.keys())
        svNames     = list(svEng[structNames[0]].keys())

        # NOTE: elements must be sorted here to match assumed ordering of
        # tree.svNodes
        elements = sorted(list(svEng[structNames[0]][svNames[0]].keys()))

        # indexers[sv][el][i] = list of indices for svEng[*][sv][el] for tree i
        indexers = {}

        struct0 = structNames[0]
        # Build dictionary of indexers for each SV
        for svName in svEng[struct0]:
            indexers[svName] = {}
            for elem in elements:
                indexers[svName][elem]  = []

                counter = 0

                for tree in self.trees:
                    treeIndices = []
                    for svNode in tree.chemistryTrees[elem].svNodes:
                        # Only update the SVNode objects of the current type
                        if svNode.description == svName:
                            treeIndices.append(counter)
                            counter += 1

                    # Reverse list here since we'll be popping from it later
                    indexers[svName][elem].append(treeIndices[::-1])

        taskArgs = []
        for struct in structNames:
            indexCopy = deepcopy(indexers)

            for ii, tree in enumerate(self.trees):
                treeArgs = []
                for elem in elements:
                    for svNode in tree.chemistryTrees[elem].svNodes:
                        svName = svNode.description

                        engDot = svEng[struct][svName][elem]
                        fcsDot = svFcs[struct][svName][elem]

                        treeArgs.append(
                            (engDot, fcsDot, indexCopy[svName][elem][ii].pop())
                        )

                taskArgs.append(treeArgs)

        taskArgs = taskArgs[::-1]

        perTreeResults = []
        for structName in structNames:
            for t in self.trees:
                args = taskArgs.pop()

                if useDask:
                    perTreeResults.append(
                        dask.delayed(parseAndEval, pure=True, nout=2)(
                            pickle.dumps(t), args, P,
                            trueValues[structName]['forces'],
                            allSums=self.settings['allSums']
                        )
                    )
                else:
                    perTreeResults.append(
                        parseAndEval(
                            t, args, P,
                            trueValues[structName]['forces'],
                            allSums=self.settings['allSums']
                        )
                    )

        return perTreeResults

    def initializeOptimizers(self):
        import hashlib

        h5Hash = lambda t: hashlib.md5(str(t).encode()).hexdigest()

        path = os.path.join(self.settings['outputPath'], 'outcmaes', '{}/')

        # self.optimizers = [
        #     self.optimizer(
        #         tree.populate(N=1)[0],
        #         *.update(
        #             self.optimizerArgs[-1]
        #         )
        #     )
        #     for tree in self.trees
        # ]

        argsCopy = deepcopy(self.optimizerArgs)

        self.optimizers = []
        for tree in self.trees:
            d = {'verb_filenameprefix': path.format(h5Hash(tree))}
            d.update(self.optimizerArgs[-1])

            argsCopy[-1] = d

            self.optimizers.append(
                self.optimizer(
                    tree.populate(N=1)[0],
                    *argsCopy
                )
            )

    def tournament(self, topN):
        """
        Randomly return a random individual from the topN individuals in the
        population.

        Return:
            A deep copy (to avoid multiple trees pointing to the same nodes) of
            the best individual.
        """

        indices = np.arange(len(self.trees))
        costs = [t.cost for t in self.trees]

        argsort = np.argsort(costs)
        costs = costs[argsort]
        indices = indices[argsort]

        return deepcopy(self.trees[random.choice(indices[:topN])])

    def newIndividual(self):
        """
        Generates a new individual using the current population. Allow for
        random point mutation.
        """

        newTree = self.tournament(self.settings['tournamentSize'])
        donor   = self.tournament(self.settings['tournamentSize'])

        newTree.crossover(donor)

        if random.random() < self.settings['pointMutateProb']:
            newTree.pointMutate(
                self.svNodePool, self.settings['pointMutateProb']
            )

        newTree.updateSVNodes()

        return newTree

    def printTop10Header(self, regStep):
        print(regStep, flush=True)

        for treeNum, t in enumerate(self.trees):
            print(treeNum, t)

        print('\t\t\t', ''.join(['{:<10}'.format(i) for i in range(10)]))

    def generatePopulationDict(self, N):
        """
        Generates N parameter sets for each tree, then groups them by SV name
        for easy batch evaluation.

        Returns:
            populationDict: {svName: {el: stacked population}}
            rawPopulation: the parameter arrays generated by each optimizer
        """

        rawPopulations = [np.array(opt.ask(N)) for opt in self.optimizers]

        # Used for parsing later
        self.numNodes = {}

        # Now parse the populations and group them by SV type

        # populationDict[svName][elem]
        populationDict = {}
        for pop, tree in zip(rawPopulations, self.trees):
            # popDict= {el: {svName: population}}
            popDict = tree.parseArr2Dict(pop)

            for elem in popDict:
                for svName in popDict[elem]:
                    if svName not in populationDict:
                        populationDict[svName] = {}

                    if elem not in populationDict[svName]:
                        populationDict[svName][elem] = []

                    populationDict[svName][elem].append(popDict[elem][svName])

        # Count the number of each node to help with parsing results later
        for tree in self.trees:
            for elem in tree.chemistryTrees:
                for svNode in tree.chemistryTrees[elem].svNodes:
                    svName = svNode.description

                    if svName not in self.numNodes:
                        self.numNodes[svName] = {}

                    if elem not in self.numNodes[svName]:
                        self.numNodes[svName][elem] = 0

                    self.numNodes[svName][elem] += 1

        self.chunks = {}
        # Stack each group
        for svName in populationDict:
            if svName not in self.chunks:
                self.chunks[svName] = {}

            for elem, popList in populationDict[svName].items():
                dat = np.concatenate(popList, axis=0).T

                populationDict[svName][elem] = dat.astype('float32')

        return populationDict, rawPopulations

    def updateOptimizers(self, rawPopulations, costs, penalties):
        for treeIdx in range(len(self.optimizers)):
            fullCost = costs[treeIdx] + penalties[treeIdx]

            opt = self.optimizers[treeIdx]
            opt.tell(rawPopulations[treeIdx], fullCost)
            opt.logger.add()

    def checkStale(self):
        """
        Returns a list of the indices of any trees that have finished
        optimizing.
        """

        stale = []
        messages = []

        for i, opt in enumerate(self.optimizers):
            if opt.stop():
                stale.append(i)
                messages.append(opt.stop())

        return stale, messages

def buildSVNodePool(database):
    """Prepare svNodePool for use in tree construction"""

    svNodePool = []

    for svName in database:

        restrictions = None
        if 'restrictions' in database.attrs[svName]:
            restrictions = []
            resList = database.attrs[svName]['restrictions'].tolist()[::-1]
            for num in database.attrs[svName]['numRestrictions']:
                tmp = []
                for _ in range(num):
                    tmp.append(tuple(resList.pop()))
                restrictions.append(tmp)

        bondComps = sorted(set(database.attrs[svName]['components']))
        numParams = []
        restr = []

        cList = database.attrs[svName]['components'].tolist()
        for c in bondComps:
            idx = cList.index(c)
            numParams.append(database.attrs[svName]['numParams'][idx].astype(int))
            restr.append(restrictions[idx])

        if 'paramRanges' in database.attrs[svName]:
            pRanges = []
            for c in bondComps:
                idx = cList.index(c)
                pRanges.append(database.attrs[svName]['paramRanges'][idx])
        else:
            pRanges = None

        bondComps = [c.decode('utf-8') for c in bondComps]

        svNodePool.append(
            SVNode(
                description=svName,
                components=bondComps,
                constructor=[
                    c.decode('utf-8')
                    for c in database.attrs[svName]['components']
                ],
                numParams=numParams,
                restrictions=restr,
                paramRanges=pRanges,
                inputTypes=json.loads(
                    database.attrs[svName]['inputTypes'].decode('utf-8').replace("'", '"')
                )
            )
        )

    return svNodePool

def parseAndEval(
    *args,
    **kwargs,
    ):

    tree            = kwargs['tree']
    listOfIndices   = kwargs['listOfIndices']
    P               = kwargs['P']
    tvF             = kwargs['tvF']
    numChunks       = kwargs['numChunks']
    allSums         = kwargs['allSums']

    import pickle
    tree = pickle.loads(tree)

    nodeCounter     = 0
    chunkCounter    = 0
    for elem in tree.elements:
        for svNode in tree.chemistryTrees[elem].svNodes:
            svName = svNode.description

            chunkTup = []
            for _ in range(numChunks[svName][elem]):
                chunkTup.append(args[chunkCounter])
                chunkCounter += 1

            idx = listOfIndices[nodeCounter]
            nodeCounter += 1

            eng = np.concatenate([c[0] for c in chunkTup], axis=-1)
            fcs = np.concatenate([c[1] for c in chunkTup], axis=-1)

            Ne = eng.shape[0]
            Nn = eng.shape[1] // P

            eng = eng.reshape((Ne, Nn, P))
            eng = np.moveaxis(eng, 1, 0)
            eng = np.moveaxis(eng, -1, 1)

            # fcs shape: (Ne, Na, 3, P*Nn)

            if allSums:
                Na = fcs.shape[0]
                fcs = fcs.reshape(Na, 3, Nn, P)
            else:
                Na = fcs.shape[1]
                fcs = fcs.reshape(Ne, Na, 3, Nn, P)

            fcs = np.moveaxis(fcs, -2, 0)
            fcs = np.moveaxis(fcs, -1, 1)

            # fcs shape: (Nn, P, Ne, Na, 3)

            svNode.values = (eng[idx], fcs[idx])

    engResult, fcsResult = tree.eval(useDask=False, allSums=allSums)

    fcsErrors = np.average(abs(sum(fcsResult) - tvF['forces']), axis=(1,2))

    return sum(engResult), fcsErrors


###############################################################
#
# Cantera UPF : freely-propagating premixed flames
# Description : perform calculations of UPF for a range of eq. ratio
#               and compute global indices (integrated fuel, CO, NO consumption, ...)
# Author : A. Felden
# Last modified : 02/2018
#
###############################################################

#import :
from cantera import *
import numpy as np
import csv
import sys,os

#################################################################
# Prepare your run
#################################################################
#Mechanism used for the process
cas         = 'DET.cti'
pref        = 'C2H4DET_T298_P1'
restore     = 'no'
init_flame  = 'RedHighP_HighP_save.xml'

#General parameter values :
p          	= 101325.0                # pressure
tin        	= 298.0        # unburned gas temperature
phi_min 	= 0.8
phi_max 	= 1.5
npoints 	= 1

########
#STORAGE
########
# KSI
ksi = np.zeros(npoints,'d')
# YC eq
Yc_equil = np.zeros(npoints,'d')
#phiv
phi = np.zeros(npoints,'d')
#sl0phiv
sl_phi = np.zeros(npoints,'d')
#tadphiv
t_phi = np.zeros(npoints,'d')
#omega0phiv
#delta0phiv
dl0_phi = np.zeros(npoints,'d')

#########################################
# SPECS : profiles, wdot, consump, val max
# prof  : Fuel, oxy,  CO, CO2, OH, C2H2
# max   :       oxy,  CO, CO2, OH, C2H2
# wdot  : Fuel, oxy,  CO,          C2H2
# cons  : Fuel,       CO,          C2H2
#########################################
#FUEL consump
Fuel_c 		= np.zeros(npoints,'d')
#OXY consump
O2m_phi 	= np.zeros(npoints,'d')
#CO consump
COm_phi 	= np.zeros(npoints,'d')
CO_c 		= np.zeros(npoints,'d')
#CO2 consump
CO2m_phi 	= np.zeros(npoints,'d')
#OH consump
OHm_phi 	= np.zeros(npoints,'d')

#################
#Run parameters:

#Initial grids, chosen to be 15cm long :
initial_grid = np.array([0.0, 0.001, 0.005, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1], 'd')/0.1 #, 0.11, 0.115, 0.119, 0.12],'d')/0.1 # m

#Tolerance properties
tol_ss    = [1.0e-8, 1.0e-12]        # [rtol atol] for steady-state problem
tol_ts    = [1.0e-8, 1.0e-12]        # [rtol atol] for time stepping

loglevel  = 1                       # amount of diagnostic output

refine_grid = True                  # True to enable refinement

#################
#Stoechiometry :

fuel_species = 'C2H4'
stoich_O2 = 3.0
air_N2_O2_ratio = 3.76

gas  = Solution(cas)
m=gas.n_species

ifuel = gas.species_index(fuel_species)
ioh = gas.species_index('OH')
ih2 = gas.species_index('H2')
io2 = gas.species_index('O2')
in2 = gas.species_index('N2')
ico = gas.species_index('CO')
ico2 = gas.species_index('CO2')
ih2o = gas.species_index('H2O')

#################################################################
#BILGER PRE PROC :
#################################################################
#Mixt frac calculation based on Bilger
#molecular weight of CHON
C_W = 0.012
H_W = 0.001
O_W = 0.016
N_W = 0.014
#Bilger coefs pour formule
Bilger_coefs 		= np.zeros(4,'d')
Bilger_coefs[0]		= 2.0 / C_W
Bilger_coefs[1]		= 1.0 / ( 2.0 * H_W )
Bilger_coefs[2]		= -1.0 / O_W
Bilger_coefs[3]		= 0.0
#Storage nb d atomes par especes
Nb_C			= np.zeros(m,'d')
Nb_H			= np.zeros(m,'d')
Nb_O			= np.zeros(m,'d')
#Compo du fuel et oxi tank
Ftank			= np.zeros(m,'d')
Ftank[ifuel]		= 1.0
Oxtank			= np.zeros(m,'d')
Oxtank[io2]		= 0.233
Oxtank[in2]		= 0.767
#Normalisation
BetaF 			= 0.0
BetaOx 			= 0.0
for i, spec in enumerate(gas.species_names):
        Nb_C[i] = gas.n_atoms(spec,'C')
        Nb_H[i] = gas.n_atoms(spec,'H')
        Nb_O[i] = gas.n_atoms(spec,'O')
        Weightedatom = Bilger_coefs[0] * C_W * Nb_C[i] + Bilger_coefs[1] * H_W * Nb_H[i] + Bilger_coefs[2] * O_W * Nb_O[i]
        BetaF = BetaF + 1000.0 * Weightedatom * Ftank[i] / gas.molecular_weights[i]
        BetaOx = BetaOx + 1000.0 * Weightedatom * Oxtank[i] / gas.molecular_weights[i]

if (restore == 'yes'):
#################################################################
# RESTORE PREVIOUS FLAME
#################################################################
    phi[0]   = phi_min
    x        = np.zeros(m,'d')
    x[ifuel] = phi[0]
    x[io2]   = stoich_O2
    x[in2]   = stoich_O2*air_N2_O2_ratio

    gas.TPX = tin, p, x

    f = FreeFlame(gas, initial_grid)
    f.restore(init_flame, 'phi_'+str(phi[0]))
else:
#################################################################
#FIRST SIMULATION :
#################################################################
#Stoechiometry :
    phi[0]   = phi_min
    x        = np.zeros(m,'d')
    x[ifuel] = phi[0]
    x[io2]   = stoich_O2
    x[in2]   = stoich_O2*air_N2_O2_ratio

    print('  ')
    print(' ### First flame at Phi : '+str(phi[0]))

    gas.TPX = tin, p, x

#Create the free laminar premixed flame
    f = FreeFlame(gas, initial_grid)

    f.flame.set_steady_tolerances(default=tol_ss)
    f.flame.set_transient_tolerances(default=tol_ts)

    f.inlet.X = x
    f.inlet.T = tin

    f.transport_model = 'Mix'

#First flame:

      #No energy for starters
    f.energy_enabled = False

      #Refinement criteria
    f.set_refine_criteria(ratio = 5.0, slope = 1, curve = 1)

      #Max number of times the Jacobian will be used before it must be re-evaluated
    f.set_max_jac_age(30, 30)

      #Set time steps whenever Newton convergence fails
    f.set_time_step(1.0e-6, [5, 10, 20]) #s

      #Calculation
    f.solve(loglevel, refine_grid)

#################
#Second flame:

	#Energy equation enabled
    f.energy_enabled = True

      #Refinement criteria when energy equation is enabled
    f.set_refine_criteria(ratio = 5.0, slope = 0.5, curve = 0.5)

      #Calculation and save of the results
    f.solve(loglevel, refine_grid)

      #Refinement criteria when energy equation is enabled
    f.set_refine_criteria(ratio = 3.0, slope = 0.2, curve = 0.2)

      #Calculation and save of the results
    f.solve(loglevel, refine_grid)

#################
#Third flame and so on ...:

	#Refinement criteria should be changed ...
    f.set_refine_criteria(ratio = 2.0, slope = 0.04, curve = 0.04)

    f.solve(loglevel, refine_grid)

	#Refinement criteria should be changed ...
    f.set_refine_criteria(ratio = 2.0, slope = 0.02, curve = 0.02)

    f.solve(loglevel, refine_grid)

    sl_phi[0] = f.u[0]
    t_phi[0]  = f.T[f.flame.n_points-1]
    print( "Sl ? ",  sl_phi[0])
    print( "Tad ? ", t_phi[0])

	#...and saving of the results
    #f.save(str(pref)+'_save.xml','phi_'+str(phi[0]))
    f.write_csv(str(pref)+'_Phi'+str(phi[0])+'_Y.csv', species='Y', quiet=False)


from torch.optim.optimizer import Optimizer, required
import torch
import math
import numpy as np

class LRangerMod(Optimizer):

    # AMSGrad/Adam + AdaMod + QH Momentum + Iterate Averaging + Lookahead + Rule of Thumb Linear Warmup (instead of RAdam Rectification) + P from PAdam

    def __init__(self, params, lr=1e-3,
                 betas=(0.999, 0.999, 0.999),
                 nus=(0.7, 1.0),
                 p=0.5,
                 eps=1e-8,
                 k=5,
                 alpha=0.8,
                 amsgrad=True,
                 AdaMod=True,
                 warmup=True,
                 AdaMod_bias_correct=True,
                 IA=True,
                 use_gc=False,
                 IA_cycle=1000,
                 epochs=100,
                 step_per_epoch=None,
                 weight_decay=0):

        # betas = (beta1 for first order moments, beta2 for second order moments, beta3 for ema over adaptive learning rates (AdaMod))
        # nus = (nu1,nu2) (for quasi hyperbolic momentum)
        # eps = small value for numerical stability (avoid divide by zero)
        # k = lookahead cycle
        # alpha = outer learning rate (lookahead)
        # amsgrad = bool to decide whether to use amsgrad instead of adam as the core optimizer
        # AdaMod_bias_correct = bool to decide whether to add bias correction to AdaMod
        # IA = bool to decide if Iterate Averaging is ever going to be used
        # IA_cycle = Iterate Averaging Cycle (Recommended to initialize with no. of iterations in Epoch) (doesn't matter if you are not using IA)
        # epochs = No. of epochs you plan to use (Only relevant if using DEMON)
        # step_per_epoch = No. of iterations in an epoch (only relevant if using DEMON)
        # weight decay = decorrelated weight decay value

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= p <= 0.5:
            raise ValueError("Invalid p value: {}".format(p))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        if not 0.0 <= betas[2] < 1.0:
            raise ValueError("Invalid beta parameter at index 2: {}".format(betas[2]))
        if not 0.0 <= nus[0] <= 1.0:
            raise ValueError("Invalid nu parameter at index 0: {}".format(nus[0]))
        if not 0.0 <= nus[1] <= 1.0:
            raise ValueError("Invalid nu parameter at index 1: {}".format(nus[1]))
        if not 0.0 <= alpha <= 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.k = k
        self.epochs = epochs
        self.amsgrad = amsgrad
        self.warmup_period = 2/(1-betas[1])
        self.IA_cycle = IA_cycle
        self.IA = IA
        self.AdaMod = AdaMod
        self.use_gc = use_gc
        self.AdaMod_bias_correct = AdaMod_bias_correct
        self.warmup = warmup
        if step_per_epoch is None:
            self.step_per_epoch = IA_cycle
        else:
            self.step_per_epoch = step_per_epoch

        self.T = self.epochs*self.step_per_epoch

        defaults = dict(lr=lr,
                        betas=betas,
                        nus=nus,
                        eps=eps,
                        p=p,
                        alpha=alpha,
                        weight_decay=weight_decay)
        super(LRangerMod, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(LRangerMod, self).__setstate__(state)

    def apply_AdaMod(self, beta3, n_avg, n, step):
        n_avg.mul_(beta3).add_(1 - beta3, n)
        if self.AdaMod_bias_correct:
            n_avg_ = n_avg.clone()
            n_avg_.div_(1 - (beta3 ** step))
            torch.min(n, n_avg_, out=n)
        else:
            torch.min(n, n_avg, out=n)
        return n

    def step(self, activate_IA=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('LRangerMod does not support sparse gradients')

                state = self.state[p]

                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)
                    state['num_models'] = 0
                    state['cached_params'] = p.data.clone()
                    if self.amsgrad:
                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
                    if self.AdaMod:
                        state['n_avg'] = torch.zeros_like(p.data)

                state['step'] += 1

                w = min([1.0, state['step']/self.warmup_period])
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2, beta3 = group['betas']
                nu1, nu2 = group['nus']

                if self.warmup:
                    lr = w*group['lr']
                else:
                    lr = group['lr']

                wd = group['weight_decay']
                alpha = group['alpha']

                do_IA = False
                lookahead_step = False

                if self.IA and activate_IA:
                    lookahead_step = False
                    if state['step'] % self.IA_cycle == 0:
                        do_IA = True
                elif self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                momentum = exp_avg.clone()
                momentum.div_(1 - (beta1 ** state['step'])).mul_(nu1).add_(1-nu1, grad)

                if wd != 0:
                    p.data.add_(-wd*lr, p.data)

                beta2_t = beta2 ** state['step']

                if self.amsgrad and state['step'] > 1:
                    max_exp_avg_sq = state['max_exp_avg_sq']
                    # Maintains the maximum of all 2nd moment running avg. till now
                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    vt = max_exp_avg_sq.clone()
                else:
                    vt = exp_avg_sq.clone()

                bias_correction2 = 1 - beta2_t
                vt.div_(bias_correction2)
                if nu2 != 1.0:
                    vt.mul_(nu2).addcmul_(1-nu2, grad, grad)
                denom = vt.pow_(group['p']).add_(group['eps'])
                n = lr/denom
                if self.AdaMod:
                    n_avg = state['n_avg']
                    n = self.apply_AdaMod(beta3, n_avg, n, step=state['step'])

                p.data.add_(-n*momentum)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

                if do_IA:
                    p.data.add_(state["num_models"], state['cached_params']
                                ).div_(state["num_models"]+1.0)
                    state['cached_params'].copy_(p.data)
                    state["num_models"] += 1

        return loss

class DemonRanger(Optimizer):

    # Rectified-AMSGrad/RAdam + AdaMod + QH Momentum + Iterat Averaging + Lookahead + DEMON (decaying Momentum) + gradient centralization + grad noise

    def __init__(self, params, lr=1e-3,
                 betas=(0.999, 0.999, 0.999),
                 nus=(0.7, 1.0),
                 eps=1e-8,
                 k=5,
                 alpha=0.8,
                 gamma=0.55,
                 use_demon=True,
                 rectify=True,
                 amsgrad=True,
                 AdaMod=True,
                 AdaMod_bias_correct=True,
                 IA=True,
                 IA_cycle=1000,
                 epochs=100,
                 step_per_epoch=None,
                 weight_decay=0,
                 use_gc=True,
                 use_grad_noise=False):

        # betas = (beta1 for first order moments, beta2 for second order moments, beta3 for ema over adaptive learning rates (AdaMod))
        # nus = (nu1,nu2) (for quasi hyperbolic momentum)
        # eps = small value for numerical stability (avoid divide by zero)
        # k = lookahead cycle
        # alpha = outer learning rate (lookahead)
        # gamma = gradient noise control parameter (for regularization)
        # use_demon = bool to decide whether to use DEMON (Decaying Momentum) or not
        # rectify = bool to decide whether to apply the recitification term (from RAdam) or not
        # amsgrad = bool to decide whether to use amsgrad instead of adam as the core optimizer
        # AdaMod_bias_correct = bool to decide whether to add bias correction to AdaMod
        # IA = bool to decide if Iterate Averaging is ever going to be used
        # IA_cycle = Iterate Averaging Cycle (Recommended to initialize with no. of iterations in Epoch) (doesn't matter if you are not using IA)
        # epochs = No. of epochs you plan to use (Only relevant if using DEMON)
        # step_per_epoch = No. of iterations in an epoch (only relevant if using DEMON)
        # weight decay = decorrelated weight decay value
        # use_gc = bool to determine whether to use gradient centralization or not.
        # use_grad_noise = bool to determine whether to use gradient noise or not.

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        if not 0.0 <= betas[2] < 1.0:
            raise ValueError("Invalid beta parameter at index 2: {}".format(betas[2]))
        if not 0.0 <= nus[0] <= 1.0:
            raise ValueError("Invalid nu parameter at index 0: {}".format(nus[0]))
        if not 0.0 <= nus[1] <= 1.0:
            raise ValueError("Invalid nu parameter at index 1: {}".format(nus[1]))
        if not 0.0 <= alpha <= 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.use_gc = use_gc
        self.use_grad_noise = use_grad_noise
        self.k = k
        self.epochs = epochs
        self.amsgrad = amsgrad
        self.use_demon = use_demon
        self.IA_cycle = IA_cycle
        self.IA = IA
        self.rectify = rectify
        self.AdaMod = AdaMod
        self.AdaMod_bias_correct = AdaMod_bias_correct
        if step_per_epoch is None:
            self.step_per_epoch = IA_cycle
        else:
            self.step_per_epoch = step_per_epoch

        self.T = self.epochs*self.step_per_epoch

        defaults = dict(lr=lr,
                        betas=betas,
                        nus=nus,
                        eps=eps,
                        alpha=alpha,
                        gamma=gamma,
                        weight_decay=weight_decay)
        super(DemonRanger, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(DemonRanger, self).__setstate__(state)

    def apply_AdaMod(self, beta3, n_avg, n, step):
        n_avg.mul_(beta3).add_(1 - beta3, n)
        if self.AdaMod_bias_correct:
            n_avg_ = n_avg.clone()
            n_avg_.div_(1 - (beta3 ** step))
            torch.min(n, n_avg_, out=n)
        else:
            torch.min(n, n_avg, out=n)
        return n

    def step(self, activate_IA=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('DemonRanger does not support sparse gradients')

                state = self.state[p]

                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)
                    state['num_models'] = 0
                    state['cached_params'] = p.data.clone()
                    if self.amsgrad:
                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
                    if self.AdaMod:
                        state['n_avg'] = torch.zeros_like(p.data)

                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1_init, beta2, beta3 = group['betas']
                rho_inf = (2/(1-beta2)) - 1
                nu1, nu2 = group['nus']
                lr = group['lr']
                wd = group['weight_decay']
                alpha = group['alpha']
                gamma = group['gamma']

                do_IA = False
                lookahead_step = False

                if self.IA and activate_IA:
                    lookahead_step = False
                    if state['step'] % self.IA_cycle == 0:
                        do_IA = True
                elif self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if self.use_demon:
                    temp = 1-(state['step']/self.T)
                    beta1 = beta1_init * temp / ((1-beta1_init)+beta1_init*temp)
                else:
                    beta1 = beta1_init

                if self.use_grad_noise:
                    grad_var = lr/((1+state['step'])**gamma)
                    grad_noise = torch.empty_like(grad).normal_(mean=0.0, std=math.sqrt(grad_var))
                    grad.add_(grad_noise)

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                momentum = exp_avg.clone()
                momentum.div_(1 - (beta1 ** state['step'])).mul_(nu1).add_(1-nu1, grad)

                if wd != 0:
                    p.data.add_(-wd*lr, p.data)

                beta2_t = beta2 ** state['step']

                if self.amsgrad and state['step'] > 1:
                    max_exp_avg_sq = state['max_exp_avg_sq']
                    # Maintains the maximum of all 2nd moment running avg. till now
                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    vt = max_exp_avg_sq.clone()
                else:
                    vt = exp_avg_sq.clone()

                if self.rectify:
                    rho_t = rho_inf - 2 * state['step'] * beta2_t / (1 - beta2_t)

                    # more conservative since it's an approximated value
                    if rho_t >= 5:
                        R = math.sqrt(((rho_t-4)*(rho_t-2)*rho_inf) /
                                      ((rho_inf-4)*(rho_inf-2)*rho_t))
                        bias_correction2 = 1 - beta2_t
                        vt.div_(bias_correction2)
                        if nu2 != 1.0:
                            vt.mul_(nu2).addcmul_(1-nu2, grad, grad)
                        denom = vt.sqrt_().add_(group['eps'])

                        n = (lr*R)/denom

                        if self.AdaMod:
                            n_avg = state['n_avg']
                            n = self.apply_AdaMod(beta3, n_avg, n, step=state['step'])

                        p.data.add_(-n*momentum)
                    else:
                        if self.AdaMod:
                            n_avg = state['n_avg']
                            n_avg.mul_(beta3).add_(1 - beta3, lr)
                        p.data.add_(-lr, momentum)
                else:
                    bias_correction2 = 1 - beta2_t
                    vt.div_(bias_correction2)
                    if nu2 != 1.0:
                        vt.mul_(nu2).addcmul_(1-nu2, grad, grad)
                    denom = vt.sqrt_().add_(group['eps'])
                    n = lr/denom
                    if self.AdaMod:
                        n_avg = state['n_avg']
                        n = self.apply_AdaMod(beta3, n_avg, n, step=state['step'])

                    p.data.add_(-n*momentum)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

                if do_IA:
                    p.data.add_(state["num_models"], state['cached_params']
                                ).div_(state["num_models"]+1.0)
                    state['cached_params'].copy_(p.data)
                    state["num_models"] += 1

        return loss

class HyperRanger(Optimizer):

    # Nostalgic PAdam + QH Momentum + Iterate Averaging + Lookahead + DEMON (decaying Momentum) + gradient centralization + hypergradient descent on lr and nu1

    def __init__(self, params, lr=1e-3,
                 betas=(0.999, 0.999),
                 nus=(0.7, 1.0),
                 eps=1e-8,
                 gamma=0.0001,
                 nostalgia=True,
                 use_demon=True,
                 hypergrad_lr=1e-7,
                 HDM=False,
                 hypertune_nu1=False,
                 p=0.25,
                 k=5,
                 alpha=0.8,
                 IA=True,
                 IA_cycle=1000,
                 epochs=100,
                 step_per_epoch=None,
                 weight_decay=0,
                 use_gc=True):

        # betas = (beta1 for first order moments, beta2 for second order moments)
        # nus = (nu1,nu2) (for quasi hyperbolic momentum)
        # eps = small value for numerical stability (avoid divide by zero)
        # k = lookahead cycle
        # alpha = outer learning rate (lookahead)
        # gamma = used for nostalgia
        # nostalgia = bool to decide whether to use nostalgia (from Nostalgic Adam or NosAdam)
        # use_demon = bool to decide whether to use DEMON (Decaying Momentum) or not
        # hypergrad_lr = learning rate for updating hyperparameters (like lr) through hypergradient descent (probably need to increase around 0.02 if HDM is True). Set to 0.0 to disable hypergradient descent.
        # HDM = bool to decide whether to use Multiplicative rule for updating hyperparameters or not
        # hypertune_nu1 = bool to decide whether apply hypergradient descent on nu1 as well or not.
        # p = p from PAdam
        # IA = bool to decide if Iterate Averaging is ever going to be used
        # IA_cycle = Iterate Averaging Cycle (Recommended to initialize with no. of iterations in Epoch) (doesn't matter if you are not using IA)
        # epochs = No. of epochs you plan to use (Only relevant if using DEMON)
        # step_per_epoch = No. of iterations in an epoch (only relevant if using DEMON)
        # weight decay = decorrelated weight decay value
        # use_gc = bool to determine whether to use gradient centralization or not.
        # use_grad_noise = bool to determine whether to use gradient noise or not.

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= hypergrad_lr:
            raise ValueError("Invalid hypergradient learning rate: {}".format(hypergrad_lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        if not 0.0 <= nus[0] <= 1.0:
            raise ValueError("Invalid nu parameter at index 0: {}".format(nus[0]))
        if not 0.0 <= nus[1] <= 1.0:
            raise ValueError("Invalid nu parameter at index 1: {}".format(nus[1]))
        if not 0.0 <= p <= 0.5:
            raise ValueError("Invalid p parameter: {}".format(p))
        if not 0.0 <= alpha < 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.nostalgia = nostalgia
        self.use_demon = use_demon
        self.k = k
        self.IA = IA
        self.IA_cycle = IA_cycle
        self.epochs = epochs
        if step_per_epoch is None:
            self.step_per_epoch = IA_cycle
        else:
            self.step_per_epoch = step_per_epoch
        self.use_gc = use_gc
        self.T = self.epochs*self.step_per_epoch
        self.hypertune_nu1 = hypertune_nu1
        self.HDM = HDM

        defaults = dict(lr=lr,
                        betas=betas,
                        nu1=nus[0],
                        nu2=nus[1],
                        eps=eps,
                        alpha=alpha,
                        gamma=gamma,
                        p=p,
                        hypergrad_lr=hypergrad_lr,
                        weight_decay=weight_decay)
        super(HyperRanger, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(HyperRanger, self).__setstate__(state)

    def step(self, activate_IA=False, display=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('HyperRanger does not support sparse gradients')

                state = self.state[p]

                hypergrad_lr = group['hypergrad_lr']
                beta1_init, beta2 = group['betas']
                wd = group['weight_decay']
                alpha = group['alpha']
                gamma = group['gamma']

                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)
                    state['lr'] = group['lr']

                    if self.IA:
                        state['num_models'] = 0
                    if self.IA or (self.k > 0):
                        state['cached_params'] = p.data.clone()
                    if self.nostalgia:
                        state['B_old'] = 0
                        state['B_new'] = 1
                    if hypergrad_lr > 0.0:
                        state['nu'] = group['nu']
                        state['prev_lr_grad'] = torch.zeros_like(grad.view(-1))
                        if self.hypertune_nu1:
                            state['prev_nu_grad'] = torch.zeros_like(grad.view(-1))

                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']

                if self.use_demon:
                    temp = 1-(state['step']/self.T)
                    beta1 = beta1_init * temp / ((1-beta1_init)+beta1_init*temp)
                else:
                    beta1 = beta1_init

                if self.nostalgia:
                    beta2 = state['B_old']/state['B_new']
                    state['B_old'] += math.pow(state['step'], -gamma)
                    state['B_new'] += math.pow(state['step']+1, -gamma)

                do_IA = False
                lookahead_step = False

                if self.IA and activate_IA:
                    lookahead_step = False
                    if state['step'] % self.IA_cycle == 0:
                        do_IA = True
                elif self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if state['step'] > 1 and hypergrad_lr > 0.0:
                    prev_lr_grad = state['prev_lr_grad']
                    h = torch.dot(grad.view(-1), prev_lr_grad)

                    if self.HDM:
                        grad_norm = grad.view(-1).norm()
                        norm_denom = grad_norm*(prev_lr_grad.norm())
                        norm_denom.add_(group['eps'])
                        state['lr'] = state['lr']*(1-hypergrad_lr*(h/norm_denom))
                    else:
                        state['lr'] -= hypergrad_lr * h

                    torch.max(state['lr'], torch.zeros_like(state['lr']), out=state['lr'])

                    if display:
                        print("lr", state['lr'])

                    if self.hypertune_nu1:
                        prev_nu_grad = state['prev_nu_grad']
                        h = torch.dot(grad.view(-1), prev_nu_grad)
                        if self.HDM:
                            norm_denom = grad_norm*(prev_nu_grad.norm())
                            norm_denom.add_(group['eps'])
                            state['nu1'] = state['nu1']*(1-hypergrad_lr*(h/norm_denom))
                        else:
                            state['nu1'] -= hypergrad_lr * h

                        torch.max(state['nu1'], torch.zeros_like(state['nu1']), out=state['nu1'])
                        torch.min(state['nu1'], torch.ones_like(state['nu1']), out=state['nu1'])

                    if display:
                        print("nu", state['nu1'])

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                nu1 = state['nu1']
                nu2 = state['nu2']
                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                momentum = exp_avg.clone()
                bias_correction1 = 1 - (beta1 ** state['step'])
                momentum.div_(bias_correction1)

                vt = exp_avg_sq.clone()

                if not self.nostalgia:
                    vt.div_(1 - (beta2 ** state['step']))
                if nu2 != 1.0:
                    vt.mul_(nu2).addcmul_(1-nu2, grad, grad)

                denom = vt.pow_(group['p']).add_(group['eps'])

                n = state['lr']/denom

                if lookahead_step:
                    dalpha = alpha
                elif do_IA:
                    dalpha = (1/(state["num_models"]+1.0))
                else:
                    dalpha = 1.0

                if hypergrad_lr > 0.0 and self.hypertune_nu1:
                    state['prev_nu_grad'] = (-dalpha*n*(momentum - grad)).view(-1)

                momentum.mul_(nu1).add_(1-nu1, grad)  # quasi hyperbolic momentum

                if hypergrad_lr > 0.0:
                    temp = dalpha*(-(momentum/denom) - wd*p.data)
                    state['prev_lr_grad'] = temp.view(-1)

                p.data.add_(-n*momentum)

                if wd != 0:
                    p.data.add_(-wd*state['lr'], p.data)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

                if do_IA:
                    p.data.add_(state["num_models"], state['cached_params']
                                ).div_(state["num_models"]+1.0)
                    state['cached_params'].copy_(p.data)
                    state["num_models"] += 1

        return loss

class HyperRangerMod(Optimizer):

    # Different from HyperRanger integrates AdaMod, and hypergradient descent through it. Slower, however.
    # doesn't have hypertunability for nu1 though.

    def __init__(self, params, lr=1e-3,
                 betas=(0.999, 0.999, 0.999),
                 nus=(0.7, 1.0),
                 eps=1e-8,
                 AdaMod_bias_correct=True,
                 gamma=0.0001,
                 nostalgia=True,
                 use_demon=True,
                 hypergrad_lr=1e-7,
                 p=0.5,
                 k=5,
                 alpha=0.8,
                 IA=True,
                 IA_cycle=1000,
                 epochs=100,
                 step_per_epoch=None,
                 weight_decay=0,
                 use_gc=True):

        # betas = (beta1 for first order moments, beta2 for second order moments, beta3 for AdaMod) # set beta3 = 0 to disable AdaMod
        # nus = (nu1,nu2) (for quasi hyperbolic momentum)
        # eps = small value for numerical stability (avoid divide by zero)
        # AdaMod_bias_correct = bool to determine whether to apply bias correction on AdaMod or not
        # k = lookahead cycle
        # alpha = outer learning rate (lookahead)
        # gamma = used for nostalgia
        # nostalgia = bool to decide whether to use nostalgia (from Nostalgic Adam or NosAdam)
        # use_demon = bool to decide whether to use DEMON (Decaying Momentum) or not
        # hypergrad_lr = learning rate for updating hyperparameters (like lr) through hypergradient descent (probably need to increase around 0.02 if HDM is True). Set to 0.0 to disable hypergradient descent.
        # HDM = bool to decide whether to use Multiplicative rule for updating hyperparameters or not
        # hypertune_nu1 = bool to decide whether apply hypergradient descent on nu1 as well or not.
        # p = p from PAdam
        # IA = bool to decide if Iterate Averaging is ever going to be used
        # IA_cycle = Iterate Averaging Cycle (Recommended to initialize with no. of iterations in Epoch) (doesn't matter if you are not using IA)
        # epochs = No. of epochs you plan to use (Only relevant if using DEMON)
        # step_per_epoch = No. of iterations in an epoch (only relevant if using DEMON)
        # weight decay = decorrelated weight decay value
        # use_gc = bool to determine whether to use gradient centralization or not.
        # use_grad_noise = bool to determine whether to use gradient noise or not.

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= hypergrad_lr:
            raise ValueError("Invalid hypergradient learning rate: {}".format(hypergrad_lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        if not 0.0 <= betas[2] < 1.0:
            raise ValueError("Invalid beta parameter at index 2: {}".format(betas[2]))
        if not 0.0 <= nus[0] <= 1.0:
            raise ValueError("Invalid nu parameter at index 0: {}".format(nus[0]))
        if not 0.0 <= nus[1] <= 1.0:
            raise ValueError("Invalid nu parameter at index 1: {}".format(nus[1]))
        if not 0.0 <= p <= 0.5:
            raise ValueError("Invalid p parameter: {}".format(p))
        if not 0.0 <= alpha < 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.AdaMod_bias_correct = AdaMod_bias_correct
        self.nostalgia = nostalgia
        self.use_demon = use_demon
        self.k = k
        self.IA = IA
        self.IA_cycle = IA_cycle
        self.epochs = epochs
        if step_per_epoch is None:
            self.step_per_epoch = IA_cycle
        else:
            self.step_per_epoch = step_per_epoch
        self.use_gc = use_gc
        self.T = self.epochs*self.step_per_epoch

        defaults = dict(lr=lr,
                        betas=betas,
                        nus=nus,
                        eps=eps,
                        alpha=alpha,
                        gamma=gamma,
                        p=p,
                        hypergrad_lr=hypergrad_lr,
                        weight_decay=weight_decay)
        super(HyperRangerMod, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(HyperRangerMod, self).__setstate__(state)

    def step(self, display=False, activate_IA=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('HyperRangerMod does not support sparse gradients')

                state = self.state[p]

                hypergrad_lr = group['hypergrad_lr']
                beta1_init, beta2, beta3 = group['betas']
                nu1, nu2 = group['nus']
                wd = group['weight_decay']
                alpha = group['alpha']
                gamma = group['gamma']

                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)

                    if self.IA:
                        state['num_models'] = 0
                    if self.IA or self.k > 0:
                        state['cached_params'] = p.data.clone()
                    if beta3 > 0.0:
                        state['n_avg'] = torch.zeros_like(p.data)
                    if self.nostalgia:
                        state['B_old'] = 0
                        state['B_new'] = 1
                    if hypergrad_lr > 0.0:
                        state['cached_hypergrad_comp'] = torch.zeros_like(grad.view(-1))

                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']

                if self.use_demon:
                    temp = 1-(state['step']/self.T)
                    beta1 = beta1_init * temp / ((1-beta1_init)+beta1_init*temp)
                else:
                    beta1 = beta1_init

                if self.nostalgia:
                    beta2 = state['B_old']/state['B_new']
                    state['B_old'] += math.pow(state['step'], -gamma)
                    state['B_new'] += math.pow(state['step']+1, -gamma)

                do_IA = False
                lookahead_step = False

                if self.IA and activate_IA:
                    lookahead_step = False
                    if state['step'] % self.IA_cycle == 0:
                        do_IA = True
                elif self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if state['step'] > 1 and hypergrad_lr > 0.0:
                    du = state['cached_hypergrad_comp']
                    h = torch.dot(grad.view(-1), du)
                    state['lr'] -= hypergrad_lr * h
                    torch.max(state['lr'], torch.zeros_like(state['lr']), out=state['lr'])
                    if display:
                        print(state['lr'])

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                momentum = exp_avg.clone()
                momentum.div_(1 - (beta1 ** state['step'])).mul_(nu1).add_(1-nu1, grad)
                vt = exp_avg_sq.clone()

                if not self.nostalgia:
                    vt.div_(1 - (beta2 ** state['step']))
                if nu2 != 1.0:
                    vt.mul_(nu2).addcmul_(1-nu2, grad, grad)

                denom = vt.pow_(group['p']).add_(group['eps'])

                n = state['lr']/denom

                if beta3 > 0.0:  # apply AdaMod
                    n_avg = state['n_avg']
                    n_avg.mul_(beta3).add_(1 - beta3, n)
                    if self.AdaMod_bias_correct:
                        n_avg_ = n_avg.clone()
                        bias_correction3 = 1 - (beta3 ** state['step'])
                        n_avg_.div_(bias_correction3)
                        torch.min(n, n_avg_, out=n)
                    else:
                        torch.min(n, n_avg, out=n)

                p.data.add_(-n*momentum)

                if lookahead_step:
                    dalpha = alpha
                elif do_IA:
                    dalpha = (1/(state["num_models"]+1.0))
                else:
                    dalpha = 1.0

                if hypergrad_lr > 0.0:

                    if beta3 > 0.0:
                        grad_from_n = dalpha*(-(momentum/denom) - wd*p.data)

                        if self.AdaMod_bias_correct:
                            grad_from_n_avg_ = dalpha * \
                                (-((1-beta3)/bias_correction3)*(momentum/denom) - wd*p.data)
                            du = torch.where(n_avg_ < n,
                                             grad_from_n_avg_,
                                             grad_from_n)
                        else:
                            grad_from_n_avg = dalpha*(-(1-beta3)*(momentum/denom) - wd*p.data)
                            du = torch.where(n_avg < n,
                                             grad_from_n_avg,
                                             grad_from_n)

                    else:
                        du = dalpha*(-(momentum/denom) - wd*p.data)

                    state['cached_hypergrad_comp'] = du.view(-1)

                if wd != 0:
                    p.data.add_(-wd*state['lr'], p.data)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

                if do_IA:
                    p.data.add_(state["num_models"], state['cached_params']
                                ).div_(state["num_models"]+1.0)
                    state['cached_params'].copy_(p.data)
                    state["num_models"] += 1

        return loss

class HDQHSGDW(Optimizer):
    def __init__(self, params, lr=1e-3,
                 beta=0.999,
                 nu=0.7,
                 hypergrad_lr=1e-3,
                 HDM=False,
                 k=5,
                 alpha=0.5,
                 eps=1e-8,
                 weight_decay=0,
                 use_gc=True):

        # BASIC SGD + Momentum but with QHMomentum and Hypergradient descent over all beta, lr, and nu + Lookahead and decorrelated weight decay
        # they say the best of them all is still SGD + Momentum?

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= hypergrad_lr:
            raise ValueError("Invalid hypergradient learning rate: {}".format(hypergrad_lr))
        if not 0.0 <= beta < 1.0:
            raise ValueError("Invalid beta parameter: {}".format(beta))
        if not 0.0 <= nu <= 1.0:
            raise ValueError("Invalid nu parameter: {}".format(nu))
        if not 0.0 <= alpha < 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.k = k
        self.use_gc = use_gc
        self.HDM = HDM

        defaults = dict(lr=lr,
                        beta=beta,
                        nu=nu,
                        alpha=alpha,
                        hypergrad_lr=hypergrad_lr,
                        eps=eps,
                        weight_decay=weight_decay)
        super(HDQHSGDW, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(HDQHSGDW, self).__setstate__(state)

    def hyperupdate(self, update, grad, grad_comp, hypergrad_lr, eps):
        h = torch.dot(grad.view(-1), grad_comp)

        if self.HDM:
            grad_norm = grad.view(-1).norm()
            norm_denom = grad_norm*(grad_comp.norm())
            norm_denom.add_(eps)
            update = update*(1-hypergrad_lr*(h/norm_denom))
        else:
            update -= hypergrad_lr * h

        return update

    def step(self, display=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('HDQHSGDW does not support sparse gradients')

                state = self.state[p]

                hypergrad_lr = group['hypergrad_lr']
                wd = group['weight_decay']
                alpha = group['alpha']

                if len(state) == 0:
                    state['step'] = 0
                    state['lr'] = group['lr']
                    state['nu'] = group['nu']
                    state['beta'] = group['beta']
                    state['exp_avg'] = torch.zeros_like(p.data)
                    if self.k > 0:
                        state['cached_params'] = p.data.clone()
                    if hypergrad_lr > 0.0:
                        state['prev_lr_grad'] = torch.zeros_like(grad.view(-1))
                        state['prev_nu_grad'] = torch.zeros_like(grad.view(-1))
                        state['prev_beta_grad'] = torch.zeros_like(grad.view(-1))

                state['step'] += 1
                exp_avg = state['exp_avg']

                lookahead_step = False

                if self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if state['step'] > 1 and hypergrad_lr > 0.0:

                    prev_lr_grad = state['prev_lr_grad']
                    prev_beta_grad = state['prev_beta_grad']
                    prev_nu_grad = state['prev_nu_grad']

                    state['lr'] = self.hyperupdate(update=state['lr'],
                                                   grad=grad,
                                                   grad_comp=prev_lr_grad,
                                                   hypergrad_lr=hypergrad_lr,
                                                   eps=group['eps'])

                    torch.max(state['lr'], torch.zeros_like(state['lr']), out=state['lr'])

                    if display:
                        print("lr", state['lr'])

                    state['beta'] = self.hyperupdate(update=state['beta'],
                                                     grad=grad,
                                                     grad_comp=prev_beta_grad,
                                                     hypergrad_lr=hypergrad_lr,
                                                     eps=group['eps'])

                    torch.max(state['beta'], torch.zeros_like(state['beta']), out=state['beta'])
                    torch.min(state['beta'], torch.ones_like(state['beta']), out=state['beta'])

                    if display:
                        print("beta", group['beta'])

                    state['nu'] = self.hyperupdate(update=state['nu'],
                                                   grad=grad,
                                                   grad_comp=prev_beta_grad,
                                                   hypergrad_lr=hypergrad_lr,
                                                   eps=group['eps'])

                    torch.max(state['nu'], torch.zeros_like(state['nu']), out=state['nu'])
                    torch.min(state['nu'], torch.ones_like(state['nu']), out=state['nu'])

                    if display:
                        print("nu", state['nu'])

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                nu = state['nu']
                beta = state['beta']
                lr = state['lr']

                if lookahead_step:
                    dalpha = alpha
                else:
                    dalpha = 1.0

                gx = 1 - (beta ** state['step'])
                fx = beta*exp_avg + (1-beta)*grad

                if hypergrad_lr > 0.0:
                    dfx = exp_avg - grad
                    dgx = - state['step'] * beta**(state['step']-1)
                    dbeta = (gx*dfx + fx*dgx)/(math.pow(gx, 2)+group['eps'])
                    dbeta = - dalpha*lr*nu*dbeta
                    state['prev_beta_grad'] = dbeta.view(-1)

                momentum = fx/gx
                group['exp_avg'] = fx

                if hypergrad_lr > 0.0:
                    state['prev_nu_grad'] = (-dalpha*lr*(momentum - grad)).view(-1)

                momentum.mul_(nu).add_(1-nu, grad)  # quasi hyperbolic momentum

                if hypergrad_lr > 0.0:
                    temp = dalpha*(-momentum - wd*p.data)
                    state['prev_lr_grad'] = temp.view(-1)

                p.data.add_(-group['lr']*momentum)

                if wd != 0:
                    p.data.add_(-wd*group['lr'], p.data)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

        return loss

class HyperProp(Optimizer):

    # LaProp + hypergradient descent on lr and nu (for QH momentum) + QH Momentum + Decaying Momentum (DEMON) + Lookahead + Iterate Averaging + Nostalgia (from NosAdam) + P from PAdam
    # + gradient centralization + weight decay
    def __init__(self, params, lr=1e-3,
                 betas=(0.999, 0.999),
                 nu=0.7,
                 eps=1e-8,
                 gamma=0.0001,
                 nostalgia=True,
                 use_demon=True,
                 hypergrad_lr=0.02,
                 HDM=True,
                 hypertune_nu=True,
                 p=0.25,
                 k=5,
                 alpha=0.8,
                 IA=True,
                 IA_cycle=1000,
                 epochs=100,
                 step_per_epoch=None,
                 weight_decay=0,
                 use_gc=True):

        # betas = (beta1 for first order moments, beta2 for second order moments)
        # nu = for quasi hyperbolic momentum
        # eps = small value for numerical stability (avoid divide by zero)
        # k = lookahead cycle
        # alpha = outer learning rate (lookahead)
        # gamma = used for nostalgia
        # nostalgia = bool to decide whether to use nostalgia (from Nostalgic Adam or NosAdam)
        # use_demon = bool to decide whether to use DEMON (Decaying Momentum) or not
        # hypergrad_lr = learning rate for updating hyperparameters (like lr) through hypergradient descent (probably need to increase around 0.02 if HDM is True). Set to 0.0 to disable hypergradient descent.
        # HDM = bool to decide whether to use Multiplicative rule for updating hyperparameters or not
        # hypertune_nu1 = bool to decide whether apply hypergradient descent on nu1 as well or not.
        # p = p from PAdam
        # IA = bool to decide if Iterate Averaging is ever going to be used
        # IA_cycle = Iterate Averaging Cycle (Recommended to initialize with no. of iterations in Epoch) (doesn't matter if you are not using IA)
        # epochs = No. of epochs you plan to use (Only relevant if using DEMON)
        # step_per_epoch = No. of iterations in an epoch (only relevant if using DEMON)
        # weight decay = decorrelated weight decay value
        # use_gc = bool to determine whether to use gradient centralization or not.
        # use_grad_noise = bool to determine whether to use gradient noise or not.

        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= hypergrad_lr:
            raise ValueError("Invalid hypergradient learning rate: {}".format(hypergrad_lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        if not 0.0 <= nu <= 1.0:
            raise ValueError("Invalid nu parameter: {}".format(nu))
        if not 0.0 <= p <= 0.5:
            raise ValueError("Invalid p parameter: {}".format(p))
        if not 0.0 <= alpha < 1.0:
            raise ValueError("Invalid alpha parameter: {}".format(alpha))

        self.nostalgia = nostalgia
        self.use_demon = use_demon
        self.k = k
        self.IA = IA
        self.IA_cycle = IA_cycle
        self.epochs = epochs
        if step_per_epoch is None:
            self.step_per_epoch = IA_cycle
        else:
            self.step_per_epoch = step_per_epoch
        self.use_gc = use_gc
        self.T = self.epochs*self.step_per_epoch
        self.hypertune_nu = hypertune_nu
        self.HDM = HDM

        defaults = dict(lr=lr,
                        betas=betas,
                        nu=nu,
                        eps=eps,
                        alpha=alpha,
                        gamma=gamma,
                        p=p,
                        hypergrad_lr=hypergrad_lr,
                        weight_decay=weight_decay)
        super(HyperProp, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(HyperProp, self).__setstate__(state)

    def step(self, activate_IA=False, display=False, closure=None):

        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data.float()
                if grad.is_sparse:
                    raise RuntimeError('HyperProp does not support sparse gradients')

                state = self.state[p]

                hypergrad_lr = group['hypergrad_lr']
                beta1_init, beta2 = group['betas']
                wd = group['weight_decay']
                alpha = group['alpha']
                gamma = group['gamma']

                if len(state) == 0:
                    state['lr'] = group['lr']
                    state['nu'] = group['nu']
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p.data)
                    state['exp_avg_sq'] = torch.zeros_like(p.data)

                    if self.IA:
                        state['num_models'] = 0
                    if self.IA or (self.k > 0):
                        state['cached_params'] = p.data.clone()
                    if self.nostalgia:
                        state['B_old'] = 0
                        state['B_new'] = 1
                    if hypergrad_lr > 0.0:
                        state['prev_lr_grad'] = torch.zeros_like(grad.view(-1))
                        if self.hypertune_nu:
                            state['prev_nu_grad'] = torch.zeros_like(grad.view(-1))

                state['step'] += 1
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']

                if self.use_demon:
                    temp = 1-(state['step']/self.T)
                    beta1 = beta1_init * temp / ((1-beta1_init)+beta1_init*temp)
                else:
                    beta1 = beta1_init

                if self.nostalgia:
                    beta2 = state['B_old']/state['B_new']
                    state['B_old'] += math.pow(state['step'], -gamma)
                    state['B_new'] += math.pow(state['step']+1, -gamma)

                do_IA = False
                lookahead_step = False

                if self.IA and activate_IA:
                    lookahead_step = False
                    if state['step'] % self.IA_cycle == 0:
                        do_IA = True
                elif self.k == 0:
                    lookahead_step = False
                else:
                    if state['step'] % self.k == 0:
                        lookahead_step = True
                    else:
                        lookahead_step = False

                if state['step'] > 1 and hypergrad_lr > 0.0:
                    prev_lr_grad = state['prev_lr_grad']
                    h = torch.dot(grad.view(-1), prev_lr_grad)

                    if self.HDM:
                        grad_norm = grad.view(-1).norm()
                        norm_denom = grad_norm*(prev_lr_grad.norm())
                        norm_denom.add_(group['eps'])
                        state['lr'] = state['lr']*(1-hypergrad_lr*(h/norm_denom))
                    else:
                        state['lr'] -= hypergrad_lr * h

                    torch.max(state['lr'], torch.zeros_like(state['lr']), out=state['lr'])

                    if display:
                        print("lr", state['lr'])

                    if self.hypertune_nu:
                        prev_nu_grad = state['prev_nu_grad']
                        h = torch.dot(grad.view(-1), prev_nu_grad)
                        if self.HDM:
                            norm_denom = grad_norm*(prev_nu_grad.norm())
                            norm_denom.add_(group['eps'])
                            state['nu'] = state['nu']*(1-hypergrad_lr*(h/norm_denom))
                        else:
                            state['nu'] -= hypergrad_lr * h

                        torch.max(state['nu'], torch.zeros_like(state['nu']), out=state['nu'])
                        torch.min(state['nu'], torch.ones_like(state['nu']), out=state['nu'])

                    if display:
                        print("nu", state['nu'])

                if self.use_gc:
                    grad.add_(-grad.mean(dim=tuple(range(1, len(list(grad.size())))), keepdim=True))

                nu = state['nu']
                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                vt = exp_avg_sq.clone()

                if not self.nostalgia:
                    vt.div_(1 - (beta2 ** state['step']))

                denom = vt.pow_(group['p']).add_(group['eps'])

                exp_avg.mul_(beta1).addcdiv_(1 - beta1, grad, denom)

                momentum = exp_avg.clone()
                bias_correction1 = 1 - (beta1 ** state['step'])
                momentum.div_(bias_correction1)

                if lookahead_step:
                    dalpha = alpha
                elif do_IA:
                    dalpha = (1/(state["num_models"]+1.0))
                else:
                    dalpha = 1.0

                if hypergrad_lr > 0.0 and self.hypertune_nu:
                    state['prev_nu_grad'] = (-dalpha*state['lr']*(momentum - grad)).view(-1)

                momentum.mul_(nu).add_(1-nu, grad)  # quasi hyperbolic momentum

                if hypergrad_lr > 0.0:
                    temp = dalpha*(-momentum - wd*p.data)
                    state['prev_lr_grad'] = temp.view(-1)

                p.data.add_(-state['lr'] * momentum)

                if wd != 0:
                    p.data.add_(-wd*state['lr'], p.data)

                if lookahead_step:
                    p.data.mul_(alpha).add_(1.0 - alpha, state['cached_params'])
                    state['cached_params'].copy_(p.data)

                if do_IA:
                    p.data.add_(state["num_models"], state['cached_params']
                                ).div_(state["num_models"]+1.0)
                    state['cached_params'].copy_(p.data)
                    state["num_models"] += 1

        return loss

# new stuffs to try: https://arxiv.org/pdf/1607.04381.pdf


import numpy as np
import matplotlib.pyplot as plt
from astropy import units as u
import scipy.integrate
from sys import float_info
import warnings

class Baseliner:
  """
  A class for interactive baseliner of spectroscopic data.

  The class works by being fed a spectrum and a matplotlib axis on which
  it should be plotted. The spectrum is then plotted to the given axis,
  and a number of interactive options are made available to the user.

  Left-clicking with the mouse for the first time starts defining a window
  from the x-axis location of the click. A second click finishes the
  window between the locations of the first and second click.
  A third click will finish selecting windows, and perform the baselining.
  Alternately, right-clicking will cancel the last left-clicking action,
  allowing misplaced windows to be adjusted.

  Two keys are also accepted:
  Pressing "q" will cause the baselining process to be canceled,
  effectively skipping the baselining of this spectrum.
  Pressing "a" will allow an additional window to be defined, assuming
  one has been defined so far (by left-clicking twice to define its
  boundaries).

  Attributes
  ----------
  windows : `list`
    A list of all the set windows.
  """
  def __init__(self,ax,spec):
    """
    Baseliner(ax,spec,order=1)

    Initialise the `Baseliner` class by giving it the target axis and
    spectrum.

    Parameters
    ----------
    ax : `matplotlib.axis`
      The matplotlib axis on which the interation will happen.
    spec : `omnifit.spectrum.BaseSpectrum`
      The spectrum which will be plotted as the visual reference on
      the given axis.
    """
    self.__ax = ax
    self.__spec = spec
    self.__x = spec.x.value
    self.__y = spec.y.value
    self.__limlo=None
    self.__limhi=None
    self.__minx=np.min(self.__x)
    self.__maxx=np.max(self.__x)
    self.__miny=np.min(self.__y)
    self.__maxy=np.max(self.__y)
    self.__ax.set_xlim(self.__minx,self.__maxx)
    self.__ax.set_ylim(self.__miny,self.__maxy)
    self.__specplot,=self.__ax.plot(self.__x,self.__y,'k-',drawstyle='steps-mid')
    self.__buttonListener = self.__ax.figure.canvas.mpl_connect('button_press_event', self.__mouse_press)
    self.__keyListener = self.__ax.figure.canvas.mpl_connect('key_press_event', self.__key_press)
    self.windows=[]
  def __key_press(self, event):
    if event.key=='q':
      self.__skip()
    if event.key=='a' and self.__limlo != None and self.__limhi != None:
      self.__addwindow(self.__limlo,self.__limhi)
      self.__ax.plot([self.__limlo,self.__limlo],[self.__miny,self.__maxy],'g-')
      self.__ax.plot([self.__limhi,self.__limhi],[self.__miny,self.__maxy],'g-')
      self.__remlim()
      self.__remlim()
      print 'Window added. Ready to receive another one.'
    else:
      return
  def __mouse_press(self, event):
    if event.button==1:
      self.__setlim(event.xdata)
    elif event.button==2:
      return
    elif event.button==3:
      self.__remlim()
  def __skip(self):
    plt.close()
  def __setlim(self,i_x):
    if self.__limlo==None:
      self.__limlo=i_x
      self.__limloplot,=self.__ax.plot([i_x,i_x],[self.__miny,self.__maxy],'b-')
      self.__ax.figure.canvas.draw()
    elif self.__limhi==None:
      self.__limhi=i_x
      self.__limhiplot,=self.__ax.plot([i_x,i_x],[self.__miny,self.__maxy],'b-')
      self.__ax.figure.canvas.draw()
      print 'Ready for finalising. Press once more to do so, or press a to add another window.'
    else:
      self.__finalise()
  def __remlim(self):
    if self.__limhi!=None:
      self.__limhi=None
      self.__limhiplot.set_ydata([self.__miny,self.__miny])
      self.__ax.figure.canvas.draw()
    elif self.__limlo!=None:
      self.__limlo=None
      self.__limloplot.set_ydata([self.__miny,self.__miny])
      self.__ax.figure.canvas.draw()
    else:
      print 'No limits to cancel.'
  def __addwindow(self,limlo,limhi):
    if limhi < limlo:
      limlo,limhi = limhi,limlo
    self.windows.append([limlo,limhi])
  def __finalise(self):
    self.__addwindow(self.__limlo,self.__limhi)
    self.__ax.figure.canvas.mpl_disconnect(self.__buttonListener)
    self.__ax.figure.canvas.mpl_disconnect(self.__keyListener)
    plt.close(self.__ax.figure)

#---------------------
#New units definitions
#---------------------
#the units themselves
unit_t = u.def_unit('transmittance units',doc='Transmittance of radiation')
unit_transmittance = unit_t
unit_abs = u.def_unit('absorbance units',doc='Absorbance of radiation')
unit_absorbance = unit_abs
unit_od = u.def_unit('optical depth units',doc='Optical depth of radiation')
unit_opticaldepth = unit_od

#the equivalencies between the units
equivalencies_absorption = [
    (unit_t,unit_abs,lambda x:-np.log10(x),lambda x:10**-x),
    (unit_od,unit_abs,lambda x:x/np.log(10),lambda x:x*np.log(10)),
    (unit_od,unit_t,lambda x:10**(-x/np.log(10)),lambda x:-np.log10(x)*np.log(10))
    ]

#------------------------------------------------------
#Functions related to light scattering and transmission
#------------------------------------------------------
def cde_correct(freq,m):
  """
  cde_correct(freq,m)

  Generate a CDE-corrected spectrum from a complex refractive index
  spectrum.

  Parameters
  ----------
  freq : `numpy.ndarray`
    The frequency data of the input spectrum, in reciprocal
    wavenumbers (cm^-1).
  m : `numpy.ndarray`
    The complex refractive index spectrum.

  Returns
  -------
  A list containing the following numpy arrays, in given order:
    * The spectrum of the absorption cross section of the simulated grain.
    * The spectrum of the absorption cross section of the simulated grain,
      normalized by the volume distribution of the grain. This parameter
      is the equivalent of optical depth in most cases.
    * The spectrum of the scattering cross section of the simulated grain,
      normalized by the volume distribution of the grain.
    * The spectrum of the total cross section of the simulated grain.
  """
  wl=1.e4/freq
  m2=m**2.0
  im_part=((m2/(m2-1.0))*np.log(m2)).imag
  cabs_vol=(4.0*np.pi/wl)*im_part
  cabs=freq*(2.0*m.imag/(m.imag-1))*np.log10(m.imag)
  cscat_vol=(freq**3.0/(6.0*np.pi))*cabs
  ctot=cabs+cscat_vol
  return cabs,cabs_vol,cscat_vol,ctot

def complex_transmission_reflection(in_m0,in_m1,in_m2):
  """
  complex_transmission_reflection(in_m0,in_m1,in_m2)

  Calculate the complex transmission and reflection coefficients between
  media 0, 1, and 2 given their complex refractive indices.
  In the Kramers-Kronig implementation (in which this is most likely used
  in the context of Omnifit) media 0, 1, and 2 correspond
  respectively to the vacuum, ice, and substrate.

  Parameters
  ----------
  in_m0 : `complex` or `numpy.ndarray`
    The complex refractive index of medium 0.
  in_m1 : `complex` or `numpy.ndarray`
    The complex refractive index of medium 1.
  in_m2 : `complex` or `numpy.ndarray`
    The complex refractive index of medium 2.

  Returns
  -------
  A tuple containing the following elements:
    * The complex transmission coefficient between media 0 and 1
    * The complex transmission coefficient between media 0 and 2
    * The complex transmission coefficient between media 1 and 2
    * The complex reflection coefficient between media 0 and 1
    * The complex reflection coefficient between media 0 and 2
    * The complex reflection coefficient between media 1 and 2
  """
  complex_transmission = lambda m1,m2: (2.*m1.real)/(m1+m2)
  complex_reflection = lambda m1,m2: (m1-m2)/(m1+m2)
  return (
          complex_transmission(in_m0,in_m1),
          complex_transmission(in_m0,in_m2),
          complex_transmission(in_m1,in_m2),
          complex_reflection(in_m0,in_m1),
          complex_reflection(in_m0,in_m2),
          complex_reflection(in_m1,in_m2)
        )

def kramers_kronig(freq,transmittance,m_substrate,d_ice,m0,freq_m0,m_guess=1.0+0.0j,tol=0.001,maxiter=100,ignore_fraction=0.1,force_kkint_unity=False,precalc=False):
  """
  kramers_kronig(freq,transmittance,m_substrate,d_ice,m0,freq_m0,
                 m_guess=1.0+0.0j,tol=0.001,maxiter=100,ignore_fraction=0.1,
                 force_kkint_unity=False,precalc=False)

  Kramers-Kronig relation.
  This is an implementation of the Kramers-Kronig relation calculation
  presented in Hudgins et al 1993 (1993ApJS...86..713H), with an improved
  integration method adapted from Trotta et al 1996
  (The Cosmic Dust Connection, 1996 169-184)

  Parameters
  ----------
  wn : `astropy.units.Quantity` or `numpy.ndarray`
    The frequency data of the input spectrum. If no units are given, this
    is assumed to be in reciprocal wavenumbers (cm^-1).
  transmittance : `astropy.units.Quantity` or `numpy.ndarray`
    The transmittance data of the input spectrum. This can be given in
    units other than transmittance, as long as they can be converted to
    transmittance by making use of the `utils.equivalencies_absorption`
    equivalency information. If no units are given, transmittance is
    assumed.
  m_substrate : `complex`
    The complex refractive index of the substrate on which the ice being
    studied was grown.
  d_ice : `astropy.units.Quantity` or `float`
    The thickness of the ice which is being studied. If no units are given,
    centimeters are assumed.
  m0 : `complex`
    The complex refractive index of the ice at the reference frequency
    defined by `freq_m0` (see below).
  freq_m0 : `astropy.units.Quantity` or `float`
    The frequency at which the reference complex refractive index `m0`
    (see above) is defined. Best results are usually achieved if this
    frequency is high compared to the frequency range being probed by
    the spectrum.
    If this is not defined as `astropy.units.Quantity` in spectroscopic
    units, it is assumed to be in reciprocal wavenumbers (cm^-1).
  m_guess : `complex` or `numpy.ndarray`
    The starting guess of the complex refractive index of the ice. This
    can either be a single number (in which case it is assumed to be this
    number throughout the entire spectrum) or an array
  tol : `float`
    The square-sum of the residual between the original transmittance and
    the transmittance modeled with the iterated complex refractive index
    of the ice must be below this value for the iteration to converge. In
    other words, the smaller this number is, the better the final result
    will be at the expense of extra iterations.
  maxiter : `int`
    The maximum number of iterations allowed. If this number is reached,
    the iteration is considered to not have converged, and an exception is
    raised.
  ignore_fraction : `float` between 0 and 0.5
    The edges of the spectrum are blanked out (and replaced with the
    non-blanked value closest to the edge) during iteration to avoid edge
    effects arising from the usage of a non-infinite integration range.
    This parameter controls how large of a fraction of the edges is blanked
    out.
  force_kkint_unity : `bool`
    The results of the Kramers-Kronig integration are responsible for
    determining the real part of the complex refractive index i.e. the
    one which represents refraction. Normally this number should not drop
    below unity, and unexpected behaviour can arise if it does.
    Usually this means that there is something wrong with the input
    parameters, but sometimes forcing the result to always be greater or
    equal to unity can help. It should be noted, however, that the
    accuracy of the results of an integration forced in this way are
    suspect at best.
  precalc : `bool`
    The Kramers-Kronig iteration can be a very computationally intensive
    operation. In some situations it may result in a faster iteration to
    pre-calculate the large denominator which is part of the
    Kramers-Kronig integration instead of computing new values of it in a
    for loop. This denominator can be, however, a very
    large variable as it contains a number of elements equal to the size
    of the spectrum squared. Pre-calculating this can outright fail on
    lower-end computers as Python runs out of available memory.
    High-end systems may benefit from such pre-calculation, though.

  Returns
  -------
  A `numpy.ndarray` which contains the complex refractive index of the
  ice, in order of increasing frequency.
  """
  #set up constants
  m_vacuum = 1.0+0.0j
  #make sure the input array units are correct; convert if necessary
  if type(freq) != u.quantity.Quantity:
    warnings.warn('No units detected in input freq. Assuming kayser.',RuntimeWarning)
    freq *= u.kayser
  else:
    with u.set_enabled_equivalencies(u.equivalencies.spectral()):
      freq=freq.to(u.kayser)
  if type(transmittance) != u.quantity.Quantity:
    warnings.warn('No units detected in input transmittance. Assuming transmittance units.',RuntimeWarning)
    transmittance *= unit_t
  else:
    with u.set_enabled_equivalencies(equivalencies_absorption):
      transmittance = transmittance.to(unit_t)
  if type(d_ice) != u.quantity.Quantity:
    warnings.warn('No units detected in input d_ice. Assuming centimeters.',RuntimeWarning)
    d_ice *= u.cm
  else:
    d_ice = d_ice.to(u.cm)
  #sort the arrays and get rid of units; won't need them after this
  initial_sorter = np.argsort(freq)
  freq = freq[initial_sorter].value
  transmittance = transmittance[initial_sorter].value
  d_ice = d_ice.value
  #initialise complex refractive index and alpha arrays
  m = np.full_like(freq,np.nan+np.nan*1j,dtype=complex)
  alpha = np.full_like(freq,np.nan+np.nan*1j,dtype=complex)
  #initial guess at m at first index
  if type(m_guess)==complex:
    m_ice = np.full_like(freq,m_guess,dtype=complex)
  else:
    m_ice = m_guess
  #find top and bottom fraction indices. These will be replaced with dummy values after each integration to get rid of edge effects
  if ignore_fraction > 0.5 or ignore_fraction < 0:
    raise RuntimeError('ignore_fraction must be between 0.0 and 0.5')
  bot_fraction = round(ignore_fraction*len(freq))
  top_fraction = len(freq)-bot_fraction
  #pre-calculate the large denominator component of the KK integration, if desired
  if precalc:
    try:
      sfreq=(freq).reshape(len(freq),1)
      kkint_deno1 = freq**2-sfreq**2
      kkint_deno1[kkint_deno1!=0] = 1./kkint_deno1[kkint_deno1!=0]
      precalc = True
    #or at least try to do so; if run out of memory, switch to the slower no-precalc mode
    except MemoryError:
      precalc = False
  #some other parts can always be precalced
  kkint_mul = 1./(2*np.pi*np.pi)
  kkint_deno2 = freq**2-freq_m0**2
  kkint_deno2[kkint_deno2!=0] = 1./kkint_deno2[kkint_deno2!=0]
  #calculate alpha at freq0
  alpha0 = m0.imag/(4*np.pi*freq)
  #iteration begin!
  niter = 0
  squaresum_diff = tol+1
  while squaresum_diff > tol and niter < maxiter:
    #calculate transmission and relfection coefficients
    #in these 0 means vacuum, 1 means ice, 2 means substrate
    t01,t02,t12,r01,r02,r12 = complex_transmission_reflection(m_vacuum,m_ice,m_substrate)
    #the reflection component
    # reflection_component = np.abs((t01*t12/t02)/(1.+r01*r12*np.exp(4.j*np.pi*d_ice*m_ice*freq)))**2.)
    #this is an evil equation. do NOT touch it
    #it calculates the lambert absorption coefficient using the current best guess at m_ice
    alpha = (1./d_ice)*(-np.log(transmittance)+np.log(np.abs((t01*t12/t02)/(1.+r01*r12*np.exp(4.j*np.pi*d_ice*m_ice*freq)))**2.))
    #using the new alpha, calculate a new n (and thus m) for the ice
    #this is done in a parallel for loop, to avoid killing the computer when dealing with large amounts of data
    kkint_nomi = alpha-alpha0
    kkint = np.full_like(alpha,m0.real)
    numcols = kkint_nomi.shape[0]
    for current_col in range(numcols):
      if precalc:
        kkint[current_col]+=kkint_mul*scipy.integrate.simps((alpha-alpha[current_col])*kkint_deno1[current_col,:]-kkint_nomi*kkint_deno2)
      else:
        kkint_deno1 = freq[current_col]**2-freq**2
        kkint_deno1[kkint_deno1!=0] = 1./kkint_deno1[kkint_deno1!=0]
        kkint[current_col]+=kkint_mul*scipy.integrate.simps((alpha-alpha[current_col])*kkint_deno1-kkint_nomi/(freq**2-freq_m0**2))
    if np.any(kkint<1):
      if np.any(kkint<0):
        warnings.warn('KK integration is producing negative refractive indices! This will most likely produce nonsensical results.',RuntimeWarning)
      else:
        warnings.warn('KK integration is producing refractive indices below unity! This may result in unexpected behaviour.',RuntimeWarning)
      if force_kkint_unity:
        kkint[kkint<1]=1.
    m_ice = kkint+1j*alpha/(4*np.pi*freq)
    if np.any(np.isnan(m_ice.real)) or np.any(np.isnan(m_ice.imag)):
      raise RuntimeError('Produced complex refractive index contains NaNs. Check your input parameters.')
    #replace top and bottom fractions of m_ice with the value closest to that edge
    #this is done to combat edge effects arising from integrating over a non-infinite range
    m_ice[:bot_fraction] = m_ice[bot_fraction]
    m_ice[top_fraction:] = m_ice[top_fraction]
    #calculate transmission and relfection coefficients (again)
    #in these 0 means vacuum, 1 means ice, 2 means substrate
    t01,t02,t12,r01,r02,r12 = complex_transmission_reflection(m_vacuum,m_ice,m_substrate)
    #model a transmittance using given m_ice and alpha
    #yes, this is another evil equation
    transmittance_model = np.exp(-alpha*d_ice)*np.abs((t01*t12/t02)/(1.+r01*r12*np.exp(4.j*np.pi*d_ice*m_ice*freq)))**2.
    diff = transmittance - transmittance_model
    diff[:bot_fraction] = 0. #ignore top...
    diff[top_fraction:] = 0. #...and bottom fraction differences
    squaresum_diff = np.sum(diff**2) #square sum of difference
    niter += 1
  #at this point we are done
  if niter>=maxiter:
    raise RuntimeError('Maximum number of iterations reached before convergence criterion was met.')
  return m_ice