#!/usr/bin/env python
#
# pKaTool - analysis of systems of titratable groups
# Copyright (C) 2010 Jens Erik Nielsen
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact information:
# Email: Jens.Nielsen_at_gmail.com
# Normal mail:
# Jens Nielsen
# SBBS, Conway Institute
# University College Dublin
# Dublin 4, Ireland
#
from Tkinter import *
from numpy import *
from numpy.linalg import *
inverse=inv
Float=float

from titration_class import *

import math

class Optimisation_Analysis:

    def stab_and_CCPS_pop(self):
        """Driver for optimising stability and CCPS population"""
        #
        # Open the optimisation window
        #
        self.optwin=Toplevel()
        parent_geom=self.window.winfo_geometry()
        parent_geom=parent_geom.split('+')[1:]
        # Position optimisation window at same position as parent (titration curve window
        self.optwin.geometry('+%d+%d' %(int(parent_geom[0])+100,int(parent_geom[1])+200))
        self.optwin.title('System optimisation and analysis')
        #
        # Add scoring function specification interface
        #
        Label(self.optwin,text='Name').grid(row=1,column=0)
        Label(self.optwin,text='Optimise').grid(row=1,column=1)
        Label(self.optwin,text='Opt type').grid(row=1,column=2)
        Label(self.optwin,text='Cutoff value').grid(row=1,column=3)
        Label(self.optwin,text='Weight').grid(row=1,column=4)
        Label(self.optwin,text='pH range').grid(row=1,column=5)
        Label(self.optwin,text='Use profile').grid(row=1,column=6)
        Label(self.optwin,text='Optimise sum').grid(row=1,column=7)
        #
        self.properties={'CCPS population':{},
                         'CCPS2 population':{},
                         'CCPS maxrange':{},
                         'Stability':{},
                         'individual intpKa shift':{},
                         'individual intene':{}
                         }
        props=self.properties.keys()
        props.sort()
        row=1
        for prop in props:
            row=row+1
            Label(self.optwin,text=prop).grid(row=row,column=0)
            #
            self.properties[prop]['include']={'var':IntVar()}
            self.properties[prop]['include']['widget']=Checkbutton(self.optwin,variable=self.properties[prop]['include']['var'],onvalue=1,offvalue=0)
            self.properties[prop]['include']['widget'].grid(row=row,column=1)
            self.properties[prop]['include']['var'].set(1)
            #
            self.properties[prop]['opttype']={'var':StringVar()}
            self.properties[prop]['opttype']['widget']=Menubutton(self.optwin,textvariable=self.properties[prop]['opttype']['var'],relief='raised')
            self.properties[prop]['opttype']['menu']=Menu(self.properties[prop]['opttype']['widget'],tearoff=0)
            (self.properties[prop]['opttype']['widget'])['menu']=self.properties[prop]['opttype']['menu']
            #
            # Set defaults
            #
            if prop=='CCPS population':
                value=90.0
                defopttype='maximise'
                weight=1.0
            elif prop=='CCPS2 population':
                value=10.0
                defopttype='maximise'
                weight=1.0
            elif prop=='Stability':
                value=0.0
                defopttype='maximise'
                weight=1.0
            elif prop=='individual intpKa shift':
                value=2.0
                defopttype='keep below'
                weight=10000.0
            elif prop=='individual intene':
                value=4.6
                defopttype='keep below'
                weight=10000.0
            elif prop=='CCPS maxrange':
                value=0.20
                defopttype='maximise'
                weight=1.0
            else:
                raise Exception
            #
            # Add all opttypes
            #
            for opttype in ['maximise','minimise','keep at','keep above','keep below']:
                self.properties[prop]['opttype']['menu'].add_radiobutton(label=opttype,
                                                                         variable=self.properties[prop]['opttype']['var'],
                                                                         value=opttype,
                                                                         indicatoron=1)
            self.properties[prop]['opttype']['widget'].grid(row=row,column=2)
            self.properties[prop]['opttype']['var'].set(defopttype)
            #
            self.properties[prop]['cutoff']={'var':DoubleVar()}

            self.properties[prop]['cutoff']['var'].set(value)
            self.properties[prop]['cutoff']['widget']=Entry(self.optwin,textvariable=self.properties[prop]['cutoff']['var'],width=10)
            self.properties[prop]['cutoff']['widget'].grid(row=row,column=3)
            #
            self.properties[prop]['weight']={'var':DoubleVar()}
            self.properties[prop]['weight']['var'].set(weight)
            self.properties[prop]['weight']['widget']=Entry(self.optwin,textvariable=self.properties[prop]['weight']['var'],width=10)
            self.properties[prop]['weight']['widget'].grid(row=row,column=4)
            #
            self.properties[prop]['pH']={'var':StringVar()}
            self.properties[prop]['pH']['var'].set('6.5-7.5')
            self.properties[prop]['pH']['widget']=Entry(self.optwin,textvariable=self.properties[prop]['pH']['var'],width=10)
            self.properties[prop]['pH']['widget'].grid(row=row,column=5)
            #
            # Add a profile selector
            #
            if prop in ['CCPS maxrange','CCPS population','CCPS2 population','Stability']:
                self.properties[prop]['profile']={'var':IntVar()}
                self.properties[prop]['profile']['var'].set(0)
                self.properties[prop]['profile']['widget']=Checkbutton(self.optwin,
                    variable=self.properties[prop]['profile']['var'],onvalue=1,offvalue=0)
                self.properties[prop]['profile']['widget'].grid(row=row,column=6)
            #
            # Selector for individual opt of protonation state populations
            #
            if prop in ['CCPS maxrange','CCPS population','CCPS2 population']:
                self.properties[prop]['sumopt']={'var':IntVar()}
                self.properties[prop]['sumopt']['var'].set(1)
                self.properties[prop]['sumopt']['widget']=Checkbutton(self.optwin,
                                                                       variable=self.properties[prop]['sumopt']['var'],onvalue=1,offvalue=0)
                self.properties[prop]['sumopt']['widget'].grid(row=row,column=7)

        #
        # Open the microscopic population states
        #
        self.micro_var.set(1)
        self.update_pkasystem_curves(0)
        #
        # Open the stability window
        #
        self.stab_button.select()
        self.stability_on_off()
        #
        # Continue with the optimisation window
        #
        row=row+1
        Label(self.optwin,text='CCPS chosen').grid(row=row,column=0)
        self.chosen_CCPS=StringVar()
        Label(self.optwin,textvariable=self.chosen_CCPS).grid(row=row,column=1)
        self.find_chosen_CCPS()
        Button(self.optwin,text='update',command=self.find_chosen_CCPS).grid(row=row,column=2)
        #
        # Add action buttons
        #
        row=row+1
        Button(self.optwin,text='Load primary pH-activity profile',command=self.pre_load_pH_activity_profile).grid(row=row,column=0)
        Button(self.optwin,text='Load second pH-activity profile',command=self.pre_load_2nd_pH_activity_profile).grid(row=row,column=1)
        Button(self.optwin,text='Load pH-stability profile',state=DISABLED).grid(row=row,column=2)
        #
        row=row+1
        Button(self.optwin,text='Randomize system',fg='pink',command=self.randomizesystem).grid(row=row,column=0)
        Button(self.optwin,text='Start optimisation',fg='green',command=self.optimise_system).grid(row=row,column=1)
        Button(self.optwin,text='Stop optimisation',fg='red',command=self.stopopt).grid(row=row,column=2)

        #
        # Status bar/window
        #
        row=row+1
        self.status=StringVar()
        Label(self.optwin,textvariable=self.status,bg='white').grid(row=row,column=0,columnspan=4)
        #
        # Adjust the resolution of sliders
        #
        for group in self.groups:
            self.groups[group].intpka_scale.configure(resolution=0.01)
            for group2 in self.groups[group].intenes.keys():
                self.groups[group].intenes_scales[group2].configure(resolution=0.01)
        #
        # Set pH step
        #
        self.pHstep.set(0.1)
        return

    #
    # ----
    #

    def pre_load_pH_activity_profile(self):
        """Handle the pH_activity profile load"""
        self.load_pH_activity_profile(parent=self.optwin)
        self.properties['CCPS population']['profile']['var'].set(1)
        return

    def pre_load_2nd_pH_activity_profile(self):
        """Handle the pH_activity profile load"""
        self.load_pH_activity_profile(parent=self.optwin)
        self.properties['CCPS2 population']['profile']['var'].set(1)
        return

    #
    # ----
    #

    def stopopt(self):
        """Stop the current fitting operation"""
        self.keep_running_opt=0
        return

    #
    # -----
    #

    def find_chosen_CCPS(self):
        """Sets the variable that displays the protonation states that are chosen for the CCPS"""
        self.chosen_CCPS.set(self.get_CCPSs())
        return

    #
    # -----
    #

    def randomizesystem(self):
        return

    #
    # -----
    #

    def optimise_system(self):
        """Optimise the system according to the values given by the user"""
        #
        # Did the bozo choose the CCPS?
        #
        self.find_chosen_CCPS()
        if self.get_CCPSs()==[]:
            import tkMessageBox
            tkMessageBox.showwarning('No CCPS selected',
                                     'You have to chose the CCPS from the microscopic states',
                                     parent=self.optwin)
            return
        #
        # Find the properties we will use
        #
        self.active_properties=[]
        props=self.properties.keys()
        for prop in props:
            if self.properties[prop]['include']['var'].get()==1:
                self.active_properties.append(prop)
        #
        print 'I am optimising these properties:',self.active_properties

        #
        # Get the variables
        #
        self.intpkas=[]
        self.intenes=[]
        for group in self.groups:
            self.intpkas.append(self.groups[group].intpka)
            for group2 in self.groups[group].intenes.keys():
                self.intenes.append(self.groups[group].intenes[group2])
        #
        # Perform all variations within the limits
        #
        self.intpka_limit=self.properties['individual intpKa shift']['cutoff']['var'].get()
        self.intene_limit=self.properties['individual intene']['cutoff']['var'].get()
        print 'Intpka limit: %5.3f' %self.intpka_limit
        print 'intene limit: %5.3f' %self.intene_limit
        #
        # Make copies of all the variables
        #
        self.vars=[]
        for group in self.groups:
            self.vars.append([self.groups[group].intpka,'intpka',group])
            for group2 in self.groups[group].intenes.keys():
                self.vars.append([self.groups[group].intenes[group2],'intene',group])
        #
        # Set damper
        #
        self.LM_damper = 0.1
        #
        # Update the counter
        #
        self.count=0
        self.keep_running_opt=1
        #
        # Init vars for the scoring function
        #
        self.CCPS_pH_vals=None
        self.CCPSmax_pH_vals=None
        self.Stability_pH_vals=None
        #
        # Start off the iterations
        #
        scores=self.get_system_score().values()
        old_diff=10000.0
        now_diff=0.0
        for val in self.get_system_score().values():
            now_diff=now_diff+val
        min_steps=100
        for x in range(1,10000):
            self.fit_LM_optimise()
            #
            # Update the counter, the scales and the titration curves
            #
            self.status.set("Optimisation running. Step %5d" %self.count)
            self.count=self.count+1
            self.window.update()
            self.optwin.update()

            #self.optwin.update_idletasks()
            self.update_scales_from_fit()
            self.titwin.update() # This statement is essential for propagating changes in variable values
            #
            #
            # recalculate the differences
            #
            now_diff=0.0
            for val in self.get_system_score().values():
                now_diff=now_diff+val
            #
            # Check convergence
            #
            if abs(now_diff-old_diff)<0.000005 and x>min_steps:
                print 'Converged',now_diff
                break
            else:
                old_diff=now_diff
            #
            # Should we stop?
            #
            if self.keep_running_opt==0:
                print 'User abort'
                break
        self.update_pkasystem_curves()
        return

    #
    # ----
    #

    def get_jacobian_optimise(self):
        """Get the Jacobian matrix and errors of the data points"""
        #
        # We have one score for each property included
        #
        no_data_points=len(self.active_properties)
        if 'CCPS2 population' in self.active_properties:
            no_data_points=no_data_points-1
        errors = resize(array(0,float),[no_data_points])
        jacobian = resize(array(0,float),[no_data_points,len(self.vars)])
        #
        # Precalculate the variation of all parameters
        #

        variations=[]
        step = 1e-8
        for var in range(len(self.vars)):
            self.vars[var][0].set(self.vars[var][0].get()+step)
            variations.append(self.get_system_score())
            self.vars[var][0].set(self.vars[var][0].get()-step)
        #
        # construct jacobian
        #
        data_id=0
        #
        # Find the errors
        #
        now=self.get_system_score()
        for prop in self.active_properties:
            if prop=='CCPS2 population':
                continue
            y=0 # All functions always have minimum at zero
            errors[data_id] =0-now[prop]
            #
            # Find the derivatives
            #
            diff=resize(array(0,float),[len(self.vars)])
            count=0
            for variation in variations:
                diff[count]=(now[prop]-variation[prop])/step
                count=count+1
            jacobian[data_id]=diff
            data_id=data_id+1
        return jacobian,errors

    #
    # ----
    #

    def calc_sq_diff_optimise(self,charges):
        """Calculate the square difference"""
        diff=0.0
        import math
        for pH,crg in charges:
            y=crg
            fit_val=eval(self.funktion)
            diff=diff+math.pow(fit_val-y,2)
        return diff

    #
    # ------
    #

    def fit_LM_optimise(self):
        """Do Levenberg-Marquardt fitting"""
        J,E =self.get_jacobian_optimise()
        JT = transpose(J)
        JTE = dot(JT,E)
        JTJ = dot(JT,J)
        JTJd = JTJ + self.LM_damper*identity(shape(JTJ)[0])
        invJTJd = inverse(JTJd)
        q = -dot(JTE,invJTJd)
        #
        # Change the values
        #
        for var in range(len(self.vars)):
            #
            # Update the vars if they do not get outside the bounds
            #
            new_val=self.vars[var][0].get()+(q[var]/abs(q[var]))*min(abs(q[var]),0.2)
            group=self.vars[var][2]
            ok=None
            if self.vars[var][1]=='intpka':
                dintpka=abs(new_val-self.unfolded_groups[group].intpka.get())
                if dintpka<=self.intpka_limit and new_val>=0.0:
                    ok=1
                if self.properties['individual intpKa shift']['include']['var'].get()==0:
                    #
                    # If no restriction then the shift is always ok
                    #
                    ok=1
            elif self.vars[var][1]=='intene':
                if new_val<=self.intene_limit and new_val>=0.0:
                    ok=1
                if self.properties['individual intene']['include']['var'].get()==0:
                    #
                    # If no restriction then the shift is always ok
                    #
                    ok=1
            else:
                print 'Unknown variable'
                raise Exception
            #
            # If ok, then update vars
            #
            if ok:
                self.vars[var][0].set(new_val)
        return

    #
    # -----
    #

    def get_system_score(self):
        """Get value of the scoring function specified by the user"""

        self.X,pKa_values=self.calc_pKas(self.groups)
        curve=titration_curve(self.X.prot_states)
        self.stability=self.do_stab_curve(self.X)

        self.act_prof={}
        for pH in self.phvals:
            act=0.0
            for state in self.states:
                if self.act_state[state].get()==1:
                    act=act+self.X.all_states[pH][state]['pop']
            self.act_prof[pH]=act
        #
        # Secondary activity profile
        #
        self.act_prof2={}
        for pH in self.phvals:
            act=0.0
            for state in self.states:
                if self.act_state2[state].get()==1:
                    act=act+self.X.all_states[pH][state]['pop']
            self.act_prof2[pH]=act
        #
        # Score everything
        #
        scores={}
        CCPS2_score=None
        for prop in self.active_properties:
            thisprop=self.properties[prop]
            function=getattr(self,'_'+prop.replace(' ','_'))
            if prop=='CCPS2 population':
                CCPS2_score=function(thisprop)
            else:
                scores[prop]=function(thisprop)
        #
        # Score CCPSs under one
        #
        if CCPS2_score:
            scores['CCPS population']=max(CCPS2_score,scores['CCPS population'])
            print 'Doing max'
        #
        #
        props=scores.keys()
        props.sort()
        for prop in props:
            print '%10s score: %9.5f,' %(prop,scores[prop])
        print
        return scores

    #
    # -----
    #

    def _CCPS_population_base(self,variable,primary,spec):
        """Function that will give a score for how well one of the CCPS population meets the targets"""
        #
        # Loaded profile or entered values
        #
        if spec['profile']['var'].get()==0:
            #
            # Entered values
            #
            # Get the pH values where we want to impose the criterium
            #
            if not self.CCPS_pH_vals:
                self.CCPS_pH_vals=self.get_pH_range(spec['pH']['var'].get())
            print 'Monitoring activity at these pH values',self.CCPS_pH_vals
            #
            # pH values have already been identified
            #
            opttype=spec['opttype']['var'].get()
            cutoff=spec['cutoff']['var'].get()
            if opttype=='maximise':
                sum_act=0.0
                count=0
                for pH in self.CCPS_pH_vals:
                    if primary:
                        if self.act_prof[pH]<cutoff:
                            sum_act=sum_act+self.act_prof[pH]
                            count=count+1
                    else:
                        if self.act_prof2[pH]<cutoff:
                            sum_act=sum_act+self.act_prof2[pH]
                            count=count+1

                if count>1:
                    sum_act=sum_act/float(count)
                sum_act=max(cutoff/100.0-sum_act,0)
            else:
                raise Exception
        else:
            #
            # Loaded profile
            #
            sum_act=0.0
            #
            # Optimse profile, or each protonation state on its own
            #
            if spec['sumopt']['var'].get()==1:
                for pH in variable.keys():
                    if primary:
                        sum_act=sum_act+abs(variable[pH]-self.act_prof[pH])
                    else:
                        sum_act=sum_act+abs(variable[pH]-self.act_prof2[pH])
            else:
                #
                # Find the difference between each protonation state
                #
                diffs=[]
                for state in self.states:
                    thisdiff=0.0
                    if (primary and self.act_state[state].get()==1) or (not primary and self.act_state2[state].get()==1):
                        for pH in variable.keys():
                            thisdiff=thisdiff+abs(self.X.all_states[pH][state]['pop']-variable[pH])
                        diffs.append(thisdiff)
                sum_act=max(diffs)
        #
        # Scale the sum by the weight
        #
        sum_act=sum_act*spec['weight']['var'].get()
        return sum_act

    #
    # ----
    #

    def _CCPS_population(self,spec):
        """Get the score for the primary CCPS population"""
        return self._CCPS_population_base(self.activity_data,1,spec)

    def _CCPS2_population(self,spec):
        """Get the score for the seconary CCPS population"""
        return self._CCPS_population_base(self.secondary_activity_data,1,spec)

    #
    # ----
    #

    def _CCPS_maxrange(self,spec):
        """Score for having the maximum of the CCPS in a range"""
        #
        # Get the pH values where we want to impose the criterium
        #
        if not self.CCPSmax_pH_vals:
            self.CCPSmax_pH_vals=self.get_pH_range(spec['pH']['var'].get())
            print 'Monitoring activity at these pH values',self.CCPSmax_pH_vals
        #
        # pH values have already been identified
        #
        opttype=spec['opttype']['var'].get()
        cutoff=spec['cutoff']['var'].get()
        if opttype=='maximise':
            #
            # Find the maximum activity
            #
            maxact=0.0
            for pH in self.act_prof.keys():
                maxact=max(maxact,self.act_prof[pH])
            #
            # Report the fraction of points that are within 10%
            #
            count=0
            ok=0.0
            for pH in self.CCPSmax_pH_vals:
                if maxact>0.000001:
                    if abs((self.act_prof[pH]-maxact)/maxact)<=cutoff:
                        pass
                    else:
                        ok=ok+abs((self.act_prof[pH]-maxact)/maxact)
                        print pH,(self.act_prof[pH]-maxact)/maxact
                else:
                    ok=ok+10*(0.000001-maxact)
                count=count+1
            if count>0:
                ok=ok/float(count)
            #
        else:
            raise Exception
        ok=ok*spec['weight']['var'].get()
        return ok

    #
    # ----
    #

    def _Stability(self,spec):
        """Score the stability of the system"""
        if not self.Stability_pH_vals:
            self.Stability_pH_vals=self.get_pH_range(spec['pH']['var'].get())
        #
        # Score the stability
        #
        opttype=spec['opttype']['var'].get()
        cutoff=spec['cutoff']['var'].get()
        if opttype=='maximise':
            # Maximising stability means minimising the dG of folding
            sum_stab=0
            count=0
            for pH in self.Stability_pH_vals:
                if self.stability[pH]>cutoff:
                    sum_stab=sum_stab+self.stability[pH]-cutoff
                    count=count+1
            if count>1:
                sum_stab=sum_stab/float(count)
        else:
            raise Exeption
        import math
        sum_stab=sum_stab*spec['weight']['var'].get()
        #print 'Stability',math.exp(sum_stab)
        return sum_stab

    #
    # ----
    #

    def _individual_intpKa_shift(self,spec):
        """Score the total pKa shift"""
        return 0.0
        opttype=spec['opttype']['var'].get()
        cutoff=spec['cutoff']['var'].get()
        if opttype=='keep below':
            score=0.0
            for group in self.groups:
                dintpka=abs(self.groups[group].intpka.get()-self.unfolded_groups[group].intpka.get())
                if dintpka>cutoff:
                    score=score+math.pow(dintpka-cutoff,4)
        else:
            raise Exception
        return score*spec['weight']['var'].get()

    #
    # ----
    #

    def _individual_intene(self,spec):
        return 0.0
        """Keep all interaction energies below a certain kT limit"""
        import math
        opttype=spec['opttype']['var'].get()
        cutoff=spec['cutoff']['var'].get()
        if opttype=='keep below':
            score=0.0
            for group in self.groups:
                for num in self.groups[group].intenes:
                    intene=self.groups[group].intenes[num].get()
                    if intene>cutoff:
                        score=score+math.pow(intene-cutoff,4)
        else:
            raise Exception
        return score*spec['weight']['var'].get()
    #
    # ----
    #

    def get_pH_range(self,pH_vals):
        """Return a list of pH values that should be monitored in the individual scoring functions
        from the input given in the pH range field"""
        ranges=pH_vals.split(',')
        pH_vals=[]
        for sep in ranges:
            if sep.find('-')!=-1:
                start_end=sep.split('-')
                start=float(start_end[0])
                end=float(start_end[1])
                for pH in self.act_prof.keys():
                    if pH>=start and pH<=end:
                        pH_vals.append(pH)
            else:
                #
                # Simple number - just pick a value closer than 0.1
                #
                pH_set=float(sep)
                if self.act_prof.has_key(pH_set):
                    pH_vals.append(pH_set)
                else:
                    for pH in self.act_prof.keys():
                        if abs(pH-pHset)<0.1:
                            pH_vals.append(pH)
        pH_vals.sort()
        return pH_vals


from einops import rearrange
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
import pandas as pd
from numpy.random import binomial
from torch.cuda.amp import GradScaler, autocast
from advbench.datasets import FFCV_AVAILABLE

from advbench import attacks, networks, optimizers, perturbations
from advbench.lib import meters

ALGORITHMS = [
    'ERM',
    'PGD',
    'FGSM',
    'TRADES',
    'ALP',
    'CLP',
    'Gaussian_DALE',
    'Laplacian_DALE',
    'Discrete_DALE',
    'Gaussian_DALE_PD',
    'Gaussian_DALE_PD_Reverse',
    'MH_DALE_PD_Reverse',
    'KL_DALE_PD',
    'Worst_Of_K',
    'Augmentation',
    'Batch_Augmentation',
    'Grid_Search',
    'Batch_Grid',
    'Uniform_DALE_PD_Reverse',
    'Worst_DALE_PD_Reverse',
    'PGD_DALE_PD_Reverse'
]

class Algorithm(nn.Module):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Algorithm, self).__init__()
        self.hparams = hparams
        self.classifier = networks.Classifier(
            input_shape, num_classes, hparams)
        #summary(self.classifier.to(device), input_size=input_shape)
        if hparams['optimizer']=="SGD":
            self.optimizer = optimizers.Optimizer(
             self.classifier, hparams)
        elif hparams['optimizer']=="SFCNN":
            self.optimizer = optimizers.SFCNN_Optimizer(self.classifier, hparams)
        else:
            print("Optimizer not suported")
            raise NotImplementedError
        self.device = device

        self.meters = OrderedDict()
        self.meters['loss'] = meters.AverageMeter()
        self.meters_df = None
        self.perturbation_name = perturbation
        if FFCV_AVAILABLE:
            self.scaler = GradScaler()

        self.label_smoothing = hparams['label_smoothing']

    def step(self, imgs, labels):
        raise NotImplementedError

    def predict(self, imgs):
        return self.classifier(imgs)

    def reset_meters(self):
        for meter in self.meters.values():
            meter.reset()

    def meters_to_df(self, epoch):
        if self.meters_df is None:
            keys = []
            for key, val in self.meters.items():
                if val.print:
                    keys.append(key)
            columns = ['Epoch'] + keys
            self.meters_df = pd.DataFrame(columns=columns)
            self.meters_df_keys = keys
        metrics = []
        for key in self.meters_df_keys:
            metrics.append(self.meters[key].avg)
        values = [epoch] + metrics
        self.meters_df.loc[len(self.meters_df)] = values
        return self.meters_df
    def export(self):
        pass
    def unexport(self):
        pass

class ERM(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(ERM, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Rand_Aug(self.classifier, self.hparams, device, perturbation=perturbation)
    def step(self, imgs, labels):
        self.optimizer.zero_grad(set_to_none=True)
        if FFCV_AVAILABLE:
            with autocast():
                loss = self.classifier.loss(self.predict(imgs), labels)
                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            loss = self.classifier.loss(self.predict(imgs), labels)
            loss.backward()
        self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Adversarial(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.penalty = hparams["adv_penalty"]

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                adv_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                loss = clean_loss+adv_loss*self.penalty
                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =   self.attack(imgs, labels)
            self.optimizer.zero_grad()
            adv_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            loss = clean_loss+adv_loss*self.penalty
            loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Adversarial_PGD(Adversarial):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial_PGD, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Fo_PGD(self.classifier, self.hparams, device, perturbation=perturbation)

class Adversarial_SGD(Adversarial):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial_SGD, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Fo_SGD(self.classifier, self.hparams, device, perturbation=perturbation)

class Adversarial_Adam(Adversarial):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial_Adam, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Fo_Adam(self.classifier, self.hparams, device, perturbation=perturbation)

class Adversarial_Smoothed(Adversarial):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial_Smoothed, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.LMC_Laplacian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)

class Gaussian_DALE(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Gaussian_DALE, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.LMC_Gaussian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.meters['clean loss'] = meters.AverageMeter()
        self.meters['robust loss'] = meters.AverageMeter()

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                adv_imgs, deltas =   self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = robust_loss + self.hparams['g_dale_nu'] * clean_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =   self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            total_loss = robust_loss + self.hparams['g_dale_nu'] * clean_loss
            total_loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))

class Laplacian_DALE(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Laplacian_DALE, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.LMC_Laplacian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.meters['clean loss'] = meters.AverageMeter()
        self.meters['robust loss'] = meters.AverageMeter()

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas =   self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = robust_loss + self.hparams['g_dale_nu'] * clean_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =   self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            total_loss = robust_loss + self.hparams['l_dale_nu'] * clean_loss
            total_loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))

class PrimalDualBase(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(PrimalDualBase, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.dual_params = {'dual_var': torch.tensor(init).to(self.device)}
        self.meters['clean loss'] = meters.AverageMeter()
        self.meters['robust loss'] = meters.AverageMeter()
        self.meters['dual variable'] = meters.AverageMeter()
        self.meters['delta L1-border'] = meters.AverageMeter()
        perturbation = vars(perturbations)[perturbation](0)
        self.meters['delta hist'] = meters.WBDeltaMeter(names = perturbation.names, dims = perturbation.dim)

class Gaussian_DALE_PD(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=1.0):
        super(Gaussian_DALE_PD, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.LMC_Gaussian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.pd_optimizer = optimizers.PrimalDualOptimizer(
            parameters=self.dual_params,
            margin=self.hparams['g_dale_pd_margin'],
            eta=self.hparams['g_dale_pd_eta'])

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = robust_loss + self.hparams['g_dale_nu'] * clean_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            total_loss = robust_loss + self.dual_params['dual_var'] * clean_loss
            total_loss.backward()
            self.optimizer.step()

        self.pd_optimizer.step(clean_loss.detach())

        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].item(), n=imgs.size(0))
        self.meters['delta L1-border'].update((torch.abs(deltas)-self.hparams['epsilon']).mean().item(), n=imgs.size(0))
        self.meters['delta hist'].update(deltas.cpu())
        #print(deltas[0])

class Gaussian_DALE_PD_Reverse(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=1.0):
        super(Gaussian_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.LMC_Gaussian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.pd_optimizer = optimizers.PrimalDualOptimizer(
            parameters=self.dual_params,
            margin=self.hparams['g_dale_pd_inv_margin'],
            eta=self.hparams['g_dale_pd_inv_eta'])

    def step(self, imgs, labels):
        adv_imgs, deltas =self.attack(imgs, labels)
        self.optimizer.zero_grad()
        clean_loss = self.classifier.loss(self.predict(imgs), labels)
        robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
        total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
        total_loss.backward()
        self.optimizer.step()
        self.pd_optimizer.step(robust_loss.detach())
        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].item(), n=1)
        self.meters['delta L1-border'].update((torch.abs(deltas)-self.hparams['epsilon']).mean().item(), n=imgs.size(0))
        self.meters['delta hist'].update(deltas.cpu())

class Laplacian_DALE_PD_Reverse(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(Laplacian_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.LMC_Laplacian_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.pd_optimizer = optimizers.PrimalDualOptimizer(
            parameters=self.dual_params,
            margin=self.hparams['l_dale_pd_inv_margin'],
            eta=self.hparams['l_dale_pd_inv_eta'])

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
            total_loss.backward()
            self.optimizer.step()
        self.pd_optimizer.step(robust_loss.detach())
        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].item(), n=1)
        self.meters['delta L1-border'].update((torch.abs(deltas)-self.hparams['epsilon']).mean().item(), n=imgs.size(0))
        self.meters['delta hist'].update(deltas.cpu())

class Worst_DALE_PD_Reverse(Laplacian_DALE_PD_Reverse):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(Worst_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.Worst_Of_K(self.classifier, self.hparams, device, perturbation=perturbation)

class PGD_DALE_PD_Reverse(Laplacian_DALE_PD_Reverse):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(PGD_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.Fo_PGD(self.classifier, self.hparams, device, perturbation=perturbation)

class Adam_DALE_PD_Reverse(Laplacian_DALE_PD_Reverse):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(Adam_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.Fo_Adam(self.classifier, self.hparams, device, perturbation=perturbation)

class MH_DALE_PD_Reverse(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(MH_DALE_PD_Reverse, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.meters['acceptance rate'] = meters.AverageMeter()
        self.attack = attacks.MH(self.classifier, self.hparams, device, perturbation=perturbation, acceptance_meter=self.meters['acceptance rate'])
        self.pd_optimizer = optimizers.PrimalDualOptimizer(
            parameters=self.dual_params,
            margin=self.hparams['g_dale_pd_inv_margin'],
            eta=self.hparams['g_dale_pd_inv_eta'])

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
            total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
            total_loss.backward()
            self.optimizer.step()
        self.pd_optimizer.step(robust_loss.detach())
        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].item(), n=1)
        self.meters['delta L1-border'].update((torch.abs(deltas)-self.hparams['epsilon']).mean().item(), n=imgs.size(0))
        self.meters['delta hist'].update(deltas.cpu())

class KL_DALE_PD(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0):
        super(KL_DALE_PD, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.TRADES_Linf(self.classifier, self.hparams, device, perturbation=perturbation)
        self.kl_loss_fn = nn.KLDivLoss(reduction='batchmean')
        self.pd_optimizer = optimizers.PrimalDualOptimizer(
            parameters=self.dual_params,
            margin=self.hparams['g_dale_pd_margin'],
            eta=self.hparams['g_dale_pd_eta'])

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                adv_imgs, deltas = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.kl_loss_fn(
                F.log_softmax(self.predict(adv_imgs), dim=1),
                F.softmax(self.predict(imgs), dim=1))
                total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            adv_imgs, deltas =self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.kl_loss_fn(
                F.log_softmax(self.predict(adv_imgs), dim=1),
                F.softmax(self.predict(imgs), dim=1))
            total_loss = robust_loss + self.dual_params['dual_var'] * clean_loss
            total_loss.backward()
            self.optimizer.step()
        self.pd_optimizer.step(clean_loss.detach())

        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].item(), n=1)

class Adversarial_Worst_Of_K(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Adversarial_Worst_Of_K, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Worst_Of_K(self.classifier, self.hparams, device, perturbation=perturbation)

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                with torch.no_grad():
                    adv_imgs, deltas =   self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), labels)
                total_loss = clean_loss + self.dual_params['dual_var'] * robust_loss
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            with torch.no_grad():
                adv_imgs, deltas =   self.attack(imgs, labels)
            self.optimizer.zero_grad()
            loss = self.classifier.loss(self.predict(adv_imgs), labels)
            loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Grid_Search(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Grid_Search, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Grid_Search(self.classifier, self.hparams, device, perturbation=perturbation)

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                with torch.no_grad():
                    adv_imgs, deltas =   self.attack(imgs, labels)
                    self.optimizer.zero_grad()
                    loss = self.classifier.loss(self.predict(adv_imgs), labels)
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
        else:
            with torch.no_grad():
                adv_imgs, deltas =   self.attack(imgs, labels)
            self.optimizer.zero_grad()
            loss = self.classifier.loss(self.predict(adv_imgs), labels)
            loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Augmentation(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Augmentation, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Rand_Aug(self.classifier, self.hparams, device, perturbation=perturbation)
        self.p = hparams['augmentation_prob']
    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                if binomial(1, self.p):
                    adv_imgs, _ =   self.attack(imgs, labels)
                else:
                    adv_imgs = imgs
                self.optimizer.zero_grad()
                loss = self.classifier.loss(self.predict(adv_imgs), labels)
                self.scaler.scale(loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            if binomial(1, self.p):
                adv_imgs, _ =   self.attack(imgs, labels)
            else:
                adv_imgs = imgs
            self.optimizer.zero_grad()
            loss = self.classifier.loss(self.predict(adv_imgs), labels)
            loss.backward()
            self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Laplacian(Augmentation):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Laplacian, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Laplace_aug(self.classifier, self.hparams, device, perturbation=perturbation)
        self.p = 1

class Gaussian(Augmentation):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Gaussian, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Gaussian_aug(self.classifier, self.hparams, device, perturbation=perturbation)
        self.p = 1

class Batch_Random(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Batch_Random, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Rand_Aug_Batch(self.classifier, self.hparams, device, perturbation=perturbation)
    def step(self, imgs, labels):
        adv_imgs, deltas, new_labels =   self.attack(imgs, labels)
        self.optimizer.zero_grad()
        loss = self.classifier.loss(self.predict(adv_imgs), new_labels)
        loss.backward()
        self.optimizer.step()

        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Batch_Grid(Algorithm):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf'):
        super(Batch_Grid, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation)
        self.attack = attacks.Grid_Batch(self.classifier, self.hparams, device, perturbation=perturbation)

    def step(self, imgs, labels):
        adv_imgs, deltas, new_labels =  self.attack(imgs, labels)
        self.optimizer.zero_grad()
        loss = self.classifier.loss(self.predict(adv_imgs), new_labels)
        loss.backward()
        self.optimizer.step()
        self.meters['loss'].update(loss.item(), n=imgs.size(0))

class Discrete_DALE(PrimalDualBase):
    def __init__(self, input_shape, num_classes, hparams, device, perturbation='Linf', init=0.0, batched=False):
        if perturbation != 'SE':
            raise NotImplementedError
        super(Discrete_DALE, self).__init__(input_shape, num_classes, hparams, device, perturbation=perturbation, init=init)
        self.attack = attacks.Grid_Batch(self.classifier, hparams, device, perturbation=perturbation)
        translations = []
        for idx in (1, 2):
            eps = hparams['epsilon'][idx]
            step = 2*eps/hparams['d_num_translations']
            translations.append(torch.arange(-eps, eps, step=step, device=self.device))
        eps = hparams['epsilon'][0]
        step = 2*eps/hparams['d_num_rotations']
        self.translations = translations
        self.rotation = torch.arange(-eps, eps, step=step, device=self.device)
        grids =  [self.rotation] + translations
        self.grid = torch.cartesian_prod(*grids)
        self.attack.grid = self.grid
        self.attack.grid_size = self.grid.shape[0]
        self.dual_params = {'dual_var': torch.ones(self.grid.shape[0]).to(self.device)*init}
        self.pd_optimizer = optimizers.PrimalDualOptimizer(parameters=self.dual_params,
                                                            margin=self.hparams['d_dale_pd_inv_margin'],
                                                            eta=self.hparams['d_dale_pd_inv_eta'])
        loc0 = (int(translations[0].shape[0]//2), int(translations[1].shape[0]//2))
        # Dual plot logger
        self.meters['dual plot'] = meters.WBDualMeter(self.grid,translations, names = "Dual var vs angle",
                                                         locs = [(0,0), loc0, (-1, -1)])

    def step(self, imgs, labels):
        if FFCV_AVAILABLE:
            with autocast():
                with torch.no_grad():
                    adv_imgs, deltas, new_labels = self.attack(imgs, labels)
                self.optimizer.zero_grad()
                clean_loss = self.classifier.loss(self.predict(imgs), labels)
                robust_loss = self.classifier.loss(self.predict(adv_imgs), new_labels, reduction='none')
                robust_loss = rearrange(robust_loss, '(B S) -> B S', B = imgs.shape[0])
                total_loss = clean_loss +  torch.mean(robust_loss@self.dual_params['dual_var'].to(self.device))
                self.scaler.scale(total_loss).backward()
                self.scaler.step(self.optimizer)
                self.scaler.update()
        else:
            with torch.no_grad():
                adv_imgs, deltas, new_labels =self.attack(imgs, labels)
            self.optimizer.zero_grad()
            clean_loss = self.classifier.loss(self.predict(imgs), labels)
            robust_loss = self.classifier.loss(self.predict(adv_imgs), new_labels, reduction='none')
            robust_loss = rearrange(robust_loss, '(B S) -> B S', B = imgs.shape[0])
            total_loss = clean_loss +  torch.mean(robust_loss@self.dual_params['dual_var'].to(self.device))
            total_loss.backward()
            self.optimizer.step()
        #print("rloss before", robust_loss[:10])
        #print("dual before upd", self.dual_params['dual_var'])
        with torch.no_grad():
            self.pd_optimizer.step(torch.mean(robust_loss, 0).detach())
        #print("rloss", robust_loss[:10])
        #print("dual after upd", self.dual_params['dual_var'][:10])
        #print(f"clean {clean_loss.item()}, robust {robust_loss.mean().item()}, total {total_loss.item()}, dual {self.dual_params['dual_var'].mean().item()}")
        self.meters['loss'].update(total_loss.item(), n=imgs.size(0))
        self.meters['clean loss'].update(clean_loss.item(), n=imgs.size(0))
        self.meters['robust loss'].update(robust_loss.mean().item(), n=imgs.size(0))
        self.meters['dual variable'].update(self.dual_params['dual_var'].mean().item(), n=1)
        self.meters['dual plot'].update(self.dual_params['dual_var'])
        #print("dual after log", self.dual_params['dual_var'][:10])

    def get_grid(self, tx, ty):
        angle_grid = self.attack.grid
        ones = torch.ones_like(angle_grid)
        grid = torch.column_stack([angle_grid, tx*ones, ty*ones])
        return grid


# read snapshot and obtain multiple systems
import collections
from scipy import spatial as sp
from .base import *
from .bse import *

G_MSUN_PC_MYR=0.00449830997959438 # Msun, pc, myr
G_HENON=1 # Henon unit

class PeTarDataHeader():
    """ Petar snapshot data header
    members:
        fid: file id
        n: number of particles
        time: time of snapshot
    """

    def __init__(self, _filename=None):
        """ Initial data header

        Parameters:
        -----------
        _filename: string
            PeTar snapshot file name to read the header, if not provide, all members are initialized to zero (None)
        """
        self.fid = 0
        self.n = 0
        self.time = 0.0

        if (_filename!=None): self.read(_filename)

    def read(self, _filename):
        """ Read snapshot file to obtain the header information

        Parameters:
        -----------
        _filename: string
            PeTar snapshot file name to read the header
        """

        fp = open(_filename, 'r')
        header=fp.readline()
        file_id, n_glb, t = header.split()
        fp.close()

        self.fid = int(file_id)
        self.n = int(n_glb)
        self.time = float(t)

class SimpleParticle(DictNpArrayMix):
    """ Simple particle class with only mass, postion, velocity
    keys: (class members)
        mass (1D): mass
        pos (2D,3): postion x, y, z
        vel (2D,3): velocity vx, vy, vz
    """
    def __init__(self, _dat=None, _offset=int(0), _append=False, **kwargs):
        """ DictNpArrayMix type initialzation, see help(DictNpArrayMix.__init__)
        """
        keys = [['mass',1], ['pos',3], ['vel',3]]
        DictNpArrayMix.__init__(self, keys, _dat, _offset, _append, **kwargs)

    def calcR2(self):
        """ calculate distance square, r2, and add it as a class member
        """
        if (not 'r2' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['r2',1])
        self.r2 = vecDot(self.pos,self.pos)

    def calcEkin(self):
        """ calculate kinetic energy
        """
        if (not 'ekin' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['ekin',1])
        self.ekin = 0.5*vecDot(self.vel,self.vel)*self.mass

    def correctCenter(self, cm_pos, cm_vel):
        self.pos -= cm_pos
        self.vel -= cm_vel

class Particle(SimpleParticle):
    """ Particle class
    keys: (class members)
        The final keys are a combination of sub keys depending on kwargs of initial function

        Sub key list:
        basic: [inherit SimpleParticle]
        add: binary_state: binary interruption state
        se: radius:        (1D): radius for merger checker
            dm:            (1D): mass loss
            time_record    (1D): last time of interruption check
            time_interrupt (1D): next interruption time
        ptcl: r_search (1D): searching radius
              id       (1D): identification
              mass_bk  (1D): artificial particle parameter 1
              status   (1D): artificial particle parameter 2
              r_in     (1D): changeover function inner boundary
              r_out    (1D): changeover function outer boundary
        hermite: dt    (1D): time step
                 time  (1D): current time
                 acc   (2D,3): acceleration x, y, z
                 jerk  (2D,3): acceleration derivative x, y, z
                 pot   (1D): potential
        soft: acc_soft (2D,3): long-range interaction acceleration (particle-tree) x, y, z
              pot      (1D): total potential
              pot_soft (1D): long-range interaction potential
              n_nb:    (1D): number of neighbors (short-interaction)

        Combination:
        ends:
            kwargs['particle_type']:
                hermite:   ptcl + hermite
                hard:      ptcl
                soft (default): ptcl + soft
        Final:
        keys:
            kwargs['interrupt_mode']:
                base:      basic + add + se + ends
                bse:       basic + add + se + ['star',SSEStarParameter] + ends
                none (default): basic + add + ends

    """

    def __init__ (self, _dat=None, _offset=int(0), _append=False, **kwargs):
        """ DictNpArrayMix type initialzation, see help(DictNpArrayMix.__init__)

        Parameters
        ----------
        keyword arguments:
            particle_type: basic particle type: hermite, hard, soft (soft)
            interrupt_mode: PeTar interrupt mode: base, bse, none (none)
        """

        keys_add = [['binary_state',1]]
        keys_se  = [['radius',1],['dm',1],['time_record',1],['time_interrupt',1]]
        keys_ptcl_add = [['r_search',1], ['id',1], ['mass_bk',1], ['status',1], ['r_in',1], ['r_out',1]]
        keys_hermite_add = [['dt',1],['time',1],['acc',3],['jerk',3],['pot',1]]
        keys_soft_add = [['acc_soft',3], ['pot',1], ['pot_soft',1], ['n_nb',1]]
        keys_end =  keys_ptcl_add + keys_soft_add
        if ('particle_type' in kwargs.keys()):
            if (kwargs['particle_type']=='hermite'):
                keys_end = keys_ptcl_add + keys_hermite_add
            elif (kwargs['particle_type']=='hard'):
                keys_end = keys_ptcl_add
        keys=keys_add+keys_end
        if ('interrupt_mode' in kwargs.keys()):
            if (kwargs['interrupt_mode']=='base'):
                keys = keys_add+keys_se+keys_end
            elif (kwargs['interrupt_mode']=='bse'):
                keys = keys_add+keys_se+[['star',SSEStarParameter]]+keys_end

        SimpleParticle.__init__(self, _dat, _offset, _append, **kwargs)
        DictNpArrayMix.__init__(self, keys, _dat, _offset+self.ncols, True, **kwargs)

    def calcEtot(self):
        """ Calculate total energy and add it as the member, etot
        """
        if (not 'etot' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['etot',1])
        self.etot = self.ekin + self.mass*self.pot

def calculateParticleCMDict(pcm, _p1, _p2):
    """ Calculate the center-of-the-mass of two particle sets

    Parameters
    ----------
    _p1: inherited SimpleParticle
        particle set 1
    _p2: inherited SimpleParticle
        particle set 2, should have the same size as _p1
    pcm: dict
        particle center-of-the-mass, should include keys: 'mass','pos','vel'.
    """
    if (issubclass(type(_p1), SimpleParticle)) & (issubclass(type(_p2),SimpleParticle)):
        pcm['mass'] = _p1.mass + _p2.mass
        pcm['pos']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1.mass, _p1.pos, _p2.mass, _p2.pos)))
        pcm['vel']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1.mass, _p1.vel, _p2.mass, _p2.vel)))
    elif (isinstance(_p1, collections.OrderedDict)) & (isinstance(_p2,collections.OrderedDict)) | (isinstance(_p1, dict)) & (isinstance(_p2, dict)):
        pcm['mass'] = _p1['mass'] + _p2['mass']
        pcm['pos']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1['mass'], _p1['pos'], _p2['mass'], _p2['pos'])))
        pcm['vel']  = np.array(list(map(lambda m1,x1,m2,x2:(m1*x1+m2*x2)/(m1+m2), _p1['mass'], _p1['vel'], _p2['mass'], _p2['vel'])))
    else:
        raise ValueError('Initial fail, date type should be Particle or collections.OrderDict, given',type(_p1))

class Binary(DictNpArrayMix):
    """ Binary class
    Keys:
        The final keys depends on kwargs of initial function

        kwargs['simple_mode'] (bool)
            True: (default)
                mass (1D): total mass of two components
                pos  (2D,3): c.m. position x, y, z
                vel  (2D,3): c.m. velocity vx, vy, vz
                rrel (1D): relative distance
                semi (1D): semi-major axis
                ecc  (1D): eccentricity
                p1   (member_particle_type) component one
                p2   (member_particle_type) component two
            False:
                mass (1D): total mass of two components
                pos  (2D,3): c.m. position x, y, z
                vel  (2D,3): c.m. velocity vx, vy, vz
                m1   (1D): component 1 mass
                m2   (1D): component 2 mass
                rrel (1D): relative distance
                semi (1D): semi-major axis
                am   (2D,3): specific angular momemtum x, y, z
                L    (2D,3): angular momemtum x, y, z
                eccvec  (2D,3): eccentric vector
                incline (1D): inclination
                rot_horizon (1D): frame rotational angle in x-y plane (longitude of ascending node)
                ecc  (1D): eccentricity
                rot_self (1D): frame rotational angle in orbital plane (argument of periapsis)
                ecca (1D): eccentric anomaly
                period (1D): period
                t_peri (1D): time to peri-center
                p1 (member_particle_type) component one
                p2 (member_particle_type) component two

        member_particle_type is given by kwargs['member_particle_type'], in default it is SimpleParticle

    """
    def __init__ (self, _p1=None, _p2=None, _offset=int(0), _append=False, **kwargs):
        """
        Parameters
        ----------
        _p1: inherited SimpleParticle | 2D numpy.ndarray | Binary | None
            If the type is inherited SimpleParticle, it is the first component of binary (_p2 should be the same type).
            If the type is Binary, the class instance is initialized by copy the data of _p1.
            If it is None, initialize class with empty data
        _p2: inherited SimpleParticle | None
            If the type is inherited SimpleParticle, it is the second component of binary
            If it is None, _p1 should be either 2D numpy.ndarray or Bina
        _offset: int (0)
            Reading column offset of _dat if it is 2D np.ndarray
        _append: bool (False)
            If true, append keys and ncols to the current class instead of create new class members
        kwaygs: dict ()
            keyword arguments:
                simple_mode: only calculate semi and ecc, save computing time significantly (True)
                G: gravitational constant (1.0)
                member_particle_type: type of component particle (SimpleParticle)
        """
        G=1
        simple_mode=True
        member_particle_type=SimpleParticle

        if 'G' in kwargs.keys(): G=kwargs['G']
        if 'simple_mode' in kwargs.keys(): simple_mode=kwargs['simple_mode']
        if 'member_particle_type' in kwargs.keys(): member_particle_type=kwargs['member_particle_type']

        if (issubclass(type(_p1), SimpleParticle)) & (issubclass(type(_p2),SimpleParticle)):
            member_particle_type = type(_p1)
            if (simple_mode):
                self.keys = [['mass',1],['pos',3],['vel',3],['rrel',1],['semi',1],['ecc',1],['p1',member_particle_type], ['p2', member_particle_type]]
                self.particleToSemiEcc(_p1, _p2, G)
                self.ncols= int(10)
            else:
                self.keys = [['mass',1],['pos',3],['vel',3],['m1',1],['m2',1],['rrel',1],['semi',1],['am',3],['L',3],['eccvec',3],['incline',1],['rot_horizon',1],['ecc',1],['rot_self',1],['ecca',1],['period',1],['t_peri',1],['p1', member_particle_type],['p2', member_particle_type]]
                self.particleToBinary(_p1, _p2, G)
                self.ncols= int(27)
            self.p1 = _p1
            self.p2 = _p2
            self.size = _p1.size
            self.ncols += self.p1.ncols + self.p2.ncols
        elif (_p2==None):
            if (simple_mode):
                keys = [['mass',1],['pos',3],['vel',3],['rrel',1],['semi',1],['ecc',1],['p1',member_particle_type], ['p2', member_particle_type]]
                DictNpArrayMix.__init__(self, keys, _p1, _offset, _append, **kwargs)
            else:
                keys=[['mass',1],['pos',3],['vel',3],['m1',1],['m2',1],['rrel',1],['semi',1],['am',3],['L',3],['eccvec',3],['incline',1],['rot_horizon',1],['ecc',1],['rot_self',1],['ecca',1],['period',1],['t_peri',1],['p1', member_particle_type],['p2', member_particle_type]]
                DictNpArrayMix.__init__(self, keys, _p1, _offset, _append, **kwargs)
        else:
            raise ValueError('Initial fail, date type should be Particle (2), Binary (1) or no argument (0)')
        self.initargs = kwargs.copy()
        self.initargs['G'] = G

    def calcEkin(self):
        """ Calculate c.m. kinetic energy, ekin, and add it as a member
        """
        if (not 'ekin' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['ekin',1])
        self.ekin = 0.5*vecDot(self.vel,self.vel)*self.mass

    def calcEtot(self):
        """ Calculate c.m. total energy (binary energy is excluded) , etot, and add it as a member
        """
        if (not 'etot' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['etot',1])
        self.etot = self.ekin + self.mass*self.pot

    def calcR2(self, member_also=False):
        """ Calculate c.m. distance square, r2, and add it as a member
        """
        if (not 'r2' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['r2',1])
        self.r2 = vecDot(self.pos,self.pos)
        if (member_also):
            ncols = self.p1.ncols + self.p2.ncols
            self.p1.calcR2()
            self.p2.calcR2()
            ncols = self.p1.ncols + self.p2.ncols - ncols
            self.ncols += ncols

    def calcEbin(self):
        """ Calculate binding energy, ebin, and add it as a member
            Notice G should be given the correct value in initialization (keyword argument 'G')
        """
        if (not 'ebin' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['ebin',1])
        self.ebin = self.initargs['G']*self.p1.mass*self.p2.mass/(2*self.semi)

    def calcPot(self):
        """ Calculate potential of c.m., pot, and add it as a member
            Notice G should be given the correct value in initialization (keyword argument 'G')
        """
        G = self.initargs['G']
        pos_b1 = self.p1.pos
        pos_b2 = self.p2.pos
        m_b1 = self.p1.mass
        m_b2 = self.p2.mass
        dr = pos_b1-pos_b2
        dr2 = vecDot(dr,dr)
        invr = 1/np.sqrt(dr2)
        pot_b1 = self.p1.pot + G*m_b2*invr
        pot_b2 = self.p2.pot + G*m_b1*invr
        if (not 'pot' in self.__dict__.keys()):
            self.ncols += 1
            self.keys.append(['pot',1])
        self.pot = (m_b2*pot_b1 + m_b1*pot_b2)/self.mass

    def correctCenter(self, cm_pos, cm_vel):
        """ Corrent c.m and component position and velocity by subtracting cm_pos and cm_vel
        """
        self.pos -= cm_pos
        self.vel -= cm_vel
        self.p1.correctCenter(cm_pos, cm_vel)
        self.p2.correctCenter(cm_pos, cm_vel)

    def particleToSemiEcc(self, _p1, _p2, _G):
        """ Calculate relative distance, semi-major axis and eccentricity from particle pairs

        Parameters
        ----------
        _p1, _p2: inherited SimpleParticle
            Particle pair data set
        _G: float
            Gravitational constant

        """
        calculateParticleCMDict(self.__dict__, _p1, _p2)

        dr = (_p1.pos - _p2.pos)
        dv = (_p1.vel - _p2.vel)

        dr2  = (dr*dr).sum(axis=1)
        dv2  = (dv*dv).sum(axis=1)
        rvdot= (dr*dv).sum(axis=1)

        dr   = np.sqrt(dr2)
        m    = (_p1.mass+_p2.mass)
        semi = 1.0/(2.0/dr - dv2/(_G*m))

        dr_semi = 1.0 - dr/semi
        ecc = np.sqrt(dr_semi*dr_semi + rvdot*rvdot/(_G*m*semi))

        self.rrel = dr
        self.semi = semi
        self.ecc  = ecc

    def particleToBinary(self, _p1, _p2, _G):
        """ Calculate binary orbit from particle pairs

        Parameters
        ----------
        _p1, _p2: inherited SimpleParticle
            Particle pair data set
        _G: float
            Gravitational constant

        """
        binary=self.__dict__

        def regular_sign(_a,_a_err):
            _a[(_a<0) & (_a>-_a_err)] *= -1

        f_err = 1e-2
        calculateParticleCMDict(binary, _p1, _p2)

        binary['m1'] = _p1.mass
        binary['m2'] = _p2.mass
        m_tot = binary['mass']
        Gm_tot = _G*m_tot

        dx = _p1.pos-_p2.pos
        dv = _p1.vel-_p2.vel
        dr2  = vecDot(dx,dx)
        dv2  = vecDot(dv,dv)
        rvdot= vecDot(dx,dv)
        dr   = np.sqrt(dr2)
        binary['rrel'] = np.sqrt(dr2)

        inv_dr = 1.0 / binary['rrel']
        binary['semi'] = 1.0 / (2.0*inv_dr - dv2 / Gm_tot)
        binary['am'] = np.array(list(map(lambda x,y:np.cross(x,y),dx,dv)))
        dp = np.array(list(map(lambda m1,x1,m2,x2:m1*x1-m2*x2,_p1.mass,_p1.vel,_p2.mass,_p2.vel)))
        binary['L'] = np.array(list(map(lambda x,y:np.cross(x,y),dx,dp)))
        binary['eccvec'] = np.array(list(map(lambda v,am,gm,dx,dr:np.cross(v,am)/gm-dx/dr,dv,binary['am'],Gm_tot,dx,dr)))

        binary['incline'] = np.arctan2(np.sqrt(binary['am'][:,0]*binary['am'][:,0]+binary['am'][:,1]*binary['am'][:,1]),binary['am'][:,2])
        binary['rot_horizon'] = np.arctan2(binary['am'][:,0],-binary['am'][:,1])
        regular_sign(binary['am'][:,0],f_err)
        regular_sign(binary['am'][:,1],f_err)
        #binary['rot_horizon'][binary['rot_horizon']<0] += np.pi
        binary['rot_horizon'][binary['am'][:,1]==0.0]=0.0

        cosOMG = np.cos(binary['rot_horizon'])
        sinOMG = np.sin(binary['rot_horizon'])
        cosinc = np.cos(binary['incline'])
        sininc = np.sin(binary['incline'])

        pos_bar_x =   dx[:,0]*cosOMG + dx[:,1]*sinOMG
        pos_bar_y = (-dx[:,0]*sinOMG + dx[:,1]*cosOMG)*cosinc + dx[:,2]*sininc
        pos_bar_z = 0.0
        vel_bar_x =   dv[:,0]*cosOMG + dv[:,1]*sinOMG
        vel_bar_y = (-dv[:,0]*sinOMG + dv[:,1]*cosOMG)*cosinc + dv[:,2]*sininc
        vel_bar_z = 0.0

        h = np.array(list(map(lambda x:np.sqrt(np.inner(x,x)),binary['am'])))
        ecccosomg =  h/Gm_tot*vel_bar_y - pos_bar_x*inv_dr
        eccsinomg = -h/Gm_tot*vel_bar_x - pos_bar_y*inv_dr
        binary['ecc'] = np.sqrt( ecccosomg*ecccosomg + eccsinomg*eccsinomg )
        regular_sign(ecccosomg,f_err)
        regular_sign(eccsinomg,f_err)
        binary['rot_self'] = np.arctan2(eccsinomg,ecccosomg)
        #binary['rot_self'][binary['rot_self']<-np.pi+1e-5] += 2*np.pi
        #binary['rot_self'][binary['rot_self']>=np.pi-1e-5] -= 2*np.pi

        regular_sign(pos_bar_y,f_err)
        regular_sign(pos_bar_x,f_err)
        phi = np.arctan2(pos_bar_y, pos_bar_x)
        #phi[phi<-np.pi+1e-5] += 2*np.pi
        #phi[phi>=np.pi-1e-5] -= 2*np.pi

        f = phi - binary['rot_self']
        binary['ecca'] = np.arctan(np.sin(f)*np.sqrt(np.abs(binary['ecc']*binary['ecc'] - 1.0))/(binary['ecc']+np.cos(f)))
        n = np.sqrt(Gm_tot/np.abs(binary['semi']*binary['semi']*binary['semi']))
        binary['period'] = 8.0*np.arctan(1.0)/n
        l = binary['ecca'] - binary['ecc']*np.sin(binary['ecca'])
        binary['t_peri'] = l / n

def findPair(_dat, _G, _rmax, use_kdtree=False, simple_binary=True):
    """  Find binaries in a particle data set
    The scipy.spatial.cKDTree is used to find pairs

    Parameters
    ----------
    _dat: inhermited SimpleParticle
        Particle data set
    _G: float
        Gravitational constant
    _rmax: float
        Maximum binary separation
    use_kdtree: bool (False)
        If True, use KDtree to find all binaries (slow); otherwise use information from PeTar, only hard binaries are detected (fast)
    simple_binary: bool (True)
        If True, only calculate semi and ecc (fast); otherwise calculating all binary parameters (slow)

    Return
    ----------
    kdt: KDtree structure if use_kdtree=True
    single: type of _dat
        single particle data set
    binary: Binary(simple_mode=simple_binary, member_particle_type=type(single), G=_G)
        binary data set
    """
    if (not issubclass(type(_dat), SimpleParticle)):
        raise ValueError("Data type wrong",type(_dat)," should be subclass of ", SimpleParticle)

    if (use_kdtree):
        # create KDTree
        #print('create KDTree')
        kdt=sp.cKDTree(_dat.pos)

        # find all close pairs
        #pairs=kdt.query_pairs(_rmax*AU2PC)

        # only check nearest index
        #pair_index=np.unique(np.transpose(np.array([np.array([x[0],x[1]]) for x in pairs])),axis=0)

        # find pair index and distance
        #print('Get index')
        r,index=kdt.query(_dat.pos,k=2)
        pair_index=np.transpose(np.unique(np.sort(index,axis=1),axis=0))
        #pair_index = np.transpose(index)

        #index = kdt.query_pairs(_rmax,output_type='ndarray')
        #pair_index = np.transpose(index)

        # two members
        p1 = _dat[pair_index[0]]
        p2 = _dat[pair_index[1]]

        # check orbits
        #print('Create binary')
        binary = Binary(p1, p2, G=_G, simple_mode=simple_binary)
        apo =binary.semi*(binary.ecc+1.0)

        bsel= ((binary.semi>0) & (apo<_rmax))
        binary = binary[bsel]

        single_mask = np.ones(_dat.size).astype(bool)
        single_mask[pair_index[0][bsel]]=False
        single_mask[pair_index[1][bsel]]=False
        single = _dat[single_mask]
        return kdt, single, binary
    else:
        idx = _dat.status.argsort()
        dat_sort = _dat[idx]
        status, index, inverse, counts = np.unique(dat_sort.status, return_index=True, return_inverse=True, return_counts=True)
        binary_i1 = index[counts==2]
        binary_i2 = binary_i1+1
        binary = Binary(dat_sort[binary_i1], dat_sort[binary_i2], _G)
        single = dat_sort[index[-1]:]

        return single, binary


"""
Module for alignment of variants to the wild type sequence.

This module is optional, and using it will dramatically increase runtime when
counting variants. It is only recommended for users who need to count
insertion and deletion variants (i.e. not coding sequences).
"""

import numpy as np

#: Default similarity matrix used by the aligner.
#: User-defined matrices must have this format.
_simple_similarity = {
    "A": {"A": 1, "C": -1, "G": -1, "T": -1, "N": 0, "X": 0},
    "C": {"A": -1, "C": 1, "G": -1, "T": -1, "N": 0, "X": 0},
    "G": {"A": -1, "C": -1, "G": 1, "T": -1, "N": 0, "X": 0},
    "T": {"A": -1, "C": -1, "G": -1, "T": 1, "N": 0, "X": 0},
    "N": {"A": 0, "C": 0, "G": 0, "T": 0, "N": 0, "X": 0},
    "X": {"A": 0, "C": 0, "G": 0, "T": 0, "N": 0, "X": 0},
    "gap": -1,
}

class Aligner(object):
    """
    Class for performing local alignment of two DNA sequences.

    This class implements `Needleman-Wunsch <http://en.wikipedia.org/wiki/
    Needleman%E2%80%93Wunsch_algorithm>`_ local alignment.

    The :py:class:`~aligner.Aligner` requires a scoring matrix when
    created. The format is a nested dictionary, with a special ``'gap'`` entry
    for the gap penalty (this value is used for both gap opening and gap
    extension).

    The ``'X'`` nucleotide is a special case for unresolvable mismatches in
    :py:class:`~overlap.OverlapSeqLib` variant data.
    """

    _MAT = 1  # match
    _INS = 2  # insertion (with respect to wild type)
    _DEL = 3  # deletion (with respect to wild type)
    _END = 4  # end of traceback

    def __init__(self, similarity=_simple_similarity):
        similarity_keys = similarity.keys()
        if "gap" in similarity_keys:
            similarity_keys.remove("gap")
        for key in similarity_keys:
            if not all(x in similarity[key] for x in similarity_keys) or len(
                similarity[key]
            ) != len(similarity_keys):
                raise ValueError("Asymmetrical alignment scoring matrix")

        self.similarity = similarity
        if "gap" not in self.similarity:
            raise ValueError("No gap penalty in alignment scoring matrix.")

        self.matrix = None
        self.seq1 = None
        self.seq2 = None
        self.calls = 0

    def align(self, seq1, seq2):
        """
        Aligns the two sequences, *seq1* and *seq2* and returns a list of
        tuples describing the differences between the sequences.

        The tuple format is ``(i, j, type, length)``, where ``i`` and ``j``
        are the positions in *seq1* and *seq2*, respectively, and type is one
        of ``"match"``, ``"mismatch"``, ``"insertion"``, or ``"deletion"``.
        For indels, the ``length`` value is the number of bases inserted or
        deleted with respect to *seq1* starting at ``i``.
        """
        self.matrix = np.ndarray(
            shape=(len(seq1) + 1, len(seq2) + 1),
            dtype=np.dtype([("score", np.int), ("trace", np.byte)]),
        )
        seq1 = seq1.upper()
        seq2 = seq2.upper()

        # build matrix of scores/traceback information
        for i in xrange(len(seq1) + 1):
            self.matrix[i, 0] = (self.similarity["gap"] * i, Aligner._DEL)
        for j in xrange(len(seq2) + 1):
            self.matrix[0, j] = (self.similarity["gap"] * j, Aligner._INS)
        for i in xrange(1, len(seq1) + 1):
            for j in xrange(1, len(seq2) + 1):
                match = (
                    self.matrix[i - 1, j - 1]["score"]
                    + self.similarity[seq1[i - 1]][seq2[j - 1]],
                    Aligner._MAT,
                )
                delete = (
                    self.matrix[i - 1, j]["score"] + self.similarity["gap"],
                    Aligner._DEL,
                )
                insert = (
                    self.matrix[i, j - 1]["score"] + self.similarity["gap"],
                    Aligner._INS,
                )
                self.matrix[i, j] = max(delete, insert, match, key=lambda x: x[0])
        self.matrix[0, 0] = (0, Aligner._END)

        # calculate alignment from the traceback
        i = len(seq1)
        j = len(seq2)
        traceback = list()
        while i > 0 or j > 0:
            if self.matrix[i, j]["trace"] == Aligner._MAT:
                if seq1[i - 1] == seq2[j - 1]:
                    traceback.append((i - 1, j - 1, "match", None))
                else:
                    traceback.append((i - 1, j - 1, "mismatch", None))
                i -= 1
                j -= 1
            elif self.matrix[i, j]["trace"] == Aligner._INS:
                traceback.append((i - 1, j - 1, "insertion", 1))
                j -= 1
            elif self.matrix[i, j]["trace"] == Aligner._DEL:
                traceback.append((i - 1, j - 1, "deletion", 1))
                i -= 1
            elif self.matrix[i, j]["trace"] == Aligner._END:
                pass
            else:
                raise RuntimeError("Invalid value in alignment traceback.")
        traceback.reverse()

        # combine indels
        indel = None
        traceback_combined = list()
        for t in traceback:
            if t[2] == "insertion" or t[2] == "deletion":
                if indel is not None:
                    if t[2] == indel[2]:
                        indel[3] += t[3]
                    else:
                        raise RuntimeError(
                            "Aligner failed to combine indels. " "Check gap penalty."
                        )
                else:
                    indel = list(t)
            else:
                if indel is not None:
                    traceback_combined.append(tuple(indel))
                    indel = None
                traceback_combined.append(t)
        if indel is not None:
            traceback_combined.append(tuple(indel))

        self.calls += 1
        return traceback_combined