import pdb
import os
import pandas as pd
import dymos as dm
import numpy as np
import datetime as dt
from typing import Union
import openmdao.api as om
from pyNA.src.settings import Settings
from pyNA.src.aircraft import Aircraft
from pyNA.src.trajectory_src.atmosphere import Atmosphere
from pyNA.src.engine import Engine
from scipy.interpolate import RegularGridInterpolator
from pyNA.src.trajectory_src.trajectory_ode import TrajectoryODE
from pyNA.src.trajectory_src.mux import Mux
from pyNA.src.trajectory_src.surrogate_noise import SurrogateNoise

class Trajectory:
    """
    The trajectory module contains the methods to compute the take-off trajectory used by pyNA.
    """

    def __init__(self, n_order):

        # Initialize path
        self.path = pd.DataFrame
        self.n_t = np.int

        # Initialize phases
        # self.phase_name_lst = ['groundroll', 'rotation', 'liftoff', 'vnrs', 'cutback']
        self.phase_name_lst = ['groundroll', 'rotation', 'liftoff', 'vnrs', 'cutback']
        self.phases = dict()

        # Compute transcription for the phases
        self.num_segments = []
        self.transcription_order = []
        self.transcription_phases = []
        self.phase_size = []
        for i, phase in enumerate(self.phase_name_lst):
            if phase == 'groundroll':
                self.num_segments.append(3)
                self.transcription_order.append(3)
            elif phase == 'rotation':
                self.num_segments.append(3)
                self.transcription_order.append(3)
            elif phase == 'liftoff':
                self.num_segments.append(4)
                self.transcription_order.append(3)
            elif phase == 'vnrs':
                self.num_segments.append(7)
                self.transcription_order.append(3)
            elif phase == 'cutback':
                self.num_segments.append(12)
                self.transcription_order.append(3)

            self.transcription_phases.append(dm.GaussLobatto(num_segments=self.num_segments[i], order=self.transcription_order[i], compressed=True, solve_segments=False))
            self.transcription_phases[i].init_grid()
            self.phase_size.append(self.num_segments[i] * self.transcription_order[i]+1)

        # Compute size of the muxed trajectory
        self.trajectory_size = Trajectory.compute_size_output_mux(size_inputs=self.phase_size)

        return

    @staticmethod
    def get_engine_variables(settings: Settings) -> Union[list, list]:
        """
        Get the engine parameters to compute during the trajectory computations.

        :param settings: pyna settings
        :type settings: Settings

        :return: (engine_var, engine_var_units)
        :rtype: (list, list)
        """
        # Engine variables
        engine_var = ['W_f', 'Tti_c', 'Pti_c']
        engine_var_units = ['kg/s', 'K', 'Pa']

        # Jet parameters
        engine_var.extend(['V_j', 'rho_j', 'A_j', 'Tt_j', 'M_j'])
        engine_var_units.extend(['m/s', 'kg/m**3', 'm**2', 'K', None])

        # Core parameters
        if settings.method_core_turb == 'GE':
            engine_var.extend(['mdoti_c', 'Ttj_c', 'DTt_des_c'])
            engine_var_units.extend(['kg/s', 'K', 'K'])
        elif settings.method_core_turb == 'PW':
            engine_var.extend(['mdoti_c', 'Ttj_c', 'DTt_des_c', 'rho_te_c', 'c_te_c', 'rho_ti_c', 'c_ti_c'])
            engine_var_units.extend(['kg/s', 'K', 'K', 'kg/m**3', 'm/s', 'kg/m**3', 'm/s'])

        # Fan parameters
        engine_var.extend(['DTt_f', 'mdot_f', 'N_f', 'A_f', 'd_f'])
        engine_var_units.extend(['K', 'kg/s', 'rpm', 'm**2', 'm'])

        return engine_var, engine_var_units

    def compute_size_output_mux(size_inputs: np.ndarray):
        """
        Compute vector size of the muxed trajectory.

        :param size_inputs:
        :type size_inputs: np.ndarray

        """

        mux_num = len(size_inputs)

        size_output = 0
        for i in range(mux_num):

            # Add input size to output vector
            if i < mux_num-1:
                size_output = size_output + (size_inputs[i]-1)
            else:
                size_output = size_output + (size_inputs[i])

        return size_output

    @staticmethod
    def compute_minimum_TS(settings: Settings, ac: Aircraft, engine: Engine) -> np.float64:
        """
        Compute minimum cutback thrust-setting meeting the 4%CG and one-engine-inoperative (OEI) airworthiness requirements.

        :param settings: pyNA settings
        :type settings: Settings
        :param ac: aircraft parameters
        :type ac: Aircraft
        :param engine: engine parameters
        :param engine: Engine
        :return: TS_max
        :rtype: np.float64

        """

        # Initialize limiting cases
        gamma_lst = np.array([0, 2.7])
        nr_engine_lst = np.array([ac.n_eng - 1, ac.n_eng])

        alpha = np.zeros(2)
        TS_lst = np.zeros(2)

        for cc, case in enumerate(['OEI', '4%CG']):
            # Compute atmospheric properties at ac.z_max
            prob_atm = om.Problem()
            prob_atm.model.add_subsystem("atm", Atmosphere(num_nodes=1, settings=settings))
            prob_atm.setup(force_alloc_complex=True)
            prob_atm.set_val('atm.z', ac.z_max)
            prob_atm.run_model()
            rho_0 = prob_atm.get_val('atm.rho_0')
            c_0 = prob_atm.get_val('atm.c_0')

            # Lift requirement for horizontal, steady climbing flight
            L = 9.80665 * ac.mtow * np.cos(gamma_lst[cc] * np.pi / 180.)
            c_l = L / (0.5* rho_0 * ac.v_max ** 2 * ac.af_S_w)

            # Compute required angle of attack to meet lift coefficient
            c_l_interp = RegularGridInterpolator((ac.aero['alpha'], ac.aero['theta_flaps'], ac.aero['theta_slats']), ac.aero['c_l'])
            c_l_data = c_l_interp((ac.aero['alpha'], settings.theta_flaps, settings.theta_slats))
            alpha[cc] = np.interp(c_l, c_l_data, ac.aero['alpha'])

            # Compute corresponding drag coefficient
            c_d_interp = RegularGridInterpolator((ac.aero['alpha'], ac.aero['theta_flaps'], ac.aero['theta_slats']), ac.aero['c_d'])
            c_d_data = c_d_interp((ac.aero['alpha'], settings.theta_flaps, settings.theta_slats))
            c_d = np.interp(alpha[cc], ac.aero['alpha'], c_d_data)

            # Compute thrust requirement
            D = (c_d * 0.5 * rho_0 * ac.v_max ** 2 * ac.af_S_w) + ac.mtow * 9.80065 * np.sin(gamma_lst[cc] * np.pi / 180.)
            F_req = D / nr_engine_lst[cc]

            # Compute thrust available
            F_n_interp = RegularGridInterpolator((engine.deck['z'], engine.deck['M_0'], engine.deck['TS']), engine.deck['F_n'])
            F_avl = F_n_interp([ac.z_max, ac.v_max / c_0, 1.])[0]

            # Compute minimum thrust setting
            TS_lst[cc] = F_req / F_avl
            # Print results
            print(case, 'engine thrust-setting requirement: ', np.round(TS_lst[cc], 3))

        # Compute TS_max
        TS_max = max(TS_lst)

        return TS_max

    def load_time_series(self, settings: Settings) -> None:
        """
        Loads predefined trajectory timeseries.

        :param settings: pyna settings
        :type settings: Settings

        :return: None

        """

        # Load trajectory data for the specific observer
        # Source: validation noise assessment data set of NASA STCA (Berton et al., 2019)
        self.path = pd.read_csv(settings.pyNA_directory + '/cases/' + settings.case_name + '/trajectory/' + settings.output_directory_name + '/' + settings.trajectory_file_name)
        self.n_t = np.size(self.path['t_source [s]'])

        return None

    def load_operating_point(self, settings:Settings, time_step: int) -> None:
        """
        Loads predefined trajectory timeseries.

        :param settings: pyna settings
        :type settings: Settings
        :param time_step: time step of the operating point in the trajectory time series
        :type time_step: int

        :return: None

        """

        # Load trajectory data for the specific observer
        # Source: validation noise assessment data set of NASA STCA (Berton et al., 2019)
        self.path = pd.read_csv(settings.pyNA_directory + '/cases/' + settings.case_name + '/trajectory/' + settings.output_directory_name + '/' + settings.trajectory_file_name)

        # Select operating point
        cols = self.path.columns
        op_point = pd.DataFrame(np.reshape(self.path.values[time_step, :], (1, len(cols))))
        op_point.columns = cols

        # Duplicate operating for theta range (np.linspace(0, 180, 19))
        self.path = pd.DataFrame()
        for i in np.arange(19):
            self.path = self.path.append(op_point)

        self.n_t = 19

        return None

    def setup(self, problem: om.Problem, settings: Settings, ac: Aircraft, engine: Engine, trajectory_mode: str, objective: str) -> None:
        """
        Setup take-off trajectory module using the following phases:

        * ``Ground roll``:  acceleration from V=0
        * ``Flaps down``:   accelerate further to V=kVstall while deploying flaps during x seconds
        * ``Rotation``:     rotate at dalpha/dt = cnst until load factor n=1;
        * ``Lift off``:     climb until obstacle is cleared (35ft).
        * ``VNRS``:         any VNRS can be applied in this phase, e.g. PTCB or PHLD
        * ``cutback``:      a pilot-initiated thrust cut-back to a constant thrust-setting is applied

        :param problem: openmdao problem
        :type problem: om.Problem
        :param settings: pyna settings
        :type settings: Settings
        :param ac: aircraft parameters
        :type ac: Aircraft
        :param engine: engine parameters
        :param engine: Engine
        :param objective: optimization objective
        :type objective: str

        :return: None

        """

        # Set solver settings for the problem
        problem.driver = om.pyOptSparseDriver(optimizer='IPOPT')
        problem.driver.opt_settings['print_level'] = 5
        problem.driver.opt_settings['nlp_scaling_method'] = 'gradient-based'

        problem.driver.declare_coloring(tol=1e-12)
        problem.model.linear_solver = om.LinearRunOnce()
        problem.driver.opt_settings['output_file'] = settings.pyNA_directory + '/cases/' + settings.case_name + '/output/' + settings.output_directory_name + '/IPOPT_trajectory_convergence.out'

        if objective == 'noise':
            problem.driver.opt_settings['tol'] = 1e-2
            problem.driver.opt_settings['acceptable_tol'] = 1e-1
        else:
            problem.driver.opt_settings['tol'] = settings.tol
            problem.driver.opt_settings['acceptable_tol'] = 1e-3

        problem.driver.opt_settings['max_iter'] = settings.max_iter
        problem.driver.opt_settings['mu_strategy'] = 'adaptive'
        problem.driver.opt_settings['bound_mult_init_method'] = 'mu-based'
        problem.driver.opt_settings['mu_init'] = 0.01
        problem.driver.opt_settings['constr_viol_tol'] = 1e-3
        problem.driver.opt_settings['compl_inf_tol'] = 1e-3
        problem.driver.opt_settings['acceptable_iter'] = 0
        problem.driver.opt_settings['acceptable_constr_viol_tol'] = 1e-1
        problem.driver.opt_settings['acceptable_compl_inf_tol'] = 1e-1
        problem.driver.opt_settings['acceptable_obj_change_tol'] = 1e-1

        # Setup trajectory and initialize trajectory transcription and compute number of points per phase
        traj = dm.Trajectory()
        problem.model.add_subsystem('phases', traj)

        # Get engine variables
        engine_var, engine_var_units = Trajectory.get_engine_variables(settings)

        # Compute the minimum thrust-setting based on 4% climb gradient and OEI requirement
        if settings.TS_cutback:
            TS_min = settings.TS_cutback
        else:
            TS_min = Trajectory.compute_minimum_TS(settings=settings, ac=ac, engine=engine)

        # Phase 1: ground roll
        if 'groundroll' in self.phase_name_lst:
            opts = {'phase': 'groundroll', 'ac': ac, 'engine': engine, 'settings': settings, 'objective': objective}
            self.phases['groundroll'] = dm.Phase(ode_class=TrajectoryODE, ode_init_kwargs=opts, transcription=self.transcription_phases[0])
            self.phases['groundroll'].set_time_options(fix_initial=True, duration_bounds=(0, 60), duration_ref=100.)
            self.phases['groundroll'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=True, fix_final=False, ref=1000.)
            self.phases['groundroll'].add_state('v', targets='v', rate_source='flight_dynamics.v_dot', units='m/s', fix_initial=True, fix_final=False, ref=100.)
            self.phases['groundroll'].add_state('alpha', targets='alpha', rate_source='flight_dynamics.alpha_dot', units='deg', fix_initial=True, fix_final=False, lower=ac.aero['alpha'][0], upper=ac.aero['alpha'][-1], ref=1.)
            self.phases['groundroll'].add_parameter('z', targets='z', units='m', val=0., dynamic=True,include_timeseries=True)
            self.phases['groundroll'].add_parameter('gamma', targets='gamma', units='deg', val=0., dynamic=True, include_timeseries=True)
            self.phases['groundroll'].add_parameter('TS', targets='propulsion.TS', units=None, val=settings.TS_to, dynamic=True, include_timeseries=True)
            self.phases['groundroll'].add_parameter('TS_min', units=None, val=1, dynamic=True, include_timeseries=True)
            if settings.PKROT:
                self.phases['groundroll'].add_parameter('k_rot', targets='flight_dynamics.k_rot', units=None, lower=1.1, upper=1.6, dynamic=False, val=ac.k_rot, opt=True)
            else:
                self.phases['groundroll'].add_parameter('k_rot', targets='flight_dynamics.k_rot', units=None, dynamic=False, val=ac.k_rot, opt=False)
            # PHLD
            # if objective == 'noise' and settings.PHLD:
                # self.phases['groundroll'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=ac.aero['theta_flaps_c_d_min_gr'], dynamic=True, include_timeseries=True, opt=False)
            # else:
            self.phases['groundroll'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, dynamic=True, include_timeseries=True)
            self.phases['groundroll'].add_parameter('theta_slats', targets='theta_slats', units='deg', val=settings.theta_slats, dynamic=True, include_timeseries=True)
            self.phases['groundroll'].add_timeseries('interpolated', transcription=dm.GaussLobatto(num_segments=self.phase_size[0]-1,order=3, solve_segments=False, compressed=True), subset='state_input')
            self.phases['groundroll'].add_boundary_constraint('flight_dynamics.v_rot_residual', equals=0., loc='final', ref=100, units='m/s')

        # Phase 2: rotation phase
        if 'rotation' in self.phase_name_lst:
            opts = {'phase': 'rotation', 'ac': ac, 'engine': engine, 'settings': settings, 'objective': objective}
            self.phases['rotation'] = dm.Phase(ode_class=TrajectoryODE, transcription=self.transcription_phases[1], ode_init_kwargs=opts)
            self.phases['rotation'].set_time_options(initial_bounds=(20, 60), duration_bounds=(0, 60), initial_ref=100., duration_ref=100.)
            self.phases['rotation'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=False, fix_final=False, ref=1000.)
            self.phases['rotation'].add_state('v', targets='v', rate_source='flight_dynamics.v_dot', units='m/s', fix_initial=False, fix_final=False, ref=100.)
            self.phases['rotation'].add_state('alpha', targets='alpha', rate_source='flight_dynamics.alpha_dot', units='deg', fix_initial=False, fix_final=False, lower=ac.aero['alpha'][0], upper=ac.aero['alpha'][-1], ref=10.)
            self.phases['rotation'].add_parameter('z', targets='z', units='m', val=0., dynamic=True,include_timeseries=True)
            self.phases['rotation'].add_parameter('gamma', targets='gamma', units='deg', val=0., dynamic=True, include_timeseries=True)
            self.phases['rotation'].add_parameter('TS', targets='propulsion.TS', units=None, val=settings.TS_to, dynamic=True, include_timeseries=True)
            self.phases['rotation'].add_parameter('TS_min', units=None, val=1, dynamic=True, include_timeseries=True)
            # PHLD
            # if objective == 'noise' and settings.PHLD:
                # self.phases['rotation'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, lower=ac.aero['theta_flaps'][0], upper=ac.aero['theta_flaps'][-1], dynamic=True, include_timeseries=True, opt=True, ref=10.)
                # self.phases['rotation'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=15., dynamic=True, include_timeseries=True)
            # else:
            self.phases['rotation'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, dynamic=True, include_timeseries=True)
            self.phases['rotation'].add_parameter('theta_slats', targets='theta_slats', units='deg', val=settings.theta_slats, dynamic=True, include_timeseries=True)
            self.phases['rotation'].add_timeseries('interpolated', transcription=dm.GaussLobatto(num_segments=self.phase_size[1]-1, order=3, solve_segments=False, compressed=True), subset='state_input')
            self.phases['rotation'].add_boundary_constraint('flight_dynamics.n', equals=1., loc='final', ref=1, units=None)

        # Phase 3: lift-off phase
        if 'liftoff' in self.phase_name_lst:
            opts = {'phase': 'liftoff', 'ac': ac, 'engine': engine, 'settings': settings, 'objective': objective}
            self.phases['liftoff'] = dm.Phase(ode_class=TrajectoryODE, transcription=self.transcription_phases[2], ode_init_kwargs=opts)
            self.phases['liftoff'].set_time_options(initial_bounds=(20, 150), duration_bounds=(0, 500), initial_ref=100., duration_ref=100., fix_duration=False)
            self.phases['liftoff'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=False, fix_final=False, ref=10000.)
            self.phases['liftoff'].add_state('z', rate_source='flight_dynamics.z_dot', units='m', fix_initial=False, fix_final=True, ref=10.)
            self.phases['liftoff'].add_state('v', targets='v', rate_source='flight_dynamics.v_dot', units='m/s', fix_initial=False, fix_final=False, ref=100.)
            self.phases['liftoff'].add_state('gamma', rate_source='flight_dynamics.gamma_dot', units='deg', fix_initial=False, fix_final=False, ref=10.)
            self.phases['liftoff'].add_control('alpha', targets='alpha', units='deg', lower=ac.aero['alpha'][0], upper=ac.aero['alpha'][-1], rate_continuity=True, rate_continuity_scaler=1.0, rate2_continuity=False, opt=True, ref=10.)
            self.phases['liftoff'].add_timeseries('interpolated', transcription=dm.GaussLobatto(num_segments=self.phase_size[2]-1, order=3, solve_segments=False, compressed=True), subset='state_input')
            self.phases['liftoff'].add_path_constraint(name='flight_dynamics.gamma_dot', lower=0., units='deg/s')
            self.phases['liftoff'].add_path_constraint(name='flight_dynamics.eas_dot', lower=0., units='m/s**2')
            self.phases['liftoff'].add_parameter('TS', targets='propulsion.TS', units=None, val=settings.TS_to, dynamic=True, include_timeseries=True)
            self.phases['liftoff'].add_parameter('TS_min', units=None, val=1, dynamic=True, include_timeseries=True)
            # PHLD
            # if objective == 'noise' and settings.PHLD:
                # self.phases['liftoff'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, lower=ac.aero['theta_flaps'][0], upper=ac.aero['theta_flaps'][-1], dynamic=True, include_timeseries=True, opt=True, ref=10.)
                # self.phases['liftoff'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=15, dynamic=True, include_timeseries=True, opt=True)
            # else:
            self.phases['liftoff'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, dynamic=True, include_timeseries=True)
            self.phases['liftoff'].add_parameter('theta_slats', targets='theta_slats', units='deg', val=settings.theta_slats, dynamic=True, include_timeseries=True)

        # Phase 4: vnrs phase
        if 'vnrs' in self.phase_name_lst:
            opts = {'phase': 'vnrs', 'ac': ac, 'engine': engine, 'settings': settings, 'objective': objective}
            self.phases['vnrs'] = dm.Phase(ode_class=TrajectoryODE, transcription=self.transcription_phases[3], ode_init_kwargs=opts)
            self.phases['vnrs'].set_time_options(initial_bounds=(20, 150), duration_bounds=(0, 500), initial_ref=100., duration_ref=100.)
            if trajectory_mode == 'flyover':
                self.phases['vnrs'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=False, fix_final=True, ref=10000.)
                self.phases['vnrs'].add_state('z', rate_source='flight_dynamics.z_dot', units='m', fix_initial=True, fix_final=False, ref=1000.)
            elif trajectory_mode == 'cutback':
                self.phases['vnrs'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=False, fix_final=False, ref=10000.)
                self.phases['vnrs'].add_state('z', rate_source='flight_dynamics.z_dot', units='m', fix_initial=True, fix_final=True, ref=1000.)
            self.phases['vnrs'].add_state('v', targets='v', rate_source='flight_dynamics.v_dot', units='m/s', fix_initial=False, fix_final=False, ref=100.)
            self.phases['vnrs'].add_state('gamma', rate_source='flight_dynamics.gamma_dot', units='deg', fix_initial=False, fix_final=False, ref=10.)
            self.phases['vnrs'].add_control('alpha', targets='alpha', units='deg', lower=5., upper=ac.aero['alpha'][-1], rate_continuity=True, rate_continuity_scaler=1.0, rate2_continuity=False, opt=True, ref=10.)
            self.phases['vnrs'].add_timeseries('interpolated', transcription=dm.GaussLobatto(num_segments=self.phase_size[3]-1, order=3, solve_segments=False, compressed=True), subset='state_input')
            self.phases['vnrs'].add_path_constraint(name='flight_dynamics.eas_dot', lower=0., units='m/s**2')
            self.phases['vnrs'].add_path_constraint(name='gamma', lower=0., units='deg', ref=10.)
            # PTCB
            if objective == 'noise' and settings.PTCB:
                self.phases['vnrs'].add_control('TS', targets='propulsion.TS', units=None, upper=settings.TS_to, lower=TS_min, val=TS_min, opt=True, rate_continuity=True, rate2_continuity=False, ref=1.)
                self.phases['vnrs'].add_path_constraint(name='TS', lower=TS_min, upper=1, units=None, ref=1.)
            else:
                self.phases['vnrs'].add_parameter('TS', targets='propulsion.TS', units=None, val=settings.TS_vnrs, dynamic=True, include_timeseries=True)
            # PHLD
            if objective == 'noise' and settings.PHLD:
                self.phases['vnrs'].add_control('theta_flaps', targets='theta_flaps', units='deg', val=0., lower=ac.aero['theta_flaps'][0], upper=ac.aero['theta_flaps'][-1], opt=True, rate_continuity=True, rate2_continuity=False, ref=1.)
                # self.phases['vnrs'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=0., lower=ac.aero['theta_flaps'][0], upper=ac.aero['theta_flaps'][-1], dynamic=True, include_timeseries=True, opt=True)
            else:
                self.phases['vnrs'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, dynamic=True, include_timeseries=True)
            self.phases['vnrs'].add_parameter('theta_slats', targets='theta_slats', units='deg', val=settings.theta_slats, dynamic=True, include_timeseries=True)

        # Phase 5: cutback phase
        if 'cutback' in self.phase_name_lst:
            opts = {'phase': 'cutback', 'ac': ac, 'engine': engine, 'settings': settings, 'objective': objective}
            self.phases['cutback'] = dm.Phase(ode_class=TrajectoryODE, transcription=self.transcription_phases[4], ode_init_kwargs=opts)
            self.phases['cutback'].set_time_options(initial_bounds=(20, 150), duration_bounds=(0, 500), initial_ref=100., duration_ref=100.)
            if trajectory_mode == 'flyover':
                self.phases['cutback'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=True, fix_final=False, ref=10000.)
                self.phases['cutback'].add_state('z', rate_source='flight_dynamics.z_dot', units='m', fix_initial=False, fix_final=True, ref=1000.)
            elif trajectory_mode == 'cutback':
                self.phases['cutback'].add_state('x', rate_source='flight_dynamics.x_dot', units='m', fix_initial=False, fix_final=False, ref=10000.)
                self.phases['cutback'].add_state('z', rate_source='flight_dynamics.z_dot', units='m', fix_initial=True, fix_final=True, ref=1000.)
            self.phases['cutback'].add_state('v', targets='v', rate_source='flight_dynamics.v_dot', units='m/s', fix_initial=False, fix_final=False, ref=100.)
            self.phases['cutback'].add_state('gamma', rate_source='flight_dynamics.gamma_dot', units='deg', fix_initial=False, fix_final=False, ref=10.)
            self.phases['cutback'].add_control('alpha', targets='alpha', units='deg', lower=ac.aero['alpha'][0], upper=ac.aero['alpha'][-1], rate_continuity=True, rate_continuity_scaler=1.0, rate2_continuity=False, opt=True, ref=10.)
            self.phases['cutback'].add_path_constraint(name='flight_dynamics.eas_dot', lower=0., units='m/s**2')
            self.phases['cutback'].add_path_constraint(name='flight_dynamics.gamma_dot', upper=0., units='deg/s')
            self.phases['cutback'].add_boundary_constraint('v', loc='final', upper=ac.v_max, ref=100., units='m/s')
            self.phases['cutback'].add_timeseries('interpolated', transcription=dm.GaussLobatto(num_segments=self.phase_size[4]-1, order=3, solve_segments=False, compressed=True), subset='state_input')
            # PTCB
            self.phases['cutback'].add_parameter('TS', targets='propulsion.TS', units=None, val=TS_min, dynamic=True, include_timeseries=True)
            # PHLD
            if objective == 'noise' and settings.PHLD:
                self.phases['cutback'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=0., lower=ac.aero['theta_flaps'][0], upper=ac.aero['theta_flaps'][-1], dynamic=True, include_timeseries=True, opt=True, ref=10.)
            else:
                self.phases['cutback'].add_parameter('theta_flaps', targets='theta_flaps', units='deg', val=settings.theta_flaps, dynamic=True, include_timeseries=True)
            self.phases['cutback'].add_parameter('theta_slats', targets='theta_slats', units='deg', val=settings.theta_slats, dynamic=True, include_timeseries=True)

        # Add outputs to timeseries for each phase
        for j, phase_name in enumerate(self.phase_name_lst):
            for i in np.arange(len(engine_var)):
                self.phases[phase_name].add_timeseries_output('propulsion.'+ engine_var[i], timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('aerodynamics.M_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('p_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('rho_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('I_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('drho_0_dz', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('T_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('c_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('mu_0', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('propulsion.W_f', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('emissions.mdot_NOx', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('emissions.EINOx', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('flight_dynamics.y', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('flight_dynamics.n', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('flight_dynamics.I_landing_gear', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('flight_dynamics.eas', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('flight_dynamics.n', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('aerodynamics.L', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('aerodynamics.D', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('propulsion.F_n', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('propulsion.W_f', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('clcd.c_l', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('clcd.c_l_max', timeseries='interpolated')
            self.phases[phase_name].add_timeseries_output('clcd.c_d', timeseries='interpolated')

        # Add phases to the trajectory
        for phase_name in self.phase_name_lst:
            traj.add_phase(phase_name, self.phases[phase_name])

        # Link phases
        if 'rotation' in self.phase_name_lst:
            traj.link_phases(phases=['groundroll', 'rotation'], vars=['time', 'x', 'v', 'alpha'])
            # if objective == 'noise' and settings.PHLD:
                # traj.add_linkage_constraint(phase_a='groundroll', phase_b='rotation', var_a='theta_flaps', var_b='theta_flaps', loc_a='final', loc_b='initial')

        if 'liftoff' in self.phase_name_lst:
            traj.link_phases(phases=['rotation', 'liftoff'], vars=['time', 'x', 'z', 'v', 'alpha', 'gamma'])
            # if objective == 'noise' and settings.PHLD:
                # traj.add_linkage_constraint(phase_a='rotation', phase_b='liftoff', var_a='theta_flaps', var_b='theta_flaps', loc_a='final', loc_b='initial')

        if 'vnrs' in self.phase_name_lst:
            traj.link_phases(phases=['liftoff', 'vnrs'],  vars=['time', 'x', 'v', 'alpha', 'gamma'])
            if objective == 'noise' and settings.PTCB:
                traj.add_linkage_constraint(phase_a='liftoff', phase_b='vnrs', var_a='TS', var_b='TS', loc_a='final', loc_b='initial')
            if objective == 'noise' and settings.PHLD:
                traj.add_linkage_constraint(phase_a='liftoff', phase_b='vnrs', var_a='theta_flaps', var_b='theta_flaps', loc_a='final', loc_b='initial')

        if 'cutback' in self.phase_name_lst:
            if trajectory_mode == 'flyover':
                traj.link_phases(phases=['vnrs', 'cutback'], vars=['time', 'z', 'v', 'alpha', 'gamma'])
            elif trajectory_mode == 'cutback':
                traj.link_phases(phases=['vnrs', 'cutback'], vars=['time', 'x', 'v', 'alpha', 'gamma'])
            # if objective == 'noise' and settings.PTCB:
                # traj.add_linkage_constraint(phase_a='vnrs', phase_b='cutback', var_a='TS', var_b='TS', loc_a='final', loc_b='initial')
            if objective == 'noise' and settings.PHLD:
                traj.add_linkage_constraint(phase_a='vnrs', phase_b='cutback', var_a='theta_flaps', var_b='theta_flaps', loc_a='final', loc_b='initial')

        # Mux trajectory variables
        mux_t = problem.model.add_subsystem(name='trajectory', subsys=Mux(size_inputs=np.array(self.phase_size), size_output=self.trajectory_size))
        var   = ['time', 'x', 'y', 'z', 'v', 'M_0', 'alpha', 'gamma', 'TS', 'I_landing_gear', 'theta_flaps', 'theta_slats','F_n', 'L', 'D', 'eas', 'n', 'p_0','rho_0','drho_0_dz','T_0','c_0','mu_0', 'I_0', 'W_f', 'mdot_NOx', 'EINOx', 'c_l', 'c_d', 'c_l_max']
        units = ['s', 'm', 'm', 'm', 'm/s', None, 'deg', 'deg', None, None, 'deg', 'deg', 'N', 'N', 'N', 'm/s', None, 'Pa', 'kg/m**3', 'kg/m**4', 'K', 'm/s', 'kg/m/s', 'kg/m**2/s', 'kg/s', 'kg/s', None, None, None, None]
        for i in np.arange(len(var)):
            # Add the variables to the trajectory mux
            if var[i] == 'time':
                mux_t.add_var('t_s', units=units[i])
            else:
                mux_t.add_var(var[i], units=units[i])

            # Connect phase variables to trajectory mux
            if var[i] == 'time':
                for j, phase_name in enumerate(self.phase_name_lst):
                    problem.model.connect('phases.' + phase_name + '.interpolated.' + var[i], 'trajectory.t_s_' + str(j))

            elif var[i] in ['x', 'v']:
                for j, phase_name in enumerate(self.phase_name_lst):
                    problem.model.connect('phases.' + phase_name + '.interpolated.states:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['z', 'gamma']:
                for j, phase_name in enumerate(self.phase_name_lst):
                    if phase_name in {'groundroll','rotation'}:
                        problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                    else:
                        problem.model.connect('phases.' + phase_name + '.interpolated.states:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['alpha']:
                for j, phase_name in enumerate(self.phase_name_lst):
                    if phase_name in {'groundroll', 'rotation'}:
                        problem.model.connect('phases.' + phase_name + '.interpolated.states:' + var[i],'trajectory.' + var[i] + '_' + str(j))
                    else:
                        problem.model.connect('phases.' + phase_name + '.interpolated.controls:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['TS']:
                if objective == 'noise' and settings.PTCB:
                    for j, phase_name in enumerate(self.phase_name_lst):
                        if phase_name in ['groundroll', 'rotation', 'liftoff']:
                            problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                        elif phase_name == 'vnrs':
                            problem.model.connect('phases.' + phase_name + '.interpolated.controls:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                        elif phase_name == 'cutback':
                            problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                else:
                    for j, phase_name in enumerate(self.phase_name_lst):
                        if phase_name in ['groundroll', 'rotation', 'liftoff', 'vnrs']:
                            problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                        elif phase_name == 'cutback':
                            problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['theta_slats']:
                for j, phase_name in enumerate(self.phase_name_lst):
                    problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['theta_flaps']:
                if objective == 'noise' and settings.PHLD:
                    for j, phase_name in enumerate(self.phase_name_lst):
                        if phase_name == 'vnrs':
                            problem.model.connect('phases.' + phase_name + '.interpolated.controls:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                        else:
                            problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))
                else:
                    for j, phase_name in enumerate(self.phase_name_lst):
                        problem.model.connect('phases.' + phase_name + '.interpolated.parameters:' + var[i], 'trajectory.' + var[i] + '_' + str(j))

            elif var[i] in ['y', 'F_n', 'L', 'D', 'eas', 'n', 'I_landing_gear','M_0', 'p_0','rho_0','drho_0_dz','T_0','c_0','c_bar','mu_0', 'I_0', 'W_f', 'mdot_NOx', 'EINOx', 'c_l', 'c_d', 'c_l_max']:
                for j, phase_name in enumerate(self.phase_name_lst):
                    problem.model.connect('phases.' + phase_name + '.interpolated.' + var[i], 'trajectory.' + var[i] + '_' + str(j))

        # Mux engine variables
        mux_e = problem.model.add_subsystem(name='engine', subsys=Mux(size_inputs=np.array(self.phase_size), size_output=self.trajectory_size))
        for i in np.arange(len(engine_var)):
            mux_e.add_var(engine_var[i], units=engine_var_units[i])
            for j, phase_name in enumerate(self.phase_name_lst):
                problem.model.connect('phases.' + phase_name + '.interpolated.' + engine_var[i], 'engine.' + engine_var[i] + '_' + str(j))

        return None

    def compute(self, problem: om.Problem, settings: Settings, ac: Aircraft, run_driver: bool, init_trajectory: om.Problem, trajectory_mode: str, objective: str) -> None:
        """
        Run trajectory initial guess with minimal time to climb.

        :param problem: openmdao problem
        :type problem: om.Problem
        :param settings: pyna settings
        :type settings: Settings
        :param ac: aircraft parameters
        :type ac: Aircraft
        :param run_driver: flag to enable run_driver setting for dymos run_model function
        :type run_driver: bool
        :param init_trajectory: initialization trajectory
        :type init_trajectory: om.Problem
        :param objective: optimization objective
        :type objective: str

        :return: None
        """

        # Add objective for trajectory model
        if objective == None:
            # No optimization objective required; problem is run with run_driver = False
            pass

        elif objective == 'x_end':
            problem.model.add_objective('trajectory.x', index=-1, ref=1000.)

        elif objective == 't_end':
            problem.model.add_objective('trajectory.t_s', index=-1, ref=1000.)

        elif objective == 'noise':
            # Optimization bjective is set in the noise.setup_trajectory_noise() method
            pass

        elif objective == 'x_takeoff_x_end':
            problem.model.add_subsystem('combined', om.ExecComp('objective = x_to/2000 + x_end/7000',
                                                                x_to={'val': 0., 'units': 'm'},
                                                                x_end={'val': 0., 'units': 'm'}),
                                                                promotes=['objective'])
            problem.model.connect('phases.liftoff.timeseries.states:x' , 'combined.x_to', src_indices=[-1])
            problem.model.connect('trajectory.x', 'combined.x_end', src_indices=[-1])
            problem.model.add_objective('objective', scaler=10.)

        elif objective == 'noise_surrogate':
            x_observer = np.array([6500., 0., 0.3048*4.])
            problem.model.add_subsystem('noise', SurrogateNoise(num_nodes=self.trajectory_size, x_observer=x_observer), promotes_outputs=['noise'])
            problem.model.connect('trajectory.x', 'noise.x')
            problem.model.connect('trajectory.y', 'noise.y')
            problem.model.connect('trajectory.z', 'noise.z')
            problem.model.connect('trajectory.t_s', 'noise.t_s')
            problem.model.add_objective('noise', scaler=1)

        else:
            raise ValueError('Invalid control objective specified.')

        # Run the openMDAO problem setup
        problem.setup(force_alloc_complex=True)

        # Attach a recorder to the problem to save model data
        if settings.save_results:
            problem.add_recorder(om.SqliteRecorder(settings.pyNA_directory + '/cases/' + settings.case_name + '/output/' + settings.output_directory_name + '/' + settings.output_file_name))

        # Set initial guess for the trajectory problem
        if init_trajectory is None:

            # Phase 1: groundroll
            if 'groundroll' in self.phase_name_lst:
                problem['phases.groundroll.t_initial'] = 0.0
                problem['phases.groundroll.t_duration'] = 30.0
                problem['phases.groundroll.states:x'] = self.phases['groundroll'].interp(ys=[0, 1000], nodes='state_input')
                problem['phases.groundroll.states:v'] = self.phases['groundroll'].interp(ys=[0.0, 60], nodes='state_input')
                problem['phases.groundroll.states:alpha'] = self.phases['groundroll'].interp(ys=[ac.alpha_0, ac.alpha_0], nodes='state_input')

            # Phase 2: rotation
            if 'rotation' in self.phase_name_lst:
                problem['phases.rotation.t_initial'] = 30.0
                problem['phases.rotation.t_duration'] = 10.0
                problem['phases.rotation.states:x'] = self.phases['rotation'].interp(ys=[1500, 2000], nodes='state_input')
                problem['phases.rotation.states:v'] = self.phases['rotation'].interp(ys=[100, 110.], nodes='state_input')
                problem['phases.rotation.states:alpha'] = self.phases['rotation'].interp(ys=[ac.alpha_0, 15*np.pi/180.], nodes='state_input')

            # Phase 3: lift-off
            if 'liftoff' in self.phase_name_lst:
                if trajectory_mode == 'flyover':
                    z_cutback_guess = 500.
                elif trajectory_mode == 'cutback':
                    z_cutback_guess = settings.z_cutback
                problem['phases.liftoff.t_initial'] = 40.0
                problem['phases.liftoff.t_duration'] = 2.
                problem['phases.liftoff.states:x'] = self.phases['liftoff'].interp(ys=[2000., 3500.], nodes='state_input')
                problem['phases.liftoff.states:z'] = self.phases['liftoff'].interp(ys=[0., 35*0.3048], nodes='state_input')
                problem['phases.liftoff.states:v'] = self.phases['liftoff'].interp(ys=[110., 110.], nodes='state_input')
                problem['phases.liftoff.states:gamma'] = self.phases['liftoff'].interp(ys=[0, 4.], nodes='state_input')
                problem['phases.liftoff.controls:alpha'] = self.phases['liftoff'].interp(ys=[15., 15.], nodes='control_input')

            # # Phase 4: vnrs
            if 'vnrs' in self.phase_name_lst:
                problem['phases.vnrs.t_initial'] = 50.0
                problem['phases.vnrs.t_duration'] = 50.0
                problem['phases.vnrs.states:x'] = self.phases['vnrs'].interp(ys=[3500., 6501.], nodes='state_input')
                problem['phases.vnrs.states:z'] = self.phases['vnrs'].interp(ys=[35*0.3048, z_cutback_guess], nodes='state_input')
                problem['phases.vnrs.states:v'] = self.phases['vnrs'].interp(ys=[110., 110.], nodes='state_input')
                problem['phases.vnrs.states:gamma'] = self.phases['vnrs'].interp(ys=[4., 15.], nodes='state_input')
                problem['phases.vnrs.controls:alpha'] = self.phases['vnrs'].interp(ys=[15., 15.], nodes='control_input')

            # Phase 5: cutback
            if 'cutback' in self.phase_name_lst:
                problem['phases.cutback.t_initial'] = 100.0
                problem['phases.cutback.t_duration'] = 50.0
                problem['phases.cutback.states:x'] = self.phases['cutback'].interp(ys=[6501., 20000.], nodes='state_input')
                problem['phases.cutback.states:z'] = self.phases['cutback'].interp(ys=[z_cutback_guess, ac.z_max], nodes='state_input')
                problem['phases.cutback.states:v'] = self.phases['cutback'].interp(ys=[110., 110.], nodes='state_input')
                problem['phases.cutback.states:gamma'] = self.phases['cutback'].interp(ys=[15, 15.], nodes='state_input')
                problem['phases.cutback.controls:alpha'] = self.phases['cutback'].interp(ys=[15., 15.], nodes='control_input')

        else:
            # Phase 1: groundroll
            if 'groundroll' in self.phase_name_lst:
                problem['phases.groundroll.t_initial'] = init_trajectory.get_val('phases.groundroll.t_initial')
                problem['phases.groundroll.t_duration'] = init_trajectory.get_val('phases.groundroll.t_duration')
                problem['phases.groundroll.timeseries.time'] = init_trajectory.get_val('phases.groundroll.timeseries.time')
                problem['phases.groundroll.states:x'] = init_trajectory.get_val('phases.groundroll.states:x')
                problem['phases.groundroll.states:v'] = init_trajectory.get_val('phases.groundroll.states:v')
                problem['phases.groundroll.states:alpha'] = init_trajectory.get_val('phases.groundroll.states:alpha')

            # Phase 2: rotation
            if 'rotation' in self.phase_name_lst:
                problem['phases.rotation.t_initial'] = init_trajectory.get_val('phases.rotation.t_initial')
                problem['phases.rotation.t_duration'] = init_trajectory.get_val('phases.rotation.t_duration')
                problem['phases.rotation.timeseries.time'] = init_trajectory.get_val('phases.rotation.timeseries.time')
                problem['phases.rotation.states:x'] = init_trajectory.get_val('phases.rotation.states:x')
                problem['phases.rotation.states:v'] = init_trajectory.get_val('phases.rotation.states:v')
                problem['phases.rotation.states:alpha'] = init_trajectory.get_val('phases.rotation.states:alpha')

            # Phase 3-5: liftoff-cutback
            for j, phase_name in enumerate(self.phase_name_lst[3:]):
                problem['phases.' + phase_name + '.t_initial'] = init_trajectory.get_val('phases.' + phase_name + '.t_initial')
                problem['phases.' + phase_name + '.t_duration'] = init_trajectory.get_val('phases.' + phase_name + '.t_duration')
                problem['phases.' + phase_name + '.timeseries.time'] = init_trajectory.get_val('phases.' + phase_name + '.timeseries.time')
                problem['phases.' + phase_name + '.states:x'] = init_trajectory.get_val('phases.' + phase_name + '.states:x')
                problem['phases.' + phase_name + '.states:z'] = init_trajectory.get_val('phases.' + phase_name + '.states:z')
                problem['phases.' + phase_name + '.states:v'] = init_trajectory.get_val('phases.' + phase_name + '.states:v')
                problem['phases.' + phase_name + '.states:gamma'] = init_trajectory.get_val('phases.' + phase_name + '.states:gamma')
                problem['phases.' + phase_name + '.controls:alpha'] = init_trajectory.get_val('phases.' + phase_name + '.controls:alpha')
                # if phase_name == 'vnrs' and settings.PTCB:
                    # problem['phases.' + phase_name + '.controls:TS'] = init_trajectory.get_val('phases.' + phase_name + '.controls:TS')

        # Run problem
        dm.run_problem(problem, run_driver=run_driver)

        # Save the results
        if settings.save_results:
            problem.record(case_name=settings.ac_name)

        # Write output
        return None

    @staticmethod
    def check_convergence(settings: Settings, filename: str) -> bool:
        """
        Checks convergence of case using optimizer output file.

        :param settings: pyna settings
        :type settings: Settings
        :param filename: file name of IPOPT output
        :type filename: str

        :return: converged
        :rtype: bool
        """

        # Save convergence info for trajectory
        # Read IPOPT file
        file_ipopt = open(settings.pyNA_directory + '/cases/' + settings.case_name + '/output/' + settings.output_directory_name + '/' + filename, 'r')
        ipopt = file_ipopt.readlines()
        file_ipopt.close()

        # Check if convergence summary excel file exists
        cnvg_file_name = settings.pyNA_directory + '/cases/' + settings.case_name + '/output/' + settings.output_directory_name + '/' + 'Convergence.csv'
        if not os.path.isfile(cnvg_file_name):
            file_cvg = open(cnvg_file_name, 'w')
            file_cvg.writelines("Trajectory name , Execution date/time,  Converged")
        else:
            file_cvg = open(cnvg_file_name, 'a')

        # Write convergence output to file
        # file = open(cnvg_file_name, 'a')
        if ipopt[-1] in {'EXIT: Optimal Solution Found.\n', 'EXIT: Solved To Acceptable Level.\n'}:
            file_cvg.writelines("\n" + settings.output_file_name + ", " + str(dt.datetime.now()) + ", Converged")
            converged = True
        else:
            file_cvg.writelines("\n" + settings.output_file_name + ", " + str(dt.datetime.now()) + ", Not converged")
            converged = False
        file_cvg.close()

        return converged


# !/usr/bin/env python
# -*- coding: utf-8 -*-
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2015-2018 European Synchrotron Radiation Facility, Grenoble, France
#
#    Principal author:       Jérôme Kieffer (Jerome.Kieffer@ESRF.eu)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

__author__ = "Jerome Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "10/01/2018"
__status__ = "development"
__docformat__ = 'restructuredtext'

import logging
import numpy

from . import detectors
import fabio
logger = logging.getLogger(__name__)

class Grid(object):
    """
    This class handles a regular grid in front of a detector to calibrate the
    geometrical distortion of the detector
    """
    def __init__(self, detector, image, mask=None, pitch=None, invert=False):
        """
        :param detector: instance of Detector or its name
        :param image: 2d array representing the image
        :param mask:
        :param pitch: 2-tuple representing the grid spacing in (y, x) coordinates, in meter
        :param invert: set to true if the image of the grid has regular dark spots (instead of bright points)
        """
        if isinstance(detector, detectors.Detector):
            self.detector = detectors.detector_factory(detector)
        else:
            self.detector = detector

        if isinstance(image, numpy.ndarray):
            self.image = image
        else:
            self.image = fabio.open(image).data

        if mask is not None:
            if isinstance(mask, numpy.ndarray):
                self.mask = mask
            else:
                self.mask = fabio.open(mask).data.astype(bool)
            if self.detector.mask is not None:
                self.mask = numpy.logical_or(self.detector.mask, self.mask)
        else:
            self.mask = numpy.zeros_like(self.image, bool)
        if invert:
            self.image = self.image.max() - self.image
        self.pitch = tuple(pitch[0], pitch[-1])

    def threshold(self, level=None, percentile=None):
        """
        Segment the image with a single threshold
        """
        if percentile and not level:
            data = self.image[self.mask]
            data.sort()
            level = data[int(len(data) * percentile / 100.)]
        raise NotImplementedError("TODO")


# -*- coding: utf-8 -*-
"""
refer: rlpack

Actor Critic 框架

Proximal Policy Optimization.
目标loss由三部分组成：1. clipped policy loss；2. value loss；3. entropy。
policy loss需要计算当前policy和old policy在当前state上的ratio。
需要注意的是，state分布依赖于old policy。因此，更新中的old policy是一样的。
"""

import numpy as np
import tensorflow as tf

from .base import Base

class PPO(Base):
    def __init__(self,
                 rnd=0,
                 dim_obs=None, dim_act=None,
                 policy_fn=None, value_fn=None,
                 discount=0.99, gae=0.95, clip_ratio=0.2,
                 train_epoch=40, policy_lr=1e-3, value_lr=1e-3,
                 save_path="./log", log_freq=10, save_model_freq=100):
        self._dim_obs = dim_obs
        self._dim_act = dim_act
        self._policy_fn = policy_fn
        self._value_fn = value_fn

        self._clip_ratio = clip_ratio
        self._discount = discount
        self._gae = gae
        self._train_epoch = train_epoch
        self._policy_lr = policy_lr
        self._value_lr = value_lr

        self._log_freq = log_freq
        self._save_model_freq = save_model_freq

        super().__init__(save_path=save_path, rnd=rnd)

    def _build_network(self):
        """Build tensorflow operations for algorithms."""
        self._obs = tf.placeholder(tf.float32, [None, *self._dim_obs])
        self._act = tf.placeholder(tf.int32, [None, self._dim_act])

        self._adv = tf.placeholder(tf.float32, [None])
        self._ret = tf.placeholder(tf.float32, [None])
        self._logp_old = tf.placeholder(tf.float32, [None])
        self.all_phs = [self._obs, self._act, self._adv, self._ret, self._logp_old]

        self.pi, self.logp_list = self._policy_fn(self._obs)
        self.logp = tf.gather_nd(self.logp_list, self._act)
        self.v = self._value_fn(self._obs)

    def _build_algorithm(self):
        """Build algorithms using prebuilt networks."""

        ratio = tf.exp(self.logp - self._logp_old)
        surr1 = ratio * self._adv
        surr2 = tf.clip_by_value(ratio, 1.0 - self._clip_ratio, 1.0 + self._clip_ratio) * self._adv
        self.policy_loss = -tf.reduce_mean(tf.minimum(surr1, surr2))
        self.value_loss = tf.reduce_mean((self.v - self._ret)**2)

        self._train_policy_op = tf.train.AdamOptimizer(self._policy_lr).minimize(self.policy_loss)
        self._train_value_op = tf.train.AdamOptimizer(self._value_lr).minimize(self.value_loss)

    def get_action(self, obs) -> np.ndarray:
        """Return action according to the observations.
        :param obs: the observation that could be image or real-number features
        :return: actions
        """
        a_prob = self.sess.run(self.pi, feed_dict={self._obs: obs})

        return a_prob

    def update(self, databatch):
        """
        参数:
            databatch：一个列表，分别是state, action, reward, done, early_stop, next_state。每个是矩阵或向量。
            state是状态，action是动作，reward是奖励，done是是否完结，early_stop是是否提前结束，next_state是下一个状态。
        """
        preprocess_databatch = self._parse_databatch(*databatch)

        inputs = {k: v for k, v in zip(self.all_phs, preprocess_databatch)}
        pi_l_old, v_l_old = self.sess.run([self.policy_loss, self.value_loss], feed_dict=inputs)

        # Training
        for i in range(self._train_epoch):
            self.sess.run(self._train_policy_op, feed_dict=inputs)
        for i in range(self._train_epoch):
            self.sess.run(self._train_value_op, feed_dict=inputs)

        pi_l_new, v_l_new = self.sess.run([self.policy_loss, self.value_loss], feed_dict=inputs)

        global_step, _ = self.sess.run([tf.train.get_global_step(), self.increment_global_step])
        if global_step % self._save_model_freq == 0:
            self.save_model()

    def _parse_databatch(self, states, actions, rewards, dones, earlystops, nextstates):

        batch_size = len(dones)
        oldlogproba, values = self.sess.run([self.logp, self.v], feed_dict={self._obs: states, self._act: actions})
        nextvalues = self.sess.run(self.v, feed_dict={self._obs: nextstates})

        returns = np.zeros(batch_size)
        deltas = np.zeros(batch_size)
        advantages = np.zeros(batch_size)

        for i in reversed(range(batch_size)):

            if dones[i]:
                prev_return = 0
                prev_value = 0
                prev_advantage = 0
            elif earlystops[i]:
                prev_return = nextvalues[i]
                prev_value = prev_return
                prev_advantage = 0

            returns[i] = rewards[i] + self._discount * prev_return * (1 - dones[i])
            deltas[i] = rewards[i] + self._discount * prev_value * (1 - dones[i]) - values[i]
            # ref: https://arxiv.org/pdf/1506.02438.pdf (generalization advantage estimate)
            advantages[i] = deltas[i] + self._discount * self._gae * prev_advantage * (1 - dones[i])

            prev_return = returns[i]
            prev_value = values[i]
            prev_advantage = advantages[i]

        advantages = (advantages - advantages.mean()) / advantages.std()

        return [states, actions, advantages, returns, oldlogproba]


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from dropblock import DropBlock2D

class BEV_Unet(nn.Module):

    def __init__(self,n_class,n_height,dilation = 1,group_conv=False,input_batch_norm = False,dropout = 0.,circular_padding = False, dropblock = True, use_vis_fea=False):
        super(BEV_Unet, self).__init__()
        self.n_class = n_class
        self.n_height = n_height
        if use_vis_fea:
            self.network = UNet(n_class*n_height,2*n_height,dilation,group_conv,input_batch_norm,dropout,circular_padding,dropblock)
        else:
            self.network = UNet(n_class*n_height,n_height,dilation,group_conv,input_batch_norm,dropout,circular_padding,dropblock)

    def forward(self, x):
        x,center,offset = self.network(x)

        x = x.permute(0,2,3,1)
        new_shape = list(x.size())[:3] + [self.n_height,self.n_class]
        x = x.view(new_shape)
        x = x.permute(0,4,1,2,3)

        return x,center,offset

class UNet(nn.Module):
    def __init__(self, n_class,n_height,dilation,group_conv,input_batch_norm, dropout,circular_padding,dropblock):
        super(UNet, self).__init__()
        # encoder
        self.inc = inconv(n_height, 64, dilation, input_batch_norm, circular_padding)
        self.down1 = down(64, 128, dilation, group_conv, circular_padding)
        self.down2 = down(128, 256, dilation, group_conv, circular_padding)
        self.down3 = down(256, 512, dilation, group_conv, circular_padding)
        self.down4 = down(512, 512, dilation, group_conv, circular_padding)

        # semantic decoder
        self.up1 = up(1024, 256, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.up2 = up(512, 128, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.up3 = up(256, 64, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.up4 = up(128, 64, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.dropout = nn.Dropout(p=0. if dropblock else dropout)
        # semantic head
        self.outc = outconv(64, n_class)

        # instance decoder
        # self.i_up1 = up(1024, 256, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        # self.i_up2 = up(512, 128, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        # self.i_up3 = up(256, 64, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        # self.i_up4 = up(128, 32, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.i_up4_center = up(128, 32, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        self.i_up4_offset = up(128, 32, circular_padding, group_conv = group_conv, use_dropblock=dropblock, drop_p=dropout)
        # instance head
        self.i_outc_center = outconv(32, 1)
        self.i_outc_offset = outconv(32, 2)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        # semantic
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        s_x = self.up4(x, x1)
        s_x = self.outc(self.dropout(s_x))
        # instance
        # i_x = self.i_up1(x5, x4)
        # i_x = self.i_up2(i_x, x3)
        # i_x = self.i_up3(i_x, x2)x

        # i_x = self.i_up4(i_x, x1)
        i_x_center = self.i_up4_center(x, x1)
        i_x_center = self.i_outc_center(self.dropout(i_x_center))

        i_x_offset = self.i_up4_offset(x, x1)
        i_x_offset = self.i_outc_offset(self.dropout(i_x_offset))

        return s_x, i_x_center, i_x_offset

class double_conv(nn.Module):
    '''(conv => BN => ReLU) * 2'''
    def __init__(self, in_ch, out_ch,group_conv,dilation=1):
        super(double_conv, self).__init__()
        if group_conv:
            self.conv = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1,groups = min(out_ch,in_ch)),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True),
                nn.Conv2d(out_ch, out_ch, 3, padding=1,groups = out_ch),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )
        else:
            self.conv = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True),
                nn.Conv2d(out_ch, out_ch, 3, padding=1),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )

    def forward(self, x):
        x = self.conv(x)
        return x

class double_conv_circular(nn.Module):
    '''(conv => BN => ReLU) * 2'''
    def __init__(self, in_ch, out_ch,group_conv,dilation=1):
        super(double_conv_circular, self).__init__()
        if group_conv:
            self.conv1 = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=(1,0),groups = min(out_ch,in_ch)),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )
            self.conv2 = nn.Sequential(
                nn.Conv2d(out_ch, out_ch, 3, padding=(1,0),groups = out_ch),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )
        else:
            self.conv1 = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=(1,0)),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )
            self.conv2 = nn.Sequential(
                nn.Conv2d(out_ch, out_ch, 3, padding=(1,0)),
                nn.BatchNorm2d(out_ch),
                nn.LeakyReLU(inplace=True)
            )

    def forward(self, x):
        #add circular padding
        x = F.pad(x,(1,1,0,0),mode = 'circular')
        x = self.conv1(x)
        x = F.pad(x,(1,1,0,0),mode = 'circular')
        x = self.conv2(x)
        return x

class inconv(nn.Module):
    def __init__(self, in_ch, out_ch, dilation, input_batch_norm, circular_padding):
        super(inconv, self).__init__()
        if input_batch_norm:
            if circular_padding:
                self.conv = nn.Sequential(
                    nn.BatchNorm2d(in_ch),
                    double_conv_circular(in_ch, out_ch,group_conv = False,dilation = dilation)
                )
            else:
                self.conv = nn.Sequential(
                    nn.BatchNorm2d(in_ch),
                    double_conv(in_ch, out_ch,group_conv = False,dilation = dilation)
                )
        else:
            if circular_padding:
                self.conv = double_conv_circular(in_ch, out_ch,group_conv = False,dilation = dilation)
            else:
                self.conv = double_conv(in_ch, out_ch,group_conv = False,dilation = dilation)

    def forward(self, x):
        x = self.conv(x)
        return x

class down(nn.Module):
    def __init__(self, in_ch, out_ch, dilation, group_conv, circular_padding):
        super(down, self).__init__()
        if circular_padding:
            self.mpconv = nn.Sequential(
                nn.MaxPool2d(2),
                double_conv_circular(in_ch, out_ch,group_conv = group_conv,dilation = dilation)
            )
        else:
            self.mpconv = nn.Sequential(
                nn.MaxPool2d(2),
                double_conv(in_ch, out_ch,group_conv = group_conv,dilation = dilation)
            )

    def forward(self, x):
        x = self.mpconv(x)
        return x

class up(nn.Module):
    def __init__(self, in_ch, out_ch, circular_padding, bilinear=True, group_conv=False, use_dropblock = False, drop_p = 0.5):
        super(up, self).__init__()

        #  would be a nice idea if the upsampling could be learned too,
        #  but my machine do not have enough memory to handle all those weights
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        elif group_conv:
            self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2,groups = in_ch//2)
        else:
            self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)

        if circular_padding:
            self.conv = double_conv_circular(in_ch, out_ch,group_conv = group_conv)
        else:
            self.conv = double_conv(in_ch, out_ch,group_conv = group_conv)

        self.use_dropblock = use_dropblock
        if self.use_dropblock:
            self.dropblock = DropBlock2D(block_size=7, drop_prob=drop_p)

    def forward(self, x1, x2):
        x1 = self.up(x1)

        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,
                        diffY // 2, diffY - diffY//2))

        # for padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd

        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        if self.use_dropblock:
            x = self.dropblock(x)
        return x

class outconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(outconv, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 1)

    def forward(self, x):
        x = self.conv(x)
        return x


# -*- coding: utf-8 -*-
"""
This module contains the functions necessary to correct atom probe datasets
for systematic energy deficits.  Details on the algorithm are given in a
manuscript titled:
    ``An Algorithm for Correcting Systematic Energy Deficits in the Atom Probe
    Mass Spectra of Insulating Samples''

Created on Mon Nov 25 14:02:12 2019

@author: bwc

>  NIST Public License - 2019

>  This software was developed by employees of the National Institute of
>  Standards and Technology (NIST), an agency of the Federal Government
>  and is being made available as a public service. Pursuant to title 17
>  United States Code Section 105, works of NIST employees are not subject
>  to copyright protection in the United States.  This software may be
>  subject to foreign copyright.  Permission in the United States and in
>  foreign countries, to the extent that NIST may hold copyright, to use,
>  copy, modify, create derivative works, and distribute this software and
>  its documentation without fee is hereby granted on a non-exclusive basis,
>  provided that this notice and disclaimer of warranty appears in all copies.

>  THE SOFTWARE IS PROVIDED 'AS IS' WITHOUT ANY WARRANTY OF ANY KIND,
>  EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED
>  TO, ANY WARRANTY THAT THE SOFTWARE WILL CONFORM TO SPECIFICATIONS, ANY
>  IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
>  AND FREEDOM FROM INFRINGEMENT, AND ANY WARRANTY THAT THE DOCUMENTATION
>  WILL CONFORM TO THE SOFTWARE, OR ANY WARRANTY THAT THE SOFTWARE WILL BE
>  ERROR FREE.  IN NO EVENT SHALL NIST BE LIABLE FOR ANY DAMAGES, INCLUDING,
>  BUT NOT LIMITED TO, DIRECT, INDIRECT, SPECIAL OR CONSEQUENTIAL DAMAGES,
>  ARISING OUT OF, RESULTING FROM, OR IN ANY WAY CONNECTED WITH THIS SOFTWARE,
>  WHETHER OR NOT BASED UPON WARRANTY, CONTRACT, TORT, OR OTHERWISE, WHETHER
>  OR NOT INJURY WAS SUSTAINED BY PERSONS OR PROPERTY OR OTHERWISE, AND
>  WHETHER OR NOT LOSS WAS SUSTAINED FROM, OR AROSE OUT OF THE RESULTS OF,
>  OR USE OF, THE SOFTWARE OR SERVICES PROVIDED HEREUNDER.

"""

import numpy as np
import histo_funcs as hf
from scipy.interpolate import interp1d

def get_shifts(ref, N, max_shift=150, global_shift_fun=None):
    """
    This computes the optimal shift between a 1d reference array and the
    columns of a 2d data array using an fft based cross correlation.  The
    optimal shift is assumed to be where the cross correlation is maximal.
    Because no padding or normalization is applied, this works best for
    relatively small shifts.  This is explicitly enforced by the ``max_shift``
    parameter.

    Parameters
    ----------
    ref : real numeric 1d array (shape is nx1)
        The reference array.  Should be in a nx1 column array.
    N : real numeric 2d array (shape is nxm)
        The data array.  Each column of this array will be cross-correlated to
        the reference array using an fft.
    max_shift : int (default is 150)
        The maximum number of bins that are searched for the cross-correlation.
        The cross correlation is set to zero for lags less than -max_shift and
        greater than +max_shift.
    global_shift_fun : function handle (default is None)
        This function applies an additional subtractive shift to the returned
        shift array.  I generally set this to numpy.mean or numpy.median, but
        more complex shift functions could be used (including lambdas).

    Returns
    -------
    shifts : numeric 1d array
        The shift (in units of bin index) that was required to align each
        column of data to the reference spectrum.

    """
    # Note: Use real valued fft to improve speed/memory

    # FFT the ref and data arrays
    rfft_ref = np.fft.rfft(ref, axis=1)
    rfft_N = np.fft.rfft(N, axis=1)

    # Compute the cross-correlation and take the inverse fft
    xc = np.fft.irfft(rfft_N*np.conj(rfft_ref), axis=1)

    # Set the cross correlation to zero for lags greater than max_shift
    xc[:, max_shift:xc.shape[1]-max_shift] = 0

    # Find the lags corresponding to the maximum of the cross correlation and
    # then shift them to correspond to the appropriate origin.
    max_idxs = np.argmax(xc, axis=1)
    max_idxs[max_idxs > xc.shape[1]//2] = \
        max_idxs[max_idxs > xc.shape[1]//2] - xc.shape[1]

    # Apply a global_shift_fun shift if specified
    if global_shift_fun is not None:
        shifts = max_idxs - global_shift_fun(max_idxs)
    else:
        shifts = max_idxs

    return shifts

def get_all_scale_coeffs(event_dat,
                         max_scale=1.1,
                         roi=None,
                         cts_per_chunk=2**10,
                         delta_logdat=5e-4):
    """
    This attempts to best align event data.  The basic assumption is that if
    the ``event_dat`` is binned into a 2d history (histogram) then there are
    clearly defined features (i.e. peaks) that will shift around in a
    systematic manner.  When this data is projected onto a single dimension
    then any features (peaks) will be broader than they really should be.  This
    algorithm discretizes the event_dat into `chunks' and attempts to align
    each chunk to a reference dataset using a scalar multiplicative
    coefficient.  The reference dataset is the middle 50% of the ``event_dat``.
    The alignment is performed using a logarithm based cross correlation
    approach.  Two iterations of this algorithm are performed before the result
    is returned.  In all tests performed thus far a single iteration was
    sufficient however we used two iterations in an abundance of caution.

    Parameters
    ----------
    event_dat : real float 1d array
        Event data.  Typically the data is either the mass-to-charge or
        time-of-flight of each event.  The ordering of the data is assumed to
        be chronological (i.e. the order in which they were detected).
    max_scale : real float (default is 1.1)
        The maximum possible scale factor allowed (relative to the reference
        data)
    roi : real numeric list or array (default is [0.5,200])
        The domain that the data should be evaluated over.
        Specified as [min, max] values.
    cts_per_chunk : int (default is 1024)
        The number of events to be collected into a single `chunk' of data.
    delta_logdat : real float (default is 5e-4)
        The discretization of the log(data) over the roi specified.  Smaller
        deltas are more time/memory intensive.  For deltas much less than one,
        this effectively gives a discretization/resolution of the
        multiplicative factor of 1+delta_logdat.  For the atom probe data I
        have worked with, the noise on the shift is on the order of 1e-3 and
        so setting the delta to be smaller than this, ensures that the
        discretization error is not a significant problem.

    Returns
    -------
    eventwise_scales : real float array
        An array that contains the computed scale factor for each event that
        best aligns the data.  To correct the data, just divide the event_dat
        array by the eventwise_scales array.

    """
    if roi is None:
        roi = [0.5, 200]
    log_roi = np.log(roi)

    # Take the log of data
    logdat = np.log(event_dat)

    # Create the histogram.  Compute centers and delta y
    N, seq_edges, logdat_edges = \
        hf.create_histogram(logdat,
                            roi=log_roi,
                            cts_per_chunk=cts_per_chunk,
                            delta_dat=delta_logdat)
    seq_centers, logdat_centers = hf.edges_to_centers(seq_edges, logdat_edges)
#    print('specified delta_logdat = '+str(delta_logdat))
    delta_logdat = logdat_edges[1]-logdat_edges[0]
#    print('actual delta_logdat = '+str(delta_logdat))

    # Initialize the total eventwise log(dat) shift
    eventwise_logdat_shifts = np.zeros(event_dat.size)

    # Do one iteration with the center 50% of the data as a reference
    # Note: Make it is 2d (even though it is just a single column array)
    ref = np.mean(N[N.shape[0]//4:3*N.shape[0]//4, :], axis=0)[None, :]

    # Get the maximum possible shift in bins.
    max_pixel_shift = int(np.ceil(np.log(max_scale)/delta_logdat))

    # Determine the chunkwise shifts
    chunkwise_shifts0 = delta_logdat*get_shifts(ref,
                                                N,
                                                max_shift=max_pixel_shift,
                                                global_shift_fun=np.mean)

    # Interpolate (linear) from chunkwise to eventwise shifts
    f = interp1d(seq_centers, chunkwise_shifts0, fill_value='extrapolate')

    # Accumulate the shift for the first iteration.
    eventwise_logdat_shifts += f(np.arange(event_dat.size))

    # Correct the log(data)
    logdat_corr = logdat - eventwise_logdat_shifts

    # Recompute the histogram with newly corrected log(data)
    N, seq_edges, logdat_edges = \
        hf.create_histogram(logdat_corr,
                            roi=log_roi,
                            cts_per_chunk=cts_per_chunk,
                            delta_dat=delta_logdat)
    seq_centers, logdat_centers = hf.edges_to_centers(seq_edges, logdat_edges)
    delta_logdat = logdat_edges[1]-logdat_edges[0]

    # Use the center 50% of the data as a reference
    # Note: Make it is 2d (even though it is just a single column array)
    ref = np.mean(N[N.shape[0]//4:3*N.shape[0]//4, :], axis=0)[None, :]

    # Get the maximum possible shift in bins.
    max_pixel_shift = int(np.ceil(np.log(max_scale)/delta_logdat))

    # Determine the chunkwise shifts
    chunkwise_shifts1 = delta_logdat*get_shifts(ref,
                                                N,
                                                max_shift=max_pixel_shift,
                                                global_shift_fun=np.mean)

    # Interpolate to get eventwise shifts
    f = interp1d(seq_centers, chunkwise_shifts1, fill_value='extrapolate')

    # Accumulate the shift for the second iteration.
    eventwise_logdat_shifts += f(np.arange(event_dat.size))

    # Compute total eventwise shifts for output
    eventwise_scales = np.exp(eventwise_logdat_shifts)

#    # Uncomment this to see the relative importance of the two iterations
#    import matplotlib.pyplot as plt
#    plt.figure()
#    plt.plot(np.exp(chunkwise_shifts0), label='iter 0')
#    plt.plot(np.exp(chunkwise_shifts1), label='iter 1')
#    plt.legend()

    return eventwise_scales