from copy import deepcopy
import numpy as np
from prospect.models import priors, sedmodel
from prospect.sources import CSPSpecBasis

from sedpy.observate import load_filters

# Here we are going to put together some filter names
# All these filters are available in sedpy.  If you want to use other filters,
# add their transmission profiles to sedpy/sedpy/data/filters/ with appropriate
# names (and format).  See sedpy documentation
galex = ['galex_FUV', 'galex_NUV']
sdss = ['sdss_{0}0'.format(b) for b in ['u','g','r','i','z']]
twomass = ['twomass_{}'.format(b) for b in ['J', 'H', 'Ks']]
spitzer = ['spitzer_irac_ch'+n for n in ['1','2','3','4']]

# --------------
# RUN_PARAMS
# --------------

run_params = {'verbose':True,
              'debug':False,
              'outfile':'output/demo_mock',
              'output_pickles': False,
              # Optimization parameters
              'do_powell': False,
              'ftol':0.5e-5, 'maxfev':5000,
              'do_levenberg': True,
              'nmin': 10,
              # emcee Fitter parameters
              'nwalkers':64,
              'nburn': [32, 32, 64],
              'niter': 256,
              'interval': 0.25,
              'initial_disp':0.1,
              # dynesty Fitter parameters
              'nested_bound': 'multi', # bounding method
              'nested_sample': 'unif', # sampling method
              'nested_nlive_init': 100,
              'nested_nlive_batch': 100,
              'nested_bootstrap': 0,
              'nested_dlogz_init': 0.05,
              'nested_weight_kwargs': {"pfrac": 1.0},
              'nested_stop_kwargs': {"post_thresh": 0.05},
              # Mock data parameters
              'snr': 20.0,
              'add_noise': False,
              'filterset': galex + sdss + twomass,
              # Input mock model parameters
              'mass': 1e10,
              'logzsol': -0.5,
              'tage': 12.,
              'tau': 3.,
              'dust2': 0.3,
              'zred': 0.1,
              'add_neb': False,
              # Data manipulation parameters
              'logify_spectrum':False,
              'normalize_spectrum':False,
              # SPS parameters
              'zcontinuous': 1,
              }

# --------------
# OBS
# --------------

def load_obs(snr=10.0, filterset=["sdss_g0", "sdss_r0"],
             add_noise=True, **kwargs):
    """Make a mock dataset.  Feel free to add more complicated kwargs, and put
    other things in the run_params dictionary to control how the mock is
    generated.

    :param snr:
        The S/N of the phock photometry.  This can also be a vector of same
        lngth as the number of filters.

    :param filterset:
        A list of `sedpy` filter names.  Mock photometry will be generated
        for these filters.

    :param add_noise: (optional, boolean, default: True)
        If True, add a realization of the noise to the mock spectrum
    """
    # We'll put the mock data in this dictionary, just as we would for real
    # data.  But we need to know which bands (and wavelengths if doing
    # spectroscopy) in which to generate mock data.
    mock = {}
    mock['wavelength'] = None # No spectrum
    mock['filters'] = load_filters(filterset)
    #

    # We need the models to make a mock
    sps = load_sps(**kwargs)
    mod = load_model(**kwargs)

    # Now we get the mock params from the kwargs dict
    params = {}
    for p in mod.params.keys():
        if p in kwargs:
            params[p] = np.atleast_1d(kwargs[p])

    # And build the mock
    mod.params.update(params)
    spec, phot, _ = mod.mean_model(mod.theta, mock, sps=sps)
    # Now store some output
    mock['true_spectrum'] = spec.copy()
    mock['true_maggies'] = phot.copy()
    mock['mock_params'] = deepcopy(mod.params)
    # And add noise
    pnoise_sigma = phot / snr
    if add_noise:
        pnoise = np.random.normal(0, 1, len(phot)) * pnoise_sigma
        mock['maggies'] = phot + pnoise
    else:
        mock['maggies'] = phot.copy()
    mock['maggies_unc'] = pnoise_sigma
    mock['mock_snr'] = snr
    mock['phot_mask'] = np.ones(len(phot), dtype=bool)

    return mock

# --------------
# SPS Object
# --------------

def load_sps(zcontinuous=1, **extras):
    """Instantiate and return the Stellar Population Synthesis object.

    :param zcontinuous: (default: 1)
        python-fsps parameter controlling how metallicity interpolation of the
        SSPs is acheived.  A value of `1` is recommended.
        * 0: use discrete indices (controlled by parameter "zmet")
        * 1: linearly interpolate in log Z/Z_\sun to the target metallicity
             (the parameter "logzsol".)
        * 2: convolve with a metallicity distribution function at each age.
             The MDF is controlled by the parameter "pmetals"
    """
    sps = CSPSpecBasis(zcontinuous=zcontinuous,
                       compute_vega_mags=False)
    return sps

# -----------------
# Noise Model
# ------------------

def load_gp(**extras):
    return None, None

# --------------
# MODEL_PARAMS
# --------------

def load_model(zred=0.0, add_neb=True, **extras):
    """Instantiate and return a ProspectorParams model subclass.

    :param zred: (optional, default: 0.1)
        The redshift of the model

    :param add_neb: (optional, default: False)
        If True, turn on nebular emission and add relevant parameters to the
        model.
    """

    # --- Get a basic delay-tau SFH parameter set. ---
    # This has 5 free parameters:
    #   "mass", "logzsol", "dust2", "tage", "tau"
    # And two fixed parameters
    #   "zred"=0.1, "sfh"=4
    # See the python-FSPS documentation for details about most of these
    # parameters.  Also, look at `TemplateLibrary.describe("parameteric")` to
    # view the parameters, their initial values, and the priors in detail
    from prospect.models.templates import TemplateLibrary
    model_params = TemplateLibrary["parametric_sfh"]

    # --- Adjust the basic model ----
    # Add burst parameters (fixed to zero be default)
    model_params.update(TemplateLibrary["burst_sfh"])
    # Add dust emission parameters (fixed)
    model_params.update(TemplateLibrary["dust_emission"])
    # Add nebular emission parameters and turn nebular emission on
    if add_neb:
        model_params.update(TemplateLibrary["nebular"])

    # --- Set dispersions for emcee ---
    model_params["mass"]["init_disp"] = 1e8
    model_params["mass"]["disp_floor"] = 1e7

    # --- Complexify dust attenuation ---
    # Switch to Kriek and Conroy 2013
    model_params["dust_type"] = {'N': 1, 'isfree': False,
                                 'init': 4, 'prior': None}
    # Slope of the attenuation curve, expressed as the index of the power-law
    # that modifies the base Kriek & Conroy/Calzetti shape.
    # I.e. a value of zero is basically calzetti with a 2175AA bump
    model_params["dust_index"] = {'N': 1, 'isfree': False,
                                 'init': 0.0, 'prior': None}

    # --- Set initial values ---
    model_params["zred"]["init"] = zred

    return sedmodel.SedModel(model_params)

# -------------
# Old-style Hand constructed parameters
#    These are deprecated and can be ignored, although the comments might be
#    useful
# -------------

# You'll note below that we have 5 free parameters:
#   "mass", "logzsol", "tage", "tau", "dust2"
# They are all scalars (N=1).
# The "mass" and "tau" parameters have logUniform priors (i.e. TopHat priors in
# log(mass) and log(tau)), the rest have TopHat priors.
# You should adjust the prior ranges (particularly in mass) to suit your objects.
#
# The other parameters are all fixed, but we may want to explicitly set their
# values, which can be done here, to override any defaults in python-FSPS

model_params = []

# --- Distance ---
model_params.append({'name': 'zred', 'N': 1,
                        'isfree': False,
                        'init': 0.1,
                        'units': '',
                        'prior': priors.TopHat(mini=0.0, maxi=4.0)})

# --- SFH --------
# FSPS parameter
model_params.append({'name': 'sfh', 'N': 1,
                        'isfree': False,
                        'init': 4,  # This is delay-tau
                        'units': 'type',
                        'prior': None})

model_params.append({'name': 'mass', 'N': 1,
                        'isfree': True,
                        'init': 1e10,
                        'init_disp': 1e9,
                        'units': r'M_\odot',
                        'prior': priors.LogUniform(mini=1e7, maxi=1e12)})

model_params.append({'name': 'logzsol', 'N': 1,
                        'isfree': True,
                        'init': -0.3,
                        'init_disp': 0.3,
                        'units': r'$\log (Z/Z_\odot)$',
                        'prior': priors.TopHat(mini=-2.0, maxi=0.19)})

# If zcontinuous > 1, use 3-pt smoothing
model_params.append({'name': 'pmetals', 'N': 1,
                        'isfree': False,
                        'init': -99,
                        'prior': None})

# FSPS parameter
model_params.append({'name': 'tau', 'N': 1,
                        'isfree': True,
                        'init': 1.0,
                        'init_disp': 0.5,
                        'units': 'Gyr',
                        'prior':priors.LogUniform(mini=0.101, maxi=100)})

# FSPS parameter
model_params.append({'name': 'tage', 'N': 1,
                        'isfree': True,
                        'init': 5.0,
                        'init_disp': 3.0,
                        'units': 'Gyr',
                        'prior': priors.TopHat(mini=0.101, maxi=14.0)})

model_params.append({'name': 'fage_burst', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': 'time at wich burst happens, as a fraction of `tage`',
                        'prior': priors.TopHat(mini=0.9, maxi=1.0)})

# This function transfroms from a fractional age of a burst to an absolute age.
# With this transformation one can sample in ``fage_burst`` without worry about
# the case tburst > tage.
def tburst_fage(tage=0.0, fage_burst=0.0, **extras):
    return tage * fage_burst

# FSPS parameter
model_params.append({'name': 'tburst', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': '',
                        'prior': None,})
                        #'depends_on': tburst_fage})  # uncomment if using bursts.

# FSPS parameter
model_params.append({'name': 'fburst', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': '',
                        'prior': priors.TopHat(mini=0.0, maxi=0.5)})

# --- Dust ---------
# FSPS parameter
model_params.append({'name': 'dust_type', 'N': 1,
                        'isfree': False,
                        'init': 0,  # power-laws
                        'units': 'index',
                        'prior': None})
# FSPS parameter
model_params.append({'name': 'dust2', 'N': 1,
                        'isfree': True,
                        'init': 0.35,
                        'reinit': True,
                        'init_disp': 0.3,
                        'units': 'Diffuse dust optical depth towards all stars at 5500AA',
                        'prior': priors.TopHat(mini=0.0, maxi=2.0)})

# FSPS parameter
model_params.append({'name': 'dust1', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': 'Extra optical depth towards young stars at 5500AA',
                        'prior': None,})

# FSPS parameter
model_params.append({'name': 'dust_tesc', 'N': 1,
                        'isfree': False,
                        'init': 7.0,
                        'units': 'definition of young stars for the purposes of the CF00 dust model, log(Gyr)',
                        'prior': None,})

# FSPS parameter
model_params.append({'name': 'dust_index', 'N': 1,
                        'isfree': False,
                        'init': -0.7,
                        'units': 'power law slope of the attenuation curve for diffuse dust',
                        'prior': None,})

# FSPS parameter
model_params.append({'name': 'dust1_index', 'N': 1,
                        'isfree': False,
                        'init': -1.0,
                        'units': 'power law slope of the attenuation curve for young-star dust',
                        'prior': None,})

# FSPS parameter
model_params.append({'name': 'add_dust_emission', 'N': 1,
                        'isfree': False,
                        'init': True,
                        'units': 'index',
                        'prior': None})

# An example of the parameters controlling the dust emission SED.  There are others!
model_params.append({'name': 'duste_umin', 'N': 1,
                        'isfree': False,
                        'init': 1.0,
                        'units': 'MMP83 local MW intensity',
                        'prior': None})

# --- Stellar Pops ------------
# One could imagine changing these, though doing so *during* the fitting will
# be dramatically slower.
# FSPS parameter
model_params.append({'name': 'tpagb_norm_type', 'N': 1,
                        'isfree': False,
                        'init': 2,
                        'units': 'index',
                        'prior': None})

# FSPS parameter
model_params.append({'name': 'add_agb_dust_model', 'N': 1,
                        'isfree': False,
                        'init': True,
                        'units': 'index',
                        'prior': None})

# FSPS parameter
model_params.append({'name': 'agb_dust', 'N': 1,
                        'isfree': False,
                        'init': 1,
                        'units': 'index',
                        'prior': None})

# --- Nebular Emission ------

# For speed we turn off nebular emission in the demo
model_params.append({'name': 'add_neb_emission', 'N': 1,
                        'isfree': False,
                        'init': False,
                        'prior': None})

# Here is a really simple function that takes a **dict argument, picks out the
# `logzsol` key, and returns the value.  This way, we can have gas_logz find
# the value of logzsol and use it, if we uncomment the 'depends_on' line in the
# `gas_logz` parameter definition.
#
# One can use this kind of thing to transform parameters as well (like making
# them linear instead of log, or divide everything by 10, or whatever.) You can
# have one parameter depend on several others (or vice versa).  Just remember
# that a parameter with `depends_on` must always be fixed.

def stellar_logzsol(logzsol=0.0, **extras):
    return logzsol

# FSPS parameter
model_params.append({'name': 'gas_logz', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': r'log Z/Z_\odot',
#                        'depends_on': stellar_logzsol,
                        'prior': priors.TopHat(mini=-2.0, maxi=0.5)})

# FSPS parameter
model_params.append({'name': 'gas_logu', 'N': 1,
                        'isfree': False,
                        'init': -2.0,
                        'units': '',
                        'prior': priors.TopHat(mini=-4, maxi=-1)})

# --- Calibration ---------
# Only important if using a NoiseModel
model_params.append({'name': 'phot_jitter', 'N': 1,
                        'isfree': False,
                        'init': 0.0,
                        'units': 'mags',
                        'prior': priors.TopHat(mini=0.0, maxi=0.2)})


#! /usr/bin/env python

import cvutils, moving, ml, storage, classifier

import numpy as np
import sys, argparse
#from cv2.ml import SVM_RBF, SVM_C_SVC
import cv2
from scipy.stats import norm, lognorm

# TODO add mode detection live, add choice of kernel and svm type (to be saved in future classifier format)

parser = argparse.ArgumentParser(description='The program processes indicators for all pairs of road users in the scene')
parser.add_argument('--cfg', dest = 'configFilename', help = 'name of the configuration file')
parser.add_argument('--cls', dest = 'classifierFilename', help = 'name of the classifier file', required = True)
parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file (overrides the configuration file)')
parser.add_argument('-n', dest = 'nObjects', help = 'number of objects to classify', type = int, default = None)

args = parser.parse_args()

#object is a list
objects = storage.loadTrajectoriesFromSqlite(args.databaseFilename, 'object', args.nObjects, withFeatures = True)
#features = storage.loadTrajectoriesFromSqlite(databaseFilename, 'feature')
speedclassifier = classifier.SpeedClassifier.load(args.classifierFilename)
for obj in objects:
    obj.setUserType(speedclassifier.classify(obj))

#speedclassifier.save("saveSpeed.yaml")
storage.setRoadUserTypes(args.databaseFilename, objects)


import yaml
import numpy as np
from xml.etree.ElementTree import ElementTree, Element, SubElement

from .transform import Transform
from .utils import lift_material, visual, collision, box_link, joint

class LiftDoor:
    def __init__(self, yaml_node, name, lift_size, gap, plugin=True):
        self.name = name
        self.door_type = yaml_node['door_type']
        # x & y coordinates are with respect to the centre of the cabin
        self.x = float(yaml_node['x'])
        self.y = float(yaml_node['y'])
        self.motion_axis_orientation = float(
            yaml_node['motion_axis_orientation'])
        self.width = float(yaml_node['width'])
        self.height = lift_size[2]
        self.thickness = 0.03
        self.gap = gap    # gap between cabin_door and shaft_door
        self.plugin = plugin
        self.params = {'v_max_door': 0.3,
                       'a_max_door': 0.2,
                       'a_nom_door': 0.1,
                       'dx_min_door': 1e-4,
                       'f_max_door': 35.0}

        x_diff = abs(abs(self.x) - lift_size[0] / 2)
        y_diff = abs(abs(self.y) - lift_size[1] / 2)
        offset = 0.025    # cabin door position offset
        if x_diff <= y_diff:
            if self.x > 0:
                self.side = ('right', self.y-self.width/2, self.y+self.width/2)
                self.shaft_door_pose = (lift_size[0]/2 + self.gap, self.y)
                self.cabin_door_pose = (lift_size[0]/2 - offset, self.y)
            else:
                self.side = ('left', self.y-self.width/2, self.y+self.width/2)
                self.shaft_door_pose = (-lift_size[0]/2 - self.gap, self.y)
                self.cabin_door_pose = (-lift_size[0]/2 + offset, self.y)
        else:
            if self.y > 0:
                self.side = ('front', self.x-self.width/2, self.x+self.width/2)
                self.shaft_door_pose = (self.x, lift_size[1]/2 + self.gap)
                self.cabin_door_pose = (self.x, lift_size[1]/2 - offset)
            else:
                self.side = ('back', self.x-self.width/2, self.x+self.width/2)
                self.shaft_door_pose = (self.x, -lift_size[1]/2 - self.gap)
                self.cabin_door_pose = (self.x, -lift_size[1]/2 + offset)

    def generate_cabin_door(self, lift_model_ele, name):
        door_model_ele = SubElement(lift_model_ele, 'model')
        door_model_ele.set('name', name)
        door_pose = SubElement(door_model_ele, 'pose')
        (x, y) = self.cabin_door_pose
        door_pose.text = \
            f'{x} {y} 0 0 0 {self.motion_axis_orientation}'

        self.generate_door_link_and_joint(door_model_ele, parent='platform')

        if self.plugin:
            self.generate_door_plugin(door_model_ele, name)

    def generate_shaft_door(self, world_ele, x, y, z, yaw, name):
        model_ele = SubElement(world_ele, 'model')
        model_ele.set('name', name)
        door_pose = SubElement(model_ele, 'pose')
        # tranformation
        (door_x, door_y) = self.shaft_door_pose
        x_new = x + door_x * np.cos(yaw) - door_y * np.sin(yaw)
        y_new = y + door_x * np.sin(yaw) + door_y * np.cos(yaw)
        yaw_new = yaw + self.motion_axis_orientation
        door_pose.text = f'{x_new} {y_new} {z} 0 0 {yaw_new}'

        self.generate_door_link_and_joint(model_ele)

        floor_thickness = 0.05
        ramp_depth = self.gap * 2
        ramp_size = [self.width, ramp_depth, floor_thickness]
        ramp_pose = Element('pose')
        ramp_pose.text = f'0 0 {-floor_thickness / 2} 0 0 0'
        model_ele.append(box_link('ramp',
                                  ramp_size,
                                  ramp_pose,
                                  material=lift_material(),
                                  bitmask='0x02'))
        model_ele.append(joint('ramp_joint', 'fixed', 'world', 'ramp'))

        if self.plugin:
            self.generate_door_plugin(model_ele, name)

    def generate_door_link_and_joint(self, model_ele, parent='world'):
        door_size = [self.width / 2, self.thickness, self.height]
        right_door_pose = Element('pose')
        right_door_pose.text = f'{self.width / 4} 0 {self.height / 2} 0 0 0'

        model_ele.append(box_link('right_door',
                                  door_size,
                                  right_door_pose,
                                  material=lift_material(),
                                  bitmask='0x02'))

        model_ele.append(joint('right_joint',
                               'prismatic',
                               parent,
                               'right_door',
                               joint_axis='x',
                               lower_limit=0,
                               upper_limit=self.width / 2))

        left_door_pose = Element('pose')
        left_door_pose.text = f'{-self.width / 4} 0 {self.height / 2} 0 0 0'

        model_ele.append(box_link('left_door',
                                  door_size,
                                  left_door_pose,
                                  material=lift_material(),
                                  bitmask='0x02'))

        model_ele.append(joint('left_joint',
                               'prismatic',
                               parent,
                               'left_door',
                               joint_axis='x',
                               lower_limit=-self.width / 2,
                               upper_limit=0))

    def generate_door_plugin(self, model_ele, name):
        plugin_ele = SubElement(model_ele, 'plugin')
        plugin_ele.set('name', 'door')
        plugin_ele.set('filename', 'libdoor.so')
        for param_name, param_value in self.params.items():
            ele = SubElement(plugin_ele, param_name)
            ele.text = f'{param_value}'
        door_ele = SubElement(plugin_ele, 'door')
        door_ele.set('left_joint_name', 'left_joint')
        door_ele.set('name', f'{name}')
        door_ele.set('right_joint_name', 'right_joint')
        door_ele.set('type', 'DoubleSlidingDoor')

    # TODO: remove this function once nesting model is supported in ignition.
    def generate_cabin_door_ign(self, lift_model_ele, name):
        # This is for cabin door generation for ignition gazebo as it doesn't
        # support nested models yet. Once ignition gazebo supports nested
        # models, this should be removed.
        (x, y) = self.cabin_door_pose
        yaw = self.motion_axis_orientation
        right_x = x + np.cos(yaw) * self.width/4
        left_x = x - np.cos(yaw) * self.width/4
        right_y = y + np.sin(yaw) * self.width/4
        left_y = y - np.sin(yaw) * self.width/4

        door_size = [self.width / 2, self.thickness, self.height]
        right_door_pose = Element('pose')
        right_door_pose.text = \
            f'{right_x} {right_y} {self.height / 2} 0 0 {yaw}'

        lift_model_ele.append(box_link(f'{name}_right_door',
                                       door_size,
                                       right_door_pose,
                                       material=lift_material(),
                                       bitmask='0x02'))

        lift_model_ele.append(joint(f'{name}_right_joint',
                                    'prismatic',
                                    'platform',
                                    f'{name}_right_door',
                                    joint_axis='x',
                                    lower_limit=0,
                                    upper_limit=self.width / 2))

        left_door_pose = Element('pose')
        left_door_pose.text = f'{left_x} {left_y} {self.height / 2} 0 0 {yaw}'

        lift_model_ele.append(box_link(f'{name}_left_door',
                                       door_size,
                                       left_door_pose,
                                       material=lift_material(),
                                       bitmask='0x02'))

        lift_model_ele.append(joint(f'{name}_left_joint',
                                    'prismatic',
                                    'platform',
                                    f'{name}_left_door',
                                    joint_axis='x',
                                    lower_limit=-self.width / 2,
                                    upper_limit=0))

        if self.plugin:
            plugin_ele = SubElement(lift_model_ele, 'plugin')
            plugin_ele.set('name', 'door')
            plugin_ele.set('filename', 'libdoor.so')
            for param_name, param_value in self.params.items():
                ele = SubElement(plugin_ele, param_name)
                ele.text = f'{param_value}'
            door_ele = SubElement(plugin_ele, 'door')
            door_ele.set('left_joint_name', f'{name}_left_joint')
            door_ele.set('name', f'{name}')
            door_ele.set('right_joint_name', f'{name}_right_joint')
            door_ele.set('type', 'DoubleSlidingDoor')

class Lift:
    def __init__(self, yaml_node, name, transform, levels):
        self.name = name
        print(f'parsing lift {name}')

        self.depth = float(yaml_node['depth'])
        self.width = float(yaml_node['width'])
        self.yaw = float(yaml_node['yaw'])

        if 'initial_floor_name' in yaml_node:
            self.initial_floor_name = str(yaml_node['initial_floor_name'])
        else:
            self.initial_floor_name = ''

        if 'highest_floor' in yaml_node:
            self.highest_floor = str(yaml_node['highest_floor'])
            if self.highest_floor:
                self.highest_elevation = levels[self.highest_floor].elevation
            else:
                self.highest_elevation = float('inf')
        else:
            self.highest_elevation = float('inf')

        if 'lowest_floor' in yaml_node:
            self.lowest_floor = str(yaml_node['lowest_floor'])
            if self.lowest_floor:
                self.lowest_elevation = levels[self.lowest_floor].elevation
            else:
                self.lowest_elevation = -float('inf')
        else:
            self.lowest_elevation = -float('inf')

        self.plugins = True
        if 'plugins' in yaml_node:
            self.plugins = bool(yaml_node['plugins'])

        raw_pos = (float(yaml_node['x']), -float(yaml_node['y']))
        self.x, self.y = transform.transform_point(raw_pos)
        self.cabin_height = 2.5
        self.wall_thickness = 0.05
        self.floor_thickness = 0.05
        self.gap = 0.05    # gap between lift shaft and lift cabin on each side
        self.shaft_depth = self.depth + 2 * self.gap
        self.shaft_width = self.width + 2 * self.gap

        # default params
        self.cabin_mass = 1200
        self.params = {
            'v_max_cabin': 2.0,
            'a_max_cabin': 1.2,
            'a_nom_cabin': 1.0,
            'dx_min_cabin': 0.001,
            'f_max_cabin': 25323.0}

        self.level_elevation = {}
        self.level_doors = {}
        self.level_names = []
        if 'level_doors' in yaml_node:
            for level_name, door_names in yaml_node['level_doors'].items():
                self.level_doors[level_name] = door_names
                self.level_elevation[level_name] = levels[level_name].elevation
                self.level_names.append(level_name)

        if (not self.initial_floor_name) and self.level_names:
            self.initial_floor_name = self.level_names[0]

        self.doors = []
        if 'doors' in yaml_node:
            self.doors = self.parse_lift_doors(yaml_node['doors'])

        # for wall generation
        # self.end_points stores 1-dimensional positions of endpoints of walls
        # on each side of the cabin in sorted arrays
        self.end_points = {'front': [-self.width/2, self.width/2],
                           'back': [-self.width/2, self.width/2],
                           'left': [-self.depth/2 + self.wall_thickness,
                                    self.depth/2 - self.wall_thickness],
                           'right': [-self.depth/2 + self.wall_thickness,
                                     self.depth/2 - self.wall_thickness]}

        for door in self.doors:
            side, left, right = door.side
            self.end_points[side] += [left, right]
            self.end_points[side].sort()

    def parse_lift_doors(self, yaml_node):
        doors = []
        for lift_door_name, lift_door_yaml in yaml_node.items():
            doors.append(LiftDoor(lift_door_yaml,
                                  lift_door_name,
                                  (self.width, self.depth, self.cabin_height),
                                  self.gap,
                                  self.plugins))
        return doors

    def get_lift_vertices(self):
        # parse lift shaft cavity corner vertices
        vertices = []
        d = self.depth / 2 + self.gap
        w = self.width / 2 + self.gap
        vertices.append((self.x - d * np.sin(self.yaw) - w * np.cos(self.yaw),
                         self.y + d * np.cos(self.yaw) - w * np.sin(self.yaw)))
        vertices.append((self.x + d * np.sin(self.yaw) - w * np.cos(self.yaw),
                         self.y - d * np.cos(self.yaw) - w * np.sin(self.yaw)))
        vertices.append((self.x + d * np.sin(self.yaw) + w * np.cos(self.yaw),
                         self.y - d * np.cos(self.yaw) + w * np.sin(self.yaw)))
        vertices.append((self.x - d * np.sin(self.yaw) + w * np.cos(self.yaw),
                         self.y + d * np.cos(self.yaw) + w * np.sin(self.yaw)))
        return vertices

    def generate_shaft_doors(self, world_ele):
        for level_name, door_names in self.level_doors.items():
            for door in self.doors:
                if door.name in door_names:
                    name = f'ShaftDoor_{self.name}_{level_name}_{door.name}'
                    elevation = self.level_elevation[level_name]
                    door.generate_shaft_door(
                        world_ele, self.x, self.y, elevation, self.yaw, name)

    def generate_wall(self, side, pair, name, platform):
        dims = [pair[1]-pair[0], self.wall_thickness, self.cabin_height]
        mid = (pair[0] + pair[1]) / 2
        if side == 'front':
            x, y, yaw = mid, self.depth/2 - self.wall_thickness / 2, 0
        elif side == 'back':
            x, y, yaw = mid, -self.depth/2 + self.wall_thickness / 2, 0
        elif side == 'left':
            x, y, yaw = -self.width/2 + self.wall_thickness / 2, mid, np.pi/2
        elif side == 'right':
            x, y, yaw = self.width/2 - self.wall_thickness / 2, mid, np.pi/2
        else:
            return

        pose = Element('pose')
        pose.text = f'{x} {y} {self.cabin_height / 2} 0 0 {yaw}'
        platform.append(visual(name, pose, dims, lift_material()))
        platform.append(collision(name, pose, dims, '0x01'))

    def generate_cabin(self, world_ele, options):
        # materials missing for now
        lift_model_name = f'{self.name}'
        lift_model_ele = SubElement(world_ele, 'model')
        lift_model_ele.set('name', lift_model_name)

        # main cabin link for actuation
        platform_name = 'platform'
        platform = SubElement(lift_model_ele, 'link')
        platform.set('name', platform_name)
        inertial = SubElement(platform, 'inertial')
        mass = SubElement(inertial, 'mass')
        mass.text = f'{self.cabin_mass}'
        inertial = SubElement(inertial, 'inertia')
        SubElement(inertial, 'ixx').text = \
            str(self.cabin_mass/12.0*(self.width**2 + self.floor_thickness**2))
        SubElement(inertial, 'iyy').text = \
            str(self.cabin_mass/12.0*(self.depth**2 + self.floor_thickness**2))
        SubElement(inertial, 'izz').text = \
            str(self.cabin_mass/12.0*(self.width**2 + self.depth**2))

        # visuals and collisions for floor and walls of cabin
        floor_dims = [self.width, self.depth, self.floor_thickness]
        floor_name = 'floor'
        floor_pose = Element('pose')
        floor_pose.text = f'0 0 {-self.floor_thickness / 2} 0 0 0'
        platform.append(visual(floor_name,
                               floor_pose,
                               floor_dims,
                               lift_material()))

        platform.append(collision(floor_name, floor_pose, floor_dims, '0x01'))

        # Wall generation
        # get each pair of end_points on each side, generate a section of wall
        # between the pair of points
        for side, end_points in self.end_points.items():
            assert len(end_points) % 2 == 0
            for i in range(0, len(end_points), 2):
                pair = end_points[i: i+2]
                name = f'{side}wall{i//2+1}'
                self.generate_wall(side, pair, name, platform)

        # lift cabin actuation joint
        lift_model_ele.append(joint('cabin_joint',
                                    'prismatic',
                                    'world',
                                    'platform',
                                    joint_axis='z'))

        # cabin doors
        # TODO: remove the if statement here once nesting model is supported
        # in ignition.
        if 'ignition' in options:
            for lift_door in self.doors:
                lift_door.generate_cabin_door_ign(
                    lift_model_ele, f'CabinDoor_{self.name}_{lift_door.name}')
        else:
            for lift_door in self.doors:
                lift_door.generate_cabin_door(
                    lift_model_ele, f'CabinDoor_{self.name}_{lift_door.name}')

        # lift cabin plugin
        if self.plugins:
            plugin_ele = SubElement(lift_model_ele, 'plugin')
            plugin_ele.set('name', 'lift')
            plugin_ele.set('filename', 'liblift.so')

            lift_name_ele = SubElement(plugin_ele, 'lift_name')
            lift_name_ele.text = f'{self.name}'

            for level_name, door_names in self.level_doors.items():
                floor_ele = SubElement(plugin_ele, 'floor')
                floor_ele.set('name', f'{level_name}')
                floor_ele.set(
                    'elevation', f'{self.level_elevation[level_name]}')
                for door in self.doors:
                    if door.name in door_names:
                        door_pair_ele = SubElement(floor_ele, 'door_pair')
                        door_pair_ele.set(
                            'cabin_door',
                            f'CabinDoor_{self.name}_{door.name}')
                        door_pair_ele.set(
                            'shaft_door',
                            f'ShaftDoor_{self.name}_{level_name}_{door.name}')

            initial_floor_ele = SubElement(plugin_ele, 'initial_floor')
            initial_floor_ele.text = f'{self.initial_floor_name}'
            for param_name, param_value in self.params.items():
                ele = SubElement(plugin_ele, param_name)
                ele.text = f'{param_value}'

            cabin_joint_ele = SubElement(plugin_ele, 'cabin_joint_name')
            cabin_joint_ele.text = 'cabin_joint'
        else:
            static_lift_ele = SubElement(lift_model_ele, 'static')
            static_lift_ele.text = 'true'

        # pose
        model_pose = SubElement(lift_model_ele, 'pose')
        model_pose.text = f'{self.x} {self.y} 0 0 0 {self.yaw}'


#!/usr/bin/env python
from __future__ import division
import tensorflow as tf
slim = tf.contrib.slim
import pickle
import numpy as np
import sys
from .preprocessing import image_processing
from .nets import nets_factory
from utils import *
import rospkg

class model(object):
    def __init__(self, dim_feature=[64, 2048], dim_embed=512, dim_hidden=512, n_time_step=19):
        rospack = rospkg.RosPack()
        source_dir = rospack.get_path('captioner') + "/models"
        self.word_to_idx = load_pickle(source_dir+'/train/word_to_idx.pkl')
        self.idx_to_word = {i: w for w, i in self.word_to_idx.items()}
        self.V = len(self.word_to_idx)
        self.K = dim_feature[0]
        self.D = dim_feature[1]
        self.M = dim_embed
        self.H = dim_hidden
        self.T = n_time_step
        self._start = self.word_to_idx['<START>']
        self._null = self.word_to_idx['<NULL>']
        self.CNN_MODEL = 'inception_v3'#resnet_v1_152'
        self.weight_initializer = tf.contrib.layers.xavier_initializer()
        self.const_initializer = tf.constant_initializer(0.0)
        self.emb_initializer = tf.random_uniform_initializer(minval=-1.0, maxval=1.0)

        self.pretrained_model = source_dir+'/model/4000-21'
        self.pretrained_resnet = source_dir+'/resnet/4000-20'
        self.pretrained_classifier = source_dir+'/classifier/4000-21'
        self.pretrained_human_classifier = source_dir+'/human_classifier/42-8'

        self.image_size = 299
        self.trainable_scopes_human_classifier = 'wilder'
        self.load_scopes_model = 'lstm,initial_lstm,word_embedding,image_features,attention_layer,the_second_attention_layer,Visual_Sentinel,decode_layer,attribute_attention_layer,decision_layer'
        self.trainable_scopes_classifier = 'MIL'#'MIL'
        self.checkpoint_exclude_scopes = 'InceptionV3/AuxLogits,InceptionV3/Logits,InceptionV3/Predictions,InceptionV3/PreLogits,' + self.load_scopes_model + ',' + self.trainable_scopes_classifier +',' +self.trainable_scopes_human_classifier
        self.attribute_list = np.array(load_pickle(source_dir+'/train/multi_class_labels_list.pkl'))
        self.num_attributes = 4
        self.image = tf.placeholder(tf.uint8,[None,None,3],'image')

    def init(self,sess):
        #language model saver
        model_variables_to_train = self._get_variables_to_save(self.load_scopes_model)
        saver_model = tf.train.Saver(model_variables_to_train)
        if self.pretrained_model is not None:
            print ("Start with pretrained Model..")
        saver_model.restore(sess, self.pretrained_model)
        #Attribute classifier saver
        classifier_variables_to_train = self._get_variables_to_save(self.trainable_scopes_classifier)
        saver_classifier = tf.train.Saver(classifier_variables_to_train)
        if self.pretrained_classifier is not None:
            print('Start with pretrained classifier')
            saver_classifier.restore(sess,self.pretrained_classifier)
        #resnet saver
        with tf.get_default_graph().as_default():
            print('Start with pretrained cnn')
            variables_to_restore = self.get_init_fn()
        saver_resnet = tf.train.Saver(variables_to_restore,max_to_keep = 1)
        saver_resnet.restore(sess,self.pretrained_resnet)
        #Attribute classifier saver
        human_classifier_variables_to_train = self._get_variables_to_save(self.trainable_scopes_human_classifier)
        saver_human_classifier = tf.train.Saver(human_classifier_variables_to_train)
        if self.pretrained_classifier is not None:
            print('Start with pretrained human classifier')
            saver_human_classifier.restore(sess,self.pretrained_human_classifier)

    def _get_variables_to_save(self,trainable_scopes):
        scopes = [scope.strip() for scope in trainable_scopes.split(',')]
        variables_to_save = []
        for scope in scopes:
            variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)#+tf.get_collection(tf.GraphKeys.MODEL_VARIABLES ,scope)
            variables_to_save.extend(variables)
            for variable in tf.get_collection(tf.GraphKeys.MODEL_VARIABLES ,scope):
                if variable not in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope):
                    variables_to_save.append(variable)
        return variables_to_save

    def get_init_fn(self):
        """Returns a function run by the chief worker to warm-start the training."""
        #with G.as_default():
        exclusions = [scope.strip() for scope in self.checkpoint_exclude_scopes.split(',')]
        variables_to_restore = []
        for var in slim.get_model_variables():
        #for var in tf.trainable_variables(scope ='InceptionV3' ):
            excluded = False
            for exclusion in exclusions:
                if var.op.name.startswith(exclusion):
                    excluded = True
                    break
            if not excluded:
                variables_to_restore.append(var)
        return variables_to_restore

    def _get_initial_lstm(self, features):
        '''
        features: (N,K,H)
        '''
        with tf.variable_scope('initial_lstm'):
            features_mean = tf.reduce_mean(features, 1) #(N,H)
            w_h = tf.get_variable('w_h', [self.H, self.H], initializer=self.weight_initializer)
            b_h = tf.get_variable('b_h', [self.H], initializer=self.const_initializer)
            h = tf.nn.tanh(tf.matmul(features_mean, w_h) + b_h)
            w_c = tf.get_variable('w_c', [self.H, self.H], initializer=self.weight_initializer)
            b_c = tf.get_variable('b_c', [self.H], initializer=self.const_initializer)
            c = tf.nn.tanh(tf.matmul(features_mean, w_c) + b_c)
            return c, h

    def _word_embedding(self, inputs, reuse=False):
        with tf.variable_scope('word_embedding', reuse=reuse):
            w = tf.get_variable('w', [self.V, self.M], initializer=self.emb_initializer)
            x = tf.nn.embedding_lookup(w, inputs, name='word_vector')  # (N, T, M) or (N, M)
            return x

    def _image_features(self, features ,global_feature = True,reuse = False):
        '''
        # features (N,K,H)
        # v_i = ReLu(W_a * a_i),V_g = ReLu(W_b * a_g)
        #I do not want to noisy global image representation instead
        # I would like to use high level concepts to initialize the network
        '''
        with tf.variable_scope('image_features',reuse=reuse):
            w_a = tf.get_variable('w_a', [self.D, self.H], initializer=self.weight_initializer)
            b_a = tf.get_variable('b_a', [self.H], initializer=self.const_initializer)
            features_flat = tf.reshape(features, [-1, self.D])      #(N*K, D)
            features_proj = tf.nn.relu(tf.matmul(features_flat, w_a)+b_a) #(N*K, H)
            features_proj = tf.reshape(features_proj, [-1, self.K, self.H])     #(N,K,H)
            if global_feature == True:
                w_b = tf.get_variable('w_b', [self.D, self.M], initializer=self.weight_initializer)
                b_b = tf.get_variable('b_b', [self.M], initializer=self.const_initializer)
                features_global = tf.reduce_mean(features,axis=1) #(N,D) average pooling
                features_global = tf.nn.relu(tf.matmul(features_global, w_b)+b_b,name = 'global_feature') #(N,M)
                return features_proj,features_global
        return features_proj

    def _attention_layer(self, features_proj, features_orig ,h, vs,reuse=False):
        '''
        # W_v * V + W_g * h_t
        # h (N,H); features_proj (N,K,H); vs (N,H)
        # c (N,H)
        '''
        with tf.variable_scope('attention_layer', reuse=reuse):
            w_v = tf.get_variable('w_v', [self.H, 512], initializer=self.weight_initializer)
            w_g = tf.get_variable('w_g', [self.H, 512], initializer=self.weight_initializer)
            w_s = tf.get_variable('w_s', [self.H, 512], initializer=self.weight_initializer)
            w_h = tf.get_variable('w_h', [512, 1], initializer=self.weight_initializer)

            temp_v = tf.matmul(tf.reshape(features_proj,[-1,self.H]),w_v)   #(N*K,512)
            temp_v = tf.reshape(temp_v,[-1,self.K,512]) #(N,K,512)
            content_v = tf.nn.tanh(temp_v + tf.expand_dims(tf.matmul(h,w_g),1)) #(N,K,512)

            z_t = tf.reshape(tf.matmul(tf.reshape(content_v,[-1,512]),w_h),[-1,self.K]) #(N,K)
            alpha = tf.nn.softmax(z_t) #(N,K)
            c = tf.reduce_sum(features_proj * tf.expand_dims(alpha, 2), 1, name='context') #(N,H)
            context_vector_full = tf.reduce_sum(features_orig * tf.expand_dims(alpha, 2), 1, name='context_full')

            content_s = tf.nn.tanh(tf.matmul(vs,w_s) + tf.matmul(h,w_g)) #(N,K)
            z_t_extended = tf.matmul(content_s,w_h) #(N,1)
            extended = tf.concat([z_t,z_t_extended],1)
            alpha_hat = tf.nn.softmax(extended) #(N,K+1)
            beta = tf.reshape(alpha_hat[:,-1],[-1,1],name = 'beta') #(N,1)
            c_hat = tf.multiply(vs,beta) + tf.multiply(c,(1-beta)) # (N,H)

            return context_vector_full,c_hat, alpha, beta

    def _attention_layer2(self, features_proj, features_orig ,h,alpha_in, vs,reuse=False):
        '''
        # W_v * V + W_g * h_t
        # h (N,H); features_proj (N,K,H); vs (N,H)
        # c (N,H)
        what about inverse feature (??????)
        '''
        with tf.variable_scope('the_second_attention_layer', reuse=reuse):
            w_v = tf.get_variable('w_v', [self.H, 512], initializer=self.weight_initializer)
            w_g = tf.get_variable('w_g', [self.H, 512], initializer=self.weight_initializer)
            w_s = tf.get_variable('w_s', [self.H, 512], initializer=self.weight_initializer)
            w_h = tf.get_variable('w_h', [512, 1], initializer=self.weight_initializer)

            features_proj_alpha = features_proj * tf.expand_dims(alpha_in, 2)
            temp_v = tf.matmul(tf.reshape(features_proj_alpha,[-1,self.H]),w_v)   #(N*K,512)
            temp_v = tf.reshape(temp_v,[-1,self.K,512]) #(N,K,512)
            content_v = tf.nn.tanh(temp_v + tf.expand_dims(tf.matmul(h,w_g),1)) #(N,K,512)

            z_t = tf.reshape(tf.matmul(tf.reshape(content_v,[-1,512]),w_h),[-1,self.K]) #(N,K)
            alpha = tf.nn.softmax(z_t) #(N,K)

            #alpha_mask = tf.to_float(tf.greater_equal(alpha,alpha_in))
            #alpha = tf.nn.softmax(tf.multiply(alpha_mask,z_t))

            c = tf.reduce_sum(features_proj * tf.expand_dims(alpha, 2), 1, name='context') #(N,H)
            context_vector_full = tf.reduce_sum(features_orig * tf.expand_dims(alpha, 2), 1, name='context_full')

            content_s = tf.nn.tanh(tf.matmul(vs,w_s) + tf.matmul(h,w_g)) #(N,K)
            z_t_extended = tf.matmul(content_s,w_h) #(N,1)
            extended = tf.concat([z_t,z_t_extended],1)
            alpha_hat = tf.nn.softmax(extended) #(N,K+1)
            beta = tf.reshape(alpha_hat[:,-1],[-1,1],name = 'beta') #(N,1)
            c_hat = tf.multiply(vs,beta) + tf.multiply(c,(1-beta)) # (N,H)

            return context_vector_full,c_hat, alpha, beta

    def WILDER_attention_map(self,features,reuse = False):
        with tf.variable_scope('wilder',reuse = reuse):
          w_a = tf.get_variable('w_a', [1024, 1], initializer=tf.contrib.layers.xavier_initializer())
          b_a = tf.get_variable('b_a', [1], initializer=tf.constant_initializer(0.0))
          v = tf.get_variable('v', [2048, 1024], initializer=tf.contrib.layers.xavier_initializer())
          u = tf.get_variable('u', [2048, 1024], initializer=tf.contrib.layers.xavier_initializer())
        fc_flattened = tf.reshape(features,[-1,2048]) #(c*324,2048)
        temp_v = tf.nn.tanh(tf.matmul(fc_flattened,v)) # (c*49,1024)
        temp_u = tf.nn.sigmoid(tf.matmul(fc_flattened,u)) # (c*49,1024)
        alpha_logit = tf.reshape(tf.matmul(tf.multiply(temp_v,temp_u),w_a)+b_a,[-1,64,1])
        alpha = tf.nn.softmax(alpha_logit,dim=1) #(c,49,1)
        z = tf.reduce_sum(alpha * tf.reshape(fc_flattened,[-1,64,2048]),1) #(c,2048)
        return z,alpha

    def WILDER_attention(self,z,reuse = False):
        with tf.variable_scope('wilder',reuse=reuse):
          w = tf.get_variable('w', [2048, self.num_attributes], initializer=tf.contrib.layers.xavier_initializer())
          b = tf.get_variable('b', [self.num_attributes], initializer=tf.constant_initializer(0.0))
          w_2 = tf.get_variable('w_2', [2048, 2048], initializer=tf.contrib.layers.xavier_initializer())
          b_2 = tf.get_variable('b_2', [2048], initializer=tf.constant_initializer(0.0))
        temp_logit = tf.nn.relu(tf.matmul(z,w_2)+b_2) #()
        logits = tf.matmul(temp_logit,w)+b #(N,14)
        return logits

    def WILDER_Noisy_or(self,features,attention_logit,is_training,reuse=False):
      with tf.variable_scope('wilder',reuse=reuse):
        fc6 = tf.layers.conv2d(features,filters = 4096,kernel_size = 3,activation = None, padding ='VALID') #(c,5,5,4096)
        fc6 = slim.batch_norm(fc6, is_training = is_training,activation_fn=tf.nn.relu)
        fc7 = tf.layers.conv2d(fc6,filters = 4096,kernel_size = 1,activation = None)#(c,5,5,4096)
        fc7 = slim.batch_norm(fc7, is_training = is_training,activation_fn=tf.nn.relu)

        fc8 = tf.layers.conv2d(fc7,filters = self.num_attributes,kernel_size = 1,activation = tf.nn.sigmoid,
                             bias_initializer=tf.constant_initializer(-6.58),
                             kernel_initializer=tf.truncated_normal_initializer(stddev = 0.001))#(c,5,5,14)
        fc8 = tf.reshape(fc8,[-1,36,self.num_attributes]) #(c,25,14)
        mil_1 = 1 - fc8 #(c,25,14)
        leak_prob = tf.nn.sigmoid(attention_logit)
        max_prob2 = tf.maximum(tf.reduce_max(fc8, axis = 1),leak_prob)
        prob2 = 1 - tf.multiply(tf.reduce_prod(mil_1,[1]),1-leak_prob) #(c,14)
        final_prob2 = tf.maximum(max_prob2,prob2)

        max_prob = tf.reduce_max(fc8, axis = 1) #(c,14)
        prob = 1 - tf.reduce_prod(mil_1,[1])
        final_prob = tf.maximum(max_prob,prob)
      return final_prob,final_prob2

    def _visual_sentinel(self,h,c,x,reuse = False):
        '''
        # h (N,H); c (N,H); x (N,M+H)
        '''
        input_size = self.M+self.M # if use word embedding plus global feauture
        with tf.variable_scope('Visual_Sentinel', reuse=reuse):
            w_x = tf.get_variable('w_x', [input_size, self.H], initializer=self.weight_initializer)
            w_h = tf.get_variable('w_h', [self.H, self.H], initializer=self.weight_initializer)
            gate_t = tf.matmul(x,w_x) + tf.matmul(h,w_h) #(N,H)
            vs = tf.nn.sigmoid(gate_t)*tf.nn.tanh(c) #(N,H)
        return vs

    def _decode_layer(self, h, c_hat,stop_gradient = True,reuse=False):
        '''
        detected_prob attriubtes probability: (Vc,N)
        h: (N,H); c_hat: (N,H)
        '''
        with tf.variable_scope('decode_layer', reuse=reuse):
            w_g = tf.get_variable('w_g',[self.H,self.V],initializer=self.weight_initializer)
            b_g = tf.get_variable('b_g', [self.V], initializer=self.const_initializer)
            logits = tf.matmul(c_hat+h,w_g) + b_g
        return logits

    def _static_attribute_attention_layer(self,features,reuse = False):
      with tf.variable_scope('MIL',reuse = reuse):
        w_a = tf.get_variable('w_a', [1024, 1], initializer=tf.contrib.layers.xavier_initializer())
        b_a = tf.get_variable('b_a', [1], initializer=tf.constant_initializer(0.0))
        v = tf.get_variable('v', [2048, 1024], initializer=tf.contrib.layers.xavier_initializer()) #self.H
        u = tf.get_variable('u', [2048, 1024], initializer=tf.contrib.layers.xavier_initializer()) #self.H

        fc_flattened = tf.reshape(features,[-1,2048]) #(c*324,2048)
        temp_v = tf.nn.tanh(tf.matmul(fc_flattened,v)) # (c*64,1024)
        temp_u = tf.nn.sigmoid(tf.matmul(fc_flattened,u)) # (c*64,1024)
        alpha_logit = tf.reshape(tf.matmul(tf.multiply(temp_v,temp_u),w_a)+b_a,[-1,64,1])
        alpha = tf.nn.softmax(alpha_logit,dim=1) #(c,64,1)
        z = tf.reduce_sum(alpha * tf.reshape(fc_flattened,[-1,64,2048]),1) #(c,512)
        return z,tf.squeeze(alpha,2)

    def _MIL_attention(self,beta,c,z,inital = False,reuse = False):
      with tf.variable_scope('MIL', reuse=reuse):
         w = tf.get_variable('w', [2048, 1113], initializer=tf.contrib.layers.xavier_initializer())
         b = tf.get_variable('b', [1113], initializer=tf.constant_initializer(0.0))
         w_2 = tf.get_variable('w_2', [2048, 2048], initializer=tf.contrib.layers.xavier_initializer())
         b_2 = tf.get_variable('b_2', [2048], initializer=tf.constant_initializer(0.0))
         z_hat = tf.multiply(z,beta) + tf.multiply(c,(1-beta)) # (N,H)
         temp_logit = tf.nn.relu(tf.matmul(z_hat,w_2)+b_2) #()
         logits = tf.matmul(temp_logit,w)+b #(N,1113)
         return logits

    def _MIL_Leaky_noisy_or(self,features,attention_logit,is_training,reuse=False):
        '''
        features: original resenet conv feature
        features :#(N,K,H) -> #(N,7,7,H)
        '''
        with tf.variable_scope('MIL',reuse=reuse):
            fc6 = tf.layers.conv2d(features,filters = 4096,kernel_size = 3,activation = None,padding = 'VALID') #(c,5,5,4096)
            fc6 = slim.batch_norm(fc6, is_training = is_training,activation_fn=tf.nn.relu)

            fc7 = tf.layers.conv2d(fc6,filters = 4096,kernel_size = 1,activation = None)#(c,5,5,4096)
            fc7 = slim.batch_norm(fc7, is_training = is_training,activation_fn=tf.nn.relu)

            fc8 = tf.layers.conv2d(fc7,filters = 1113,kernel_size = 1,activation = tf.nn.sigmoid,
                               bias_initializer=tf.constant_initializer(-6.8),
                               kernel_initializer=tf.truncated_normal_initializer(stddev = 0.001))#(c,5,5,1113)

            fc8 = tf.reshape(fc8,[-1,36,1113]) #(c,25,1113)
            mil_1 = 1 - fc8 #(c,25,1113)
            max_prob = tf.reduce_max(fc8, axis = 1) #(c,1113)
            prob = 1 - tf.reduce_prod(mil_1,[1])
            final_prob = tf.maximum(max_prob,prob)
            leak_prob = tf.nn.sigmoid(attention_logit)
            max_prob2 = tf.maximum(tf.reduce_max(fc8, axis = 1),leak_prob)
            prob2 = 1 - tf.multiply(tf.reduce_prod(mil_1,[1]),1-leak_prob) #(c,1113)
            final_prob2 = tf.maximum(max_prob2,prob2)

        return final_prob,final_prob2

    def _attribute_gate_layer(self,beta,c,z,h,x,prob,reuse=False):
        with tf.variable_scope('attribute_attention_layer',reuse=reuse):
            w_z = tf.get_variable('w_z', [self.M,2048 ], initializer=tf.contrib.layers.xavier_initializer())
            b_z = tf.get_variable('b_z', [2048], initializer=tf.constant_initializer(0.0))
            w_z3 = tf.get_variable('w_z3', [2048, 1113], initializer=tf.contrib.layers.xavier_initializer())
            b_z3 = tf.get_variable('b_z3', [1113], initializer=tf.constant_initializer(0.0))
            w_h = tf.get_variable('w_h', [self.H, 2048], initializer=tf.contrib.layers.xavier_initializer())
            logit_z = tf.nn.relu(tf.matmul(x,w_z) + tf.matmul(h,w_h)+b_z)
            logit_z = tf.matmul(logit_z,w_z3)+b_z3#(N,H)
            attention = tf.multiply(tf.nn.sigmoid(logit_z),prob)
        return attention

    def _attribute_vector_op(self,batch_size,attention_logits,attribute_embedding_expand):
        values_top10,_ = tf.nn.top_k(attention_logits,10) #(N,10)
        values_10 = tf.gather(values_top10,9,axis = 1) #(N,1)
        attention_mask = tf.to_float(tf.greater_equal(attention_logits,tf.expand_dims(values_10,1)))
        attribute_attention = attention_logits*attention_mask/tf.reduce_sum(attention_logits*attention_mask,axis = 1,keep_dims = True)#(N,1113)
        attribute_vector = tf.multiply(tf.expand_dims(attribute_attention,2),tf.tile(attribute_embedding_expand,[batch_size,1,1])) #(N,1113,M)
        attribute_vector = tf.reduce_sum(attribute_vector,1) #(N,M)
        return attribute_vector,attribute_attention

    def decision(self,c,h,reuse = False):
        with tf.variable_scope('decision_layer',reuse = reuse):
            w_c = tf.get_variable('w_c', [self.M, 1024], initializer=self.weight_initializer)
            b_z = tf.get_variable('b_z',[1024],initializer=tf.constant_initializer(0.0))
            w = tf.get_variable('w', [1024,1], initializer=self.weight_initializer)
            b = tf.get_variable('b',[1],initializer=tf.constant_initializer(0.0))
            logit = tf.nn.relu(tf.matmul(c+h,w_c)+b_z)
            logit = tf.matmul(logit,w)+b
        return logit

    def init_human_attributes(self,features):
        human_attributes_im,_ = self.WILDER_attention_map(features)
        human_attributes_logit = self.WILDER_attention(human_attributes_im,reuse=False)
        _,human_attributes_prob = self.WILDER_Noisy_or(features,human_attributes_logit,is_training=False,reuse=False) #(N,14)
        short_prob = tf.expand_dims(tf.gather(human_attributes_prob,indices=0,axis = 1),1)
        long_prob = tf.expand_dims(tf.gather(human_attributes_prob,indices=1,axis = 1),1)
        plaid_prob = tf.expand_dims(tf.gather(human_attributes_prob,indices=2,axis = 1),1)
        striped_prob = tf.expand_dims(tf.gather(human_attributes_prob,indices=3,axis = 1),1)
        human_attributes_prob_extended = tf.concat([short_prob,long_prob,tf.maximum(short_prob,long_prob),
                                                plaid_prob,striped_prob,tf.maximum(plaid_prob,striped_prob)],axis=1)
        self.attribute_threshold_greater = [0.5,0.5,0.5,0.5,0.5,0.5]
        self.attribute_threshold_smaller = [0.0,0.0,0.0,0.0,0.0,0.0]
        self.greater_attributes_list = [373,148,1701,1124,810,278] #short long hair plaid striped shirt
        self.smaller_attributes_list =  [0,0,0,0,0,0]
        return human_attributes_prob_extended

    def human_attributes_op(self,x,human_attributes_prob,batch_size):
        mask_greater =tf.to_int32(tf.greater_equal(human_attributes_prob,tf.tile([self.attribute_threshold_greater],[batch_size,1]))) #(N,14)
        mask_smaller = tf.to_int32(tf.greater_equal(tf.tile([self.attribute_threshold_smaller],[batch_size,1]),human_attributes_prob)) #(N,14)
        greater = self.greater_attributes_list * mask_greater
        smaller = self.smaller_attributes_list * mask_smaller
        human_attributes_list = greater + smaller
        used_mask = tf.to_int32(tf.equal(tf.expand_dims(x,1),human_attributes_list))
        for used_word in self.used_list:
            used_mask = used_mask+tf.to_int32(tf.equal(tf.expand_dims(used_word,1),human_attributes_list))
        used_mask = tf.to_int32(tf.equal(used_mask,0))
        self.used_list.append(x)
        #import ipdb; ipdb.set_trace()
        human_attributes_list = human_attributes_list*used_mask
        empty_flag = tf.to_float(tf.not_equal(tf.reduce_sum(human_attributes_list,1,keep_dims=True),0))

        human_attributes_embedding = self._word_embedding(human_attributes_list,reuse = True)#(N,6,M)
        human_attributes_embedding = tf.multiply(human_attributes_embedding,tf.cast(tf.expand_dims(
                                          used_mask*(mask_smaller+mask_greater),2),tf.float32)) #(N,14,M)
        human_attributes_mean = tf.reduce_sum(human_attributes_embedding,1)/tf.clip_by_value(tf.reduce_sum(
                                                            tf.cast(used_mask*(mask_greater+mask_smaller),tf.float32),keep_dims=True,axis=1),1.0,15.0) #(N,M)
        return human_attributes_mean,human_attributes_list,empty_flag

    def build_sampler(self,max_len=20):
        image = tf.image.convert_image_dtype(self.image, dtype=tf.float32)
        image_val = image_processing.process_image(image,is_training=False, height = self.image_size, width =self.image_size )
        network_fn = nets_factory.get_network_fn(
                      self.CNN_MODEL,
                      num_classes=1001,
                      weight_decay=0.0001,
                      is_training=False)
        logits,end_points = network_fn(tf.expand_dims(image_val,0))
        features = end_points['Mixed_7c']

        batch_size = tf.shape(features)[0]
        # batch normalize feature vectors
        with tf.variable_scope('image_features'):
            features_orig = slim.batch_norm(features, is_training = False)

        features_orig = tf.reshape(features_orig,[-1,self.K,2048])
        features_proj,features_global = self._image_features(features_orig)
        c, h = self._get_initial_lstm(features=features_proj)

        z,alpha = self._static_attribute_attention_layer(features)

        attribute_embedding = tf.expand_dims(self._word_embedding(self.attribute_list,reuse = False),0) #(1,1113,M)
        MIL_attention = self._MIL_attention(1.0,0.0,z,inital = True,reuse = False)
        _,prob2 = self._MIL_Leaky_noisy_or(features,MIL_attention,is_training=False,reuse=False)
        attribute_vector,attribute_attention = self._attribute_vector_op(batch_size,prob2,attribute_embedding)

        sampled_word_list = []
        lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(num_units=self.H)

        with tf.variable_scope('lstm', reuse=False):
            _, (c, h) = lstm_cell(inputs=tf.concat([features_global,attribute_vector],axis=1), state=[c, h])

        flag1=1.0
        human_attributes_prob_extended = self.init_human_attributes(features)
        self.used_list = []
        for t in range(max_len):
            if t == 0:
                sampled_word = tf.fill([tf.shape(features)[0]], self._start)
            x = self._word_embedding(inputs=sampled_word, reuse=True)
            x_new = tf.concat([x,attribute_vector],1)

            with tf.variable_scope('lstm', reuse=True):
                h_old = h
                c_old = c
                _, (c, h) = lstm_cell(inputs=x_new, state=[c, h]) #(N,512)
            vs = self._visual_sentinel(h_old,c,x_new,reuse=(t!=0))

            c_context,c_hat_1, alpha,beta_1 = self._attention_layer(features_proj,features_orig ,h, vs,reuse=(t!=0))
            gamma = self.decision(c_hat_1,h,reuse = (t!=0)) #(N,1)
            trigger = tf.to_float(tf.greater_equal(tf.nn.sigmoid(gamma),0.3))

            MIL_attention = self._MIL_attention(0.0,c_context,z,inital = False,reuse=True)
            _,prob2 = self._MIL_Leaky_noisy_or(features,MIL_attention,is_training=False,reuse=True)
            attention_logits= self._attribute_gate_layer(beta_1,c_context,z,h,x,prob2,reuse=(t!=0))
            attribute_vector2,attribute_attention2 = self._attribute_vector_op(batch_size,attention_logits,attribute_embedding)

            human_attributes_mean,human_attributes_list,empty_flag =self.human_attributes_op(sampled_word,human_attributes_prob_extended,batch_size)
            hyperparameter = 0.9
            attribute_vector2 = (1-trigger*flag1*hyperparameter*empty_flag)*attribute_vector2 + trigger*flag1*hyperparameter*empty_flag*human_attributes_mean
            attribute_vector2 = (1-(1-trigger)*(1-flag1)*hyperparameter*empty_flag)*attribute_vector2 + (1-trigger)*(1-flag1)*hyperparameter*empty_flag*human_attributes_mean
            flag1 = 1-trigger

            x_new = tf.concat([x,attribute_vector2],1)

            with tf.variable_scope('lstm', reuse=True):
              _, (c, h) = lstm_cell(inputs=x_new, state=[c_old, h_old])
            vs = self._visual_sentinel(h_old,c,x_new,reuse=True)
            _,c_hat, alpha,beta = self._attention_layer2(features_proj, features_orig,h,alpha, vs,reuse=(t!=0))
            logits = self._decode_layer(h, c_hat ,stop_gradient = False,reuse=(t!=0))

            sampled_word = tf.argmax(logits, 1,output_type=tf.int32)
            sampled_word_list.append(sampled_word)
        sampled_captions = tf.transpose(tf.stack(sampled_word_list), (1, 0))     # (N, max_len)
        return sampled_captions


from abc import abstractmethod
from collections import OrderedDict
from functools import partial
from typing import List

import numpy as np
import pymc3 as pm
import theano as th
import theano.tensor as tt
from pymc3.variational.updates import get_or_compute_grads

from .. import types
from ..io import io_commons
from ..models.fancy_model import GeneralizedContinuousModel

class FancyStochasticOptimizer:
    """The base class of stochastic optimizers equipped with the functionality of saving and loading the
    optimizer state to and from disk (e.g. for stateful optimizers such as ADAM and ADAMAX), and the
    possibility of utilizing the extra attributes of `GeneralizedContinuousModel` to perform structured
    parameter updates, e.g. updating only sample-specific variables while keeping global variables intact
    (see `FancyAdamax` for a concrete implementation).
    """
    @abstractmethod
    def get_optimizer(self,
                      model: GeneralizedContinuousModel=None,
                      approx: pm.MeanField=None):
        """

        Args:
            model: a generalized continuous PyMC3 model
            approx: an instance of PyMC3 mean-field approximation

        Returns:
            A callable function that upon providing `loss_or_grads` and `params`, returns an
            `OrderedDict` of shared theano tensor updates (for example, see `FancyAdamax.get_optimizer`).
        """
        raise NotImplementedError

    @staticmethod
    def get_call_kwargs(_locals_):
        _locals_ = _locals_.copy()
        _locals_.pop('loss_or_grads')
        _locals_.pop('params')
        return _locals_

    @abstractmethod
    def save(self, output_path: str):
        raise NotImplementedError

    @abstractmethod
    def load(self, input_path: str):
        raise NotImplementedError

class FancyAdamax(FancyStochasticOptimizer):
    """Adamax optimizer with saving/loading functionality and sample-specific-only update mode."""
    def __init__(self,
                 learning_rate: float = 0.002,
                 beta1: float = 0.9,
                 beta2: float = 0.999,
                 epsilon: float = 1e-8,
                 sample_specific_only: bool = False,
                 disable_bias_correction: bool = False):
        """Initializer.

        Args:
            learning_rate: learning rate
            beta1: first moment forgetting factor
            beta2: second moment forgetting factor
            epsilon: a small float for avoiding division-by-zero
            sample_specific_only: only update sample-specific variables (as specified in the generalized model)
            disable_bias_correction: disable moment estimation bias correction
        """
        self.learning_rate = learning_rate
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.sample_specific_only = sample_specific_only
        self.disable_bias_correction = disable_bias_correction

        # placeholder for first (m) and second (u) moments
        # in mean-field type approximation, ``mu`` and ``rho`` each have their own tensors
        # the list elements correspond to ``mu`` and ``rho`` moments, respectively
        self.m_tensors: List[types.TensorSharedVariable] = []
        self.u_tensors: List[types.TensorSharedVariable] = []

        # placeholder for the state of moment estimation bias corrector
        self.res_tensor: types.TensorSharedVariable = None

    def _assert_shared_tensors_available(self):
        m_u_available = len(self.m_tensors) == 2 and len(self.u_tensors) == 2
        res_available = self.disable_bias_correction or self.res_tensor is not None
        assert m_u_available and res_available, "Adamax tensors are not available yet"

    def get_mu_m(self):
        self._assert_shared_tensors_available()
        return self.m_tensors[0]

    def get_rho_m(self):
        self._assert_shared_tensors_available()
        return self.m_tensors[1]

    def get_mu_u(self):
        self._assert_shared_tensors_available()
        return self.u_tensors[0]

    def get_rho_u(self):
        self._assert_shared_tensors_available()
        return self.u_tensors[1]

    def get_res_tensor(self):
        return self.res_tensor

    @staticmethod
    def structured_adamax(loss_or_grads=None,
                          params=None,
                          model: GeneralizedContinuousModel=None,
                          approx: pm.MeanField=None,
                          learning_rate=0.002, beta1=0.9,
                          beta2=0.999, epsilon=1e-8,
                          sample_specific_only=False,
                          disable_bias_correction=False,
                          base_class: 'FancyAdamax' = None):
        """Adamax stochastic optimizer with partial sample-specific-only update functionality.

        Args:
            loss_or_grads: symbolic loss function or gradients
            params: variational parameter bundle
            model: an instance of generalized model
            approx: an instance of variational approximation for the model
            learning_rate: global learning rate
            beta1: first moment estimation forgetting factor
            beta2: second moment estimation forgetting factor
            epsilon: a small float to avoid division-by-zero
            sample_specific_only: only update parameters registered in the generalized model as sample-specific
            disable_bias_correction: disable moment estimation bias correction
            base_class: a reference to the base class to store a reference to the shared tensors (for I/O)

        Returns:
            returns the function itself if `loss_or_grads` and `params` are not given;
            otherwise, returns an ordered dict of shared tensor updates (to be used in pymc3 for compiling
            the step function)
        """
        if loss_or_grads is None and params is None:
            return partial(FancyAdamax.structured_adamax,
                           **FancyStochasticOptimizer.get_call_kwargs(locals()))
        elif loss_or_grads is None or params is None:
            raise ValueError('Please provide both `loss_or_grads` and `params` to get updates')
        assert model is not None, 'Please provide `model` to get updates'
        assert approx is not None, 'Please provide `approx` to get updates'

        all_grads = get_or_compute_grads(loss_or_grads, params)
        updates = OrderedDict()

        # indices of sample-specific vars
        if sample_specific_only:
            vmap_list = io_commons.get_var_map_list_from_meanfield_approx(approx)
            sample_specific_indices = []
            for vmap in vmap_list:
                if vmap.var in model.sample_specific_var_registry:
                    sample_specific_indices += [idx for idx in range(vmap.slc.start, vmap.slc.stop)]
            update_indices = th.shared(np.asarray(sample_specific_indices, dtype=np.int))
            num_dof = len(sample_specific_indices)

        # Using theano constant to prevent upcasting of float32
        one = tt.constant(1)

        if disable_bias_correction:
            a_t = learning_rate
        else:
            res_prev = th.shared(pm.theanof.floatX(beta1))
            res = beta1 * res_prev
            a_t = learning_rate / (one - res)
            updates[res_prev] = res
            if base_class is not None:
                base_class.res_tensor = res_prev

        for param, g_t in zip(params, all_grads):
            if sample_specific_only:
                g_t_view = g_t[update_indices]
                m_prev = th.shared(np.zeros((num_dof,), dtype=types.floatX),
                                   broadcastable=(False,))
                u_prev = th.shared(np.zeros((num_dof,), dtype=types.floatX),
                                   broadcastable=(False,))
            else:
                g_t_view = g_t
                value = param.get_value(borrow=True)
                m_prev = th.shared(np.zeros(value.shape, dtype=types.floatX),
                                   broadcastable=(False,))
                u_prev = th.shared(np.zeros(value.shape, dtype=types.floatX),
                                   broadcastable=(False,))

            # save a reference to m and u in the base class
            if base_class is not None:
                base_class.m_tensors.append(m_prev)
                base_class.u_tensors.append(u_prev)

            m_t = beta1 * m_prev + (one - beta1) * g_t_view
            u_t = tt.maximum(beta2 * u_prev, abs(g_t_view))
            step = a_t * m_t / (u_t + epsilon)

            if sample_specific_only:
                new_param = tt.inc_subtensor(param[update_indices], -step)
            else:
                new_param = param - step

            updates[m_prev] = m_t
            updates[u_prev] = u_t
            updates[param] = new_param

        return updates

    def get_optimizer(self,
                      model: GeneralizedContinuousModel=None,
                      approx: pm.MeanField=None):
        return FancyAdamax.structured_adamax(
            model=model,
            approx=approx,
            beta1=self.beta1,
            beta2=self.beta2,
            learning_rate=self.learning_rate,
            epsilon=self.epsilon,
            sample_specific_only=self.sample_specific_only,
            disable_bias_correction=self.disable_bias_correction,
            base_class=self)

    def save(self, output_path: str) -> None:
        """Saves the state of the optimizer to disk.

        Args:
            output_path: output path (must be writable directory)
        """
        from ..io import io_adamax  # lazy import to break import cycle
        io_adamax.AdamaxStateExporter(self, output_path)()

    def load(self, input_path: str):
        """Loads the state of the optimizer from disk.

        Args:
            input_path: input path (must be a readable directory)
        """
        from ..io import io_adamax  # lazy import to break import cycle
        io_adamax.AdamaxStateImporter(self, input_path)()

import numpy,import theano,import pymc3,from pymc3


import os
import pickle
import random
import re
import shutil
from typing import Dict, List, Tuple

import numpy as np
import torch

from torch.nn.utils.rnn import pad_sequence

from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler

from transformers import (

    WEIGHTS_NAME,
    AdamW,

    GPT2Config,
    GPT2LMHeadModel,
    GPT2Tokenizer,

    PreTrainedModel,
    PreTrainedTokenizer,

    get_linear_schedule_with_warmup,

)

''' config parameters '''

# Optional input sequence length after tokenization.
# The training dataset will be truncated in block of this size for training
# Default to the model max input length for single sentence inputs (take into account special tokens)
block_size: int = -1

# Number of updates steps to accumulate before performing a backward/update pass
gradient_accumulation_steps: int = 1

# The initial learning rate for Adam
learning_rate: float = 5e-5

# Weight decay if we apply some
weight_decay: float = 0.

# Epsilon for Adam optimizer
adam_epsilon: float = 1e-8

# Max gradient norm
max_grad_norm: float = 1.

# Total number of training epochs to perform
num_train_epochs: float = 1.

# If > 0: set total number of training steps to perform. Override num_train_epochs
max_steps: int = -1

# Linear warmup over warmup_steps
warmup_steps: int = 0

''' init train env '''

device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = 1

# Set seed
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if args.n_gpu > 0:
    torch.cuda.manual_seed_all(seed)

''' Load pretrained model and tokenizer '''

# load pretrained model and tokenizer
config_class, model_class, tokenizer_class = GPT2Config, GPT2LMHeadModel, GPT2Tokenizer

# init config
config = config_class()

# Our input block size will be the max possible for the model
if block_size <= 0:
    block_size = tokenizer.max_len
else:
    block_size = min(block_size, tokenizer.max_len)

# Training new model from scratch
model = model_class(config=config)

# push model to device
model.to(device)

''' Training '''

# get training dataset
file_path = '../data/lyrics/'
train_dataset = LineByLineTextDataset(tokenizer, file_path = file_path, block_size = block_size)
#train_dataset = TextDataset(tokenizer, args, file_path=file_path, block_size = block_size)

# step training
global_step, tr_loss = train(args, train_dataset, model, tokenizer)

""" Train the model """

train_batch_size = 4

def collate(examples: List[torch.Tensor]):

    if tokenizer._pad_token is None:

        return pad_sequence(examples, batch_first = True)

    return pad_sequence(examples, batch_first = True, padding_value = tokenizer.pad_token_id)

# init random sampler on dataset
train_sampler = RandomSampler(train_dataset)

# init dataloader
train_dataloader = DataLoader(
    train_dataset, sampler = train_sampler, batch_size = train_batch_size, collate_fn = collate)

if max_steps > 0:

    t_total = max_steps
    num_train_epochs = max_steps // (len(train_dataloader) // gradient_accumulation_steps) + 1

else:
    t_total = len(train_dataloader) // gradient_accumulation_steps * num_train_epochs

# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
    {
        "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
        "weight_decay": weight_decay,
    },
    {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
]

# init optimiser on model parameters
optimizer = AdamW(optimizer_grouped_parameters, lr = learning_rate, eps = adam_epsilon)

# init learning rate scheduler
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps = warmup_steps, num_training_steps = t_total)

''' perform training '''

global_step = 0
epochs_trained = 0
steps_trained_in_current_epoch = 0

tr_loss, logging_loss = 0.0, 0.0

model.resize_token_embeddings(len(tokenizer))

model.zero_grad()

for _ in epochs_trained:

    for step, batch in enumerate(train_dataloader):

        # Skip past any already trained steps if resuming training
        if steps_trained_in_current_epoch > 0:
            steps_trained_in_current_epoch -= 1
            continue

        # unpack batch data
        inputs, labels = (batch, batch)

        # push data to device
        inputs = inputs.to(args.device)
        labels = labels.to(args.device)

        # set model to train
        model.train()

        # perform forward pass through model
        outputs = model(inputs, labels=labels)

        # obtain loss; model outputs are always tuple in transformers
        loss = outputs[0]

        # gradient accumulation
        if gradient_accumulation_steps > 1:
            loss = loss / gradient_accumulation_steps

        # backprop loss
        loss.backward()

        # store loss
        tr_loss += loss.item()

        if (step + 1) % gradient_accumulation_steps == 0:

            # perform gradient normalisation
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

            # step optimiser
            optimizer.step()

            # Update learning rate schedule
            scheduler.step()

            model.zero_grad()
            global_step += 1

        if max_steps > 0 and global_step > max_steps:
            break
    if max_steps > 0 and global_step > max_steps:
        break

print(global_step, tr_loss / global_step)

''' saving '''

#model.save_pretrained(output_dir)
#tokenizer.save_pretrained(output_dir)

# Load a trained model and vocabulary that you have fine-tuned
#model = model_class.from_pretrained(args.output_dir)
#tokenizer = tokenizer_class.from_pretrained(args.output_dir)
#model.to(device)


#! /usr/bin/env python3

import numpy as np
import argparse
import json
import os
from scipy.interpolate import interp1d

parser = argparse.ArgumentParser(description="Compute log-likelihoods from light curve evaluations")
parser.add_argument("--interp-directory", help="Directory containing the interpolator evaluations and index files")
parser.add_argument("--output-file", help="Name of file to save posterior samples to")
parser.add_argument("--grid-file", help="Location of the grid file")
parser.add_argument("--lc-file", help="Location of the light curve JSON file")
parser.add_argument("--bands", nargs="+", help="Bands to use")
parser.add_argument("--distance", type=float, help="Luminosity distance (in Mpc)")
args = parser.parse_args()

# Load the files containing the grid indices corresponding to each angle
indices_fname_base = args.interp_directory + ("/" if args.interp_directory[-1] != "/" else "") + "indices_"
interp_angles = [0., 30., 45., 60., 75., 90.]
indices = {_theta_interp:np.loadtxt(indices_fname_base + str(int(_theta_interp)) + ".dat", ndmin=1).astype(int) for _theta_interp in interp_angles}

# Load the grid
grid = np.loadtxt(args.grid_file)
theta_grid = grid[:,-1] # angle is the fifth column

# Load the light curve data
with open(args.lc_file, "r") as fp:
    lc_data = json.load(fp)
for b in lc_data.keys():
    for key in lc_data[b].keys():
        lc_data[b][key] = np.array(lc_data[b][key]) # cast everything as a numpy array

interpolator_times = np.logspace(np.log10(0.125), np.log10(37.239195485411194), 264)
interpolator_times_str_to_float = {"{:.3f}".format(_t):_t for _t in interpolator_times} # this is useful since the times in the filenames are truncated to three decimal places

# Set up a data structure to hold all the interpolated magnitudes.
# This is a list with the following structure:
#   - Each element in the list corresponds to a row in the grid.
#   - This list element is a dictionary mapping (band, interpolator_angle) to *another* dictionary.
#       - This sub-dictionary maps the strings "time", "mag", and "mag_err" to lists containing those values
# The point is to get all the interpolated magnitudes in one place, since they currently exist in many different files,
# all while keeping track of which grid point each one corresponds to, its time, and its band, and to minimize lookup times going forward.
interp_data = [{} for _ in range(theta_grid.size)]

# Fill this data structure with data loaded from the interpolated magnitude files
fname_base = args.interp_directory + ("/" if args.interp_directory[-1] != "/" else "")
for fname in os.listdir(args.interp_directory):

    fname_split = fname.split("_")

    # Determine whether it's a file we want or not
    if fname_split[0] != "eval":
        continue

    fname_split = fname_split[2:]

    # Get the time, angle, and band from the filename
    t_interp = interpolator_times_str_to_float[fname_split[0]] # convert from the truncated time in the filename to the full time value
    theta_interp = float(fname_split[1])
    band = fname_split[2][0]

    if band not in args.bands:
        continue

    if indices[theta_interp].size == 0:
        continue

    # Load the interpolated magnitude data
    data = np.loadtxt(fname_base + fname)

    # "i" is the index within this data file
    # "grid_index" is the row of the grid corresponding to this particular interpolator evaluation
    for i, grid_index in enumerate(indices[theta_interp]):
        # initialize things if needed
        if (band, theta_interp) not in interp_data[grid_index].keys():
            interp_data[grid_index][(band, theta_interp)] = {"time":[], "mag":[], "mag_err":[]}
        interp_data[grid_index][(band, theta_interp)]["time"].append(t_interp)
        interp_data[grid_index][(band, theta_interp)]["mag"].append(data[i][0])
        interp_data[grid_index][(band, theta_interp)]["mag_err"].append(data[i][1])

# Go through and sort every array in interp_data by time (and also convert them to numpy arrays)
for i in range(theta_grid.size):
    for key in interp_data[i].keys():
        interp_data[i][key]["time"] = np.array(interp_data[i][key]["time"])
        interp_data[i][key]["mag"] = np.array(interp_data[i][key]["mag"])
        interp_data[i][key]["mag_err"] = np.array(interp_data[i][key]["mag_err"])
        sorted_indices = np.argsort(interp_data[i][key]["time"])
        interp_data[i][key]["time"] = interp_data[i][key]["time"][sorted_indices]
        interp_data[i][key]["mag"] = interp_data[i][key]["mag"][sorted_indices]
        interp_data[i][key]["mag_err"] = interp_data[i][key]["mag_err"][sorted_indices]

# Now build a list of interpolators corresponding to each row in the grid for each band, where the interpolation over angle has been done
lc_functions = []
for i, theta in enumerate(theta_grid):
    # Determine the angular bin
    for j in range(1, len(interp_angles)):
        if interp_angles[j] > theta:
            break
    theta_lower = interp_angles[j - 1]
    theta_upper = interp_angles[j]

    f_dict = {}
    for b in args.bands:
        # Get the interpolated values for this row in the grid
        t = interp_data[i][(b, theta_lower)]["time"]
        mag_lower = interp_data[i][(b, theta_lower)]["mag"]
        mag_err_lower = interp_data[i][(b, theta_lower)]["mag_err"]
        mag_upper = interp_data[i][(b, theta_upper)]["mag"]
        mag_err_upper = interp_data[i][(b, theta_upper)]["mag_err"]

        # Interpolate in angle
        mag = ((theta_upper - theta) / (theta_upper - theta_lower)) * mag_lower + ((theta - theta_lower) / (theta_upper - theta_lower)) * mag_upper
        mag_err = ((theta_upper - theta) / (theta_upper - theta_lower)) * mag_err_lower + ((theta - theta_lower) / (theta_upper - theta_lower)) * mag_err_upper

        # Interpolate in time and put the interpolators in the dictionary
        f_dict[b] = (interp1d(t, mag, fill_value="extrapolate"), interp1d(t, mag_err, fill_value="extrapolate"))

    lc_functions.append(f_dict)

# Now, finally, compute the likelihoods
lnL = np.zeros(theta_grid.size)
for i, f_dict in enumerate(lc_functions):
    for b in args.bands:
        mag_interpolator, mag_err_interpolator = f_dict[b]
        # make sure to account for the distance modulus when computing residuals
        residuals = mag_interpolator(lc_data[b]["time"]) + (5. * (np.log10(args.distance * 1.e6) - 1.)) - lc_data[b]["mag"]
        model_error = mag_err_interpolator(lc_data[b]["time"])
        lnL[i] += -0.5 * np.sum(residuals**2 / (model_error**2 + lc_data[b]["mag_err"]**2)
            + np.log(2. * np.pi * (model_error**2 + lc_data[b]["mag_err"]**2)))

grid[:,0] = lnL

# Get the header for the posterior sample file from the grid file, then save
with open(args.grid_file, "r") as fp:
    header = fp.readline().strip()[2:] # strip off the "# " at the beginning and the "\n" at the end
np.savetxt(args.output_file, grid, header=header)