"""
Copyright (c) 2021, Electric Power Research Institute

 All rights reserved.

 Redistribution and use in source and binary forms, with or without modification,
 are permitted provided that the following conditions are met:

     * Redistributions of source code must retain the above copyright notice,
       this list of conditions and the following disclaimer.
     * Redistributions in binary form must reproduce the above copyright notice,
       this list of conditions and the following disclaimer in the documentation
       and/or other materials provided with the distribution.
     * Neither the name of DER-VET nor the names of its contributors
       may be used to endorse or promote products derived from this software
       without specific prior written permission.

 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
"""
Storage

This Python class contains methods and attributes specific for technology analysis within StorageVet.
"""
import logging
import cvxpy as cvx
import numpy as np
import pandas as pd
from storagevet.Technology.DistributedEnergyResource import DER
from dervet.MicrogridDER.DERExtension import DERExtension
from dervet.MicrogridDER.ContinuousSizing import ContinuousSizing
from storagevet.ErrorHandling import *

class ElectricVehicle1(DER, ContinuousSizing, DERExtension):
    """ A general template for storage object

    We define "storage" as anything that can affect the quantity of load/power being delivered or used. Specific
    types of storage are subclasses. The storage subclass should be called. The storage class should never
    be called directly.

    """

    def __init__(self, params):
        """ Initialize all technology with the following attributes.

        Args:
            params (dict): Dict of parameters
        """
        TellUser.debug(f"Initializing ElectricVehicle1")
        # create generic technology object
        DER.__init__(self, params)
        ContinuousSizing.__init__(self, params)
        DERExtension.__init__(self, params)

        self.technology_type = 'Electric Vehicle'
        self.tag = 'ElectricVehicle1'

        self.ene_target = params['ene_target']
        self.ch_max_rated = params['ch_max_rated']
        self.ch_min_rated = params['ch_min_rated']

        self.plugin_time = params['plugin_time']
        self.plugout_time = params['plugout_time']

        self.capital_cost_function = params['ccost']

        self.fixed_om = params['fixed_om']
        self.incl_binary = False  # params['binary'] #TODO

        self.variable_names = {'ene', 'ch', 'uene', 'uch', 'on_c'}

        # initialize any attributes you want to set and use later
        self.plugin_times_index = None
        self.plugout_times_index = None
        self.unplugged_index = None

    # def charge_capacity(self):
    #     """

    #     Returns: the maximum charge that can be attained

    #     """
    #     return self.ch_max_rated

    def initialize_variables(self, size):
        """ Adds optimization variables to dictionary

        Variables added: (with self.unique_ess_id as a prefix to these)
            ene (Variable): A cvxpy variable for Energy collected at the end of the time step (kWh)
            ch (Variable): A cvxpy variable for Charge Power, kW during the previous time step (kW)
            on_c (Variable/Parameter): A cvxpy variable/parameter to flag for charging in previous interval (bool)

        Notes:
            CVX Parameters turn into Variable when the condition to include them is active

        Args:
            size (Int): Length of optimization variables to create

        """
        self.variables_dict = {
            'ene': cvx.Variable(shape=size, name=self.name + '-ene'),
            'ch': cvx.Variable(shape=size, name=self.name + '-ch'),
            'uene': cvx.Variable(shape=size, name=self.name + '-uene'),  # TODO you can switch to parameter, where value == np.zeros(size)  -HN
            'uch': cvx.Variable(shape=size, name=self.name + '-uch'),  # TODO you can switch to parameter, where value == np.zeros(size)  -HN
            'on_c': cvx.Parameter(shape=size, name=self.name + '-on_c', value=np.ones(size)),

        }

        if self.incl_binary:
            self.variable_names.update(['on_c'])
            self.variables_dict.update({'on_c': cvx.Variable(shape=size, boolean=True, name=self.name + '-on_c')})

    def get_state_of_energy(self, mask):
        """
        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                in the subs data set

        Returns: the state of energy as a function of time for the

        """
        return self.variables_dict['ene']

    def get_charge(self, mask):
        """
        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                in the subs data set

        Returns: the charge as a function of time for the

        """
        return self.variables_dict['ch']

    def get_capex(self, **kwargs):
        """ Returns the capex of a given technology
        """
        return self.capital_cost_function

    def get_charge_up_schedule(self, mask):
        """ the amount of charging power in the up direction (supplying power up into the grid) that
        this DER can schedule to reserve

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns: CVXPY parameter/variable

        """
        return self.variables_dict['ch'] - self.ch_min_rated

    def get_charge_down_schedule(self, mask):
        """ the amount of charging power in the up direction (pulling power down from the grid) that
        this DER can schedule to reserve

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns: CVXPY parameter/variable

        """
        return self.ch_max_rated - self.variables_dict['ch']

    def get_delta_uenegy(self, mask):
        """ the amount of energy, from the current SOE level the DER's state of energy changes
        from subtimestep energy shifting

        Returns: the energy throughput in kWh for this technology

        """
        return self.variables_dict['uene']

    def get_uenergy_increase(self, mask):
        """ the amount of energy in a timestep that is provided to the distribution grid

        Returns: the energy throughput in kWh for this technology

        """
        return self.variables_dict['uch'] * self.dt

    def get_active_times(self, mask):
        """

        Args:
            mask:

        Returns:

        """
        compute_plugin_index = pd.DataFrame(index=mask.index)
        compute_plugin_index['plugin'] = compute_plugin_index.index.hour == self.plugin_time
        compute_plugin_index['plugout'] = compute_plugin_index.index.hour == self.plugout_time
        compute_plugin_index['unplugged'] = False

        if self.plugin_time < self.plugout_time:  # plugin time and plugout time must be different
            compute_plugin_index.loc[
                (compute_plugin_index.index.hour >= self.plugin_time) * (compute_plugin_index.index.hour < self.plugout_time), 'unplugged'] = True
        elif self.plugin_time > self.plugout_time:
            compute_plugin_index.loc[
                (compute_plugin_index.index.hour >= self.plugin_time) | (compute_plugin_index.index.hour < self.plugout_time), 'unplugged'] = True

        self.plugout_times_index = compute_plugin_index['plugout']
        self.plugin_times_index = compute_plugin_index['plugin']
        self.unplugged_index = compute_plugin_index['unplugged']

    def constraints(self, mask):
        """Default build constraint list method. Used by services that do not have constraints.

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns:
            A list of constraints that corresponds the EV requirement to collect the required energy to operate. It also allows
            flexibility to provide other grid services
        """

        constraint_list = []
        self.get_active_times(mask.loc[mask])  # constructing the array that indicates whether the ev is plugged or not

        # print(self.plugin_times_index.iloc[0:24])
        # print(self.plugout_times_index.iloc[0:24])
        # print(self.unplugged_index.iloc[0:24])
        # print('Ene target :' + str(self.ene_target))
        # print('Charging max :' + str(self.ch_max_rated))
        # print('Charging min :' + str(self.ch_min_rated))

        # optimization variables
        ene = self.variables_dict['ene']
        ch = self.variables_dict['ch']
        uene = self.variables_dict['uene']
        uch = self.variables_dict['uch']
        on_c = self.variables_dict['on_c']

        # collected energy at start time is zero for all start times
        constraint_list += [cvx.Zero(ene[self.plugin_times_index])]

        # energy evolution generally for every time step

        numeric_unplugged_index = pd.Series(range(len(self.unplugged_index)), index=self.unplugged_index.index).loc[self.unplugged_index]
        ene_ini_window = 0

        if numeric_unplugged_index.iloc[0] == 0:  # energy evolution for the EV, only during plugged times
            constraint_list += [cvx.Zero(ene[numeric_unplugged_index.iloc[0]] - ene_ini_window)]
            constraint_list += [cvx.Zero(ene[list(numeric_unplugged_index.iloc[1:])] - ene[list(numeric_unplugged_index.iloc[1:] - 1)] - (
                        self.dt * ch[list(numeric_unplugged_index.iloc[1:] - 1)]))]  # - uene[list(numeric_unplugged_index.iloc[1:]-1)])]
        else:
            constraint_list += [cvx.Zero(ene[list(numeric_unplugged_index)] - ene[list(numeric_unplugged_index - 1)] - (
                        self.dt * ch[list(numeric_unplugged_index - 1)]))]  # - uene[list(numeric_unplugged_index-1)])]
        # constraint_list += [cvx.Zero(ene[1:] - ene[:-1]  - ( self.dt * ch[:-1]) - uene[:-1])]

        # energy at plugout times must be greater or equal to energy target

        numeric_plugout_time_index = pd.Series(range(len(self.plugout_times_index)), index=self.plugout_times_index.index).loc[
            self.plugout_times_index]

        # the next few lines make sure that the state of energy at the end of the chargign period is equal to the target
        if numeric_plugout_time_index[0] == 0:
            constraint_list += [cvx.Zero(self.ene_target - ene[list(numeric_plugout_time_index.iloc[1:] - 1)] - (
                        self.dt * ch[list(numeric_plugout_time_index.iloc[1:] - 1)]))]  # - uene[list(numeric_plugout_time_index.iloc[1:]-1)])]
        else:
            constraint_list += [cvx.Zero(self.ene_target - ene[list(numeric_plugout_time_index - 1)] - (
                        self.dt * ch[list(numeric_plugout_time_index - 1)]))]  # - uene[list(numeric_plugout_time_index-1)])]

        constraint_list += [cvx.Zero(ene[list(numeric_plugout_time_index)] - self.ene_target)]

        # constraints on the ch/dis power

        # make it MILP or not depending on user selection
        if self.incl_binary:
            constraint_list += [cvx.NonPos(ch - (on_c * self.ch_max_rated))]
            constraint_list += [cvx.NonPos((on_c * self.ch_min_rated) - ch)]
        else:
            constraint_list += [cvx.NonPos(ch - self.ch_max_rated)]
            constraint_list += [cvx.NonPos(- ch)]

        # constraints to make sure that the ev does nothing when it is unplugged
        constraint_list += [cvx.NonPos(ch[~self.unplugged_index])]

        # account for -/+ sub-dt energy -- this is the change in energy that the battery experiences as a result of energy option
        # constraint_list += [cvx.Zero(uene - (uch * self.dt))]
        constraint_list += [cvx.Zero(uch)]  # TODO: you can set the variable to be parameters instead  -HN
        constraint_list += [cvx.Zero(uene)]  # TODO: you can set the variable to be parameters instead  -HN
        return constraint_list

    def timeseries_report(self):
        """ Summaries the optimization results for this DER.

        Returns: A timeseries dataframe with user-friendly column headers that
            summarize the results pertaining to this instance

        """
        tech_id = self.unique_tech_id()
        results = pd.DataFrame(index=self.variables_df.index)
        solve_dispatch_opt = self.variables_df.get('ch')
        if solve_dispatch_opt is not None:
            results[tech_id + ' Charge (kW)'] = self.variables_df['ch']
            results[tech_id + ' Power (kW)'] = -self.variables_df['ch']
            results[tech_id + ' State of Energy (kWh)'] = \
                self.variables_df['ene']

            results[tech_id + ' Energy Option (kWh)'] = \
                self.variables_df['uene']
            results[tech_id + ' Charge Option (kW)'] = self.variables_df['uch']

        return results

    def proforma_report(self, apply_inflation_rate_func, fill_forward_func, results):
        """ Calculates the proforma that corresponds to participation in this value stream

        Args:
            apply_inflation_rate_func:
            fill_forward_func:
            results (pd.DataFrame):

        Returns: A DateFrame of with each year in opt_year as the index and
            the corresponding value this stream provided.

        """
        pro_forma = super().proforma_report(apply_inflation_rate_func, fill_forward_func, results)
        if self.variables_df.empty:
            return pro_forma
        analysis_years = self.variables_df.index.year.unique()
        om_costs = pd.DataFrame()
        for year in analysis_years:
            # add fixed o&m costs
            index_yr = pd.Period(year=year, freq='y')
            om_costs.loc[index_yr, self.fixed_column_name()] = -self.fixed_om
        # fill forward
        om_costs = fill_forward_func(om_costs, None)
        # apply inflation rates
        om_costs = apply_inflation_rate_func(om_costs, None, min(analysis_years))
        # append will super class's proforma
        pro_forma = pd.concat([pro_forma, om_costs], axis=1)
        return pro_forma

    def sizing_summary(self):
        """ Creates the template for sizing df that each DER must fill to report their size.

        Returns: A dictionary describe this DER's size and captial costs.

        """
        # template = pd.DataFrame(columns=)
        sizing_dict = {
            'DER': np.nan,
            'Energy Rating (kWh)': np.nan,
            'Charge Rating (kW)': np.nan,
            'Discharge Rating (kW)': np.nan,
            'Round Trip Efficiency (%)': np.nan,
            'Lower Limit on SOC (%)': np.nan,
            'Upper Limit on SOC (%)': np.nan,
            'Duration (hours)': np.nan,
            'Capital Cost ($)': np.nan,
            'Capital Cost ($/kW)': np.nan,
            'Capital Cost ($/kWh)': np.nan,
            'Power Capacity (kW)': np.nan,
            'Quantity': 1,
        }
        return sizing_dict

class ElectricVehicle2(DER, ContinuousSizing, DERExtension):
    """ A general template for storage object

    We define "storage" as anything that can affect the quantity of load/power being delivered or used. Specific
    types of storage are subclasses. The storage subclass should be called. The storage class should never
    be called directly.

    """

    def __init__(self, params):
        """ Initialize all technology with the following attributes.

        Args:
            params (dict): Dict of parameters
        """
        TellUser.debug(f"Initializing ElectricVehicle2")
        # create generic technology object
        DER.__init__(self, params)
        ContinuousSizing.__init__(self, params)
        DERExtension.__init__(self, params)

        self.technology_type = 'Electric Vehicle'
        self.tag = 'ElectricVehicle2'

        # input params
        # note: these should never be changed in simulation (i.e from degradation)

        self.max_load_ctrl = params[
                                 'max_load_ctrl'] / 100.0  # maximum amount of baseline EV load that can be shed as a percentage of the original load
        # self.qualifying_cap = params['qualifying_cap'] #capacity that can be used for 'capacity' services (DR, RA) as a percentage of the baseline load
        self.lost_load_cost = params['lost_load_cost']
        self.incl_binary = params['binary']
        self.EV_load_TS = params['EV_baseline']

        self.capital_cost_function = params['ccost']

        self.fixed_om = params['fixed_om']

        self.variable_names = {'ch'}

    def qualifying_capacity(self, event_length):
        """ Describes how much power the DER can discharge to qualify for RA or DR. Used to determine
        the system's qualifying commitment.

        Args:
            event_length (int): the length of the RA or DR event, this is the
                total hours that a DER is expected to discharge for

        Returns: int/float

        """
        return 0

    def initialize_variables(self, size):
        """ Adds optimization variables to dictionary

        Variables added: (with self.unique_ess_id as a prefix to these)
            ene (Variable): A cvxpy variable for Energy collected at the end of the time step (kWh)
            ch (Variable): A cvxpy variable for Charge Power, kW during the previous time step (kW)
            on_c (Variable/Parameter): A cvxpy variable/parameter to flag for charging in previous interval (bool)

        Notes:
            CVX Parameters turn into Variable when the condition to include them is active

        Args:
            size (Int): Length of optimization variables to create

        """
        self.variables_dict = {
            'ch': cvx.Variable(shape=size, name=self.name + '-ch')
        }

        if self.incl_binary:
            self.variable_names.update(['on_c'])
            self.variables_dict.update({'on_c': cvx.Variable(shape=size, boolean=True, name=self.name + '-on_c')})

    def get_charge(self, mask):
        """
        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                in the subs data set

        Returns: the charge as a function of time for the

        """
        return self.variables_dict['ch']

    def get_capex(self, **kwargs):
        """ Returns the capex of a given technology
        """
        return self.capital_cost_function

    def get_charge_up_schedule(self, mask):
        """ the amount of charging power in the up direction (supplying power up into the grid) that
        this DER can schedule to reserve

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns: CVXPY parameter/variable

        """
        return self.variables_dict['ch'] - (1 - self.max_load_ctrl) * self.EV_load_TS[mask]

    def get_charge_down_schedule(self, mask):
        """ the amount of charging power in the up direction (pulling power down from the grid) that
        this DER can schedule to reserve

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns: CVXPY parameter/variable

        """
        return -self.variables_dict['ch'] + self.EV_load_TS[mask]

    def objective_function(self, mask, annuity_scalar=1):
        """ Generates the objective function related to a technology. Default includes O&M which can be 0

        Args:
            mask (Series): Series of booleans used, the same length as case.power_kw
            annuity_scalar (float): a scalar value to be multiplied by any yearly cost or benefit that helps capture the cost/benefit over
                    the entire project lifetime (only to be set iff sizing, else annuity_scalar should not affect the aobject function)

        Returns:
            self.costs (Dict): Dict of objective costs
        """

        # create objective expression for variable om based on discharge activity
        ch = self.variables_dict['ch']
        costs = {
            self.name + ' fixed_om': self.fixed_om * annuity_scalar,
            self.name + ' lost_load_cost': cvx.sum(self.EV_load_TS[mask].values - ch) * self.lost_load_cost  # added to account for lost load

        }
        # add startup objective costs

        return costs

    def constraints(self, mask):
        """Default build constraint list method. Used by services that do not have constraints.

        Args:
            mask (DataFrame): A boolean array that is true for indices corresponding to time_series data included
                    in the subs data set

        Returns:
            A list of constraints that corresponds the EV requirement to collect the required energy to operate. It also allows
            flexibility to provide other grid services
        """
        constraint_list = []

        # optimization variables

        ch = self.variables_dict['ch']
        # uch = self.variables_dict['uch']

        # constraints on the ch/dis power
        constraint_list += [cvx.NonPos(ch - self.EV_load_TS[mask].values)]
        constraint_list += [cvx.NonPos((1 - self.max_load_ctrl) * self.EV_load_TS[mask].values - ch)]

        # the constraint below limits energy throughput and total discharge to less than or equal to
        # (number of cycles * energy capacity) per day, for technology warranty purposes
        # this constraint only applies when optimization window is equal to or greater than 24 hours

        return constraint_list

    def timeseries_report(self):
        """ Summaries the optimization results for this DER.

        Returns: A timeseries dataframe with user-friendly column headers that
            summarize the results pertaining to this instance

        """
        tech_id = self.unique_tech_id()
        results = pd.DataFrame(index=self.variables_df.index)
        solve_dispatch_opt = self.variables_df.get('ch')
        if solve_dispatch_opt is not None:
            results[tech_id + ' EV Fleet Baseline Load'] = self.EV_load_TS
            results[tech_id + ' Charge (kW)'] = self.variables_df['ch']
            results[tech_id + ' Power (kW)'] = -self.variables_df['ch']

        return results

    def proforma_report(self, apply_inflation_rate_func, fill_forward_func, results):
        """ Calculates the proforma that corresponds to participation in this value stream

        Args:
            apply_inflation_rate_func:
            fill_forward_func:
            results (pd.DataFrame):

        Returns: A DateFrame of with each year in opt_year as the index and
            the corresponding value this stream provided.

        """
        pro_forma = super().proforma_report(apply_inflation_rate_func, fill_forward_func, results)
        if self.variables_df.empty:
            return pro_forma
        analysis_years = self.variables_df.index.year.unique()
        om_costs = pd.DataFrame()
        for year in analysis_years:
            # add fixed o&m costs
            index_yr = pd.Period(year=year, freq='y')
            om_costs.loc[index_yr, self.fixed_column_name()] = -self.fixed_om
        # fill forward
        om_costs = fill_forward_func(om_costs, None)
        # apply inflation rates
        om_costs = apply_inflation_rate_func(om_costs, None, min(analysis_years))
        # append will super class's proforma
        pro_forma = pd.concat([pro_forma, om_costs], axis=1)
        return pro_forma

    def sizing_summary(self):
        """ Creates the template for sizing df that each DER must fill to report their size.

        Returns: A dictionary describe this DER's size and captial costs.

        """
        # template = pd.DataFrame(columns=)
        sizing_dict = {
            'DER': np.nan,
            'Energy Rating (kWh)': np.nan,
            'Charge Rating (kW)': np.nan,
            'Discharge Rating (kW)': np.nan,
            'Round Trip Efficiency (%)': np.nan,
            'Lower Limit on SOC (%)': np.nan,
            'Upper Limit on SOC (%)': np.nan,
            'Duration (hours)': np.nan,
            'Capital Cost ($)': np.nan,
            'Capital Cost ($/kW)': np.nan,
            'Capital Cost ($/kWh)': np.nan,
            'Power Capacity (kW)': np.nan,
            'Quantity': 1,
        }
        return sizing_dict


"""
Tensorflow SMPL implementation as batch.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import pickle as pickle

import tensorflow as tf
from .batch_lbs import batch_rodrigues, batch_global_rigid_transformation

# There are chumpy variables so convert them to numpy.
def undo_chumpy(x):
    return x if isinstance(x, np.ndarray) else x.r

class SMPL(object):
    def __init__(self, pkl_path, joint_type='cocoplus', dtype=tf.float64):
        """
        pkl_path is the path to a SMPL model
        """
        # -- Load SMPL params --
        with open(pkl_path, 'rb') as f:
            dd = pickle.load(f, encoding="latin1")
        self.dtype = dtype
        # Mean template vertices
        self.v_template = tf.Variable(undo_chumpy(dd['v_template']),
                                      name='v_template',
                                      dtype=self.dtype,
                                      trainable=False)
        self.f = dd['f']

        # Size of mesh [Number of vertices, 3]
        self.size = [self.v_template.shape[0].value, 3]
        self.num_betas = dd['shapedirs'].shape[-1]
        self.num_verts = dd['shapedirs']
        self.num_joints = dd['J'].shape[0]

        # Shape blend shape basis: num_verts x 3 x num_betas
        # reshaped to 3*num_verts x num_betas, transposed to num_betas x 3*num_verts
        shapedir = np.reshape(undo_chumpy(dd['shapedirs']), [-1, self.num_betas]).T
        self.shapedirs = tf.Variable(shapedir, name='shapedirs', dtype=self.dtype, trainable=False)

        # Regressor for joint locations given shape - num_verts x num_joints
        self.J_regressor = tf.Variable(dd['J_regressor'].T.todense(),
                                       name="J_regressor",
                                       dtype=self.dtype,
                                       trainable=False)

        # Pose blend shape basis: num_verts x 3 x 9*num_joints, reshaped to 3*num_verts x 9*num_joints
        num_pose_basis = dd['posedirs'].shape[-1]
        posedirs = np.reshape(undo_chumpy(dd['posedirs']), [-1, num_pose_basis]).T
        self.posedirs = tf.Variable(posedirs, name='posedirs', dtype=self.dtype, trainable=False)

        # indices of parents for each joints
        self.parents = dd['kintree_table'][0].astype(np.int32)

        # LBS weights
        self.weights = tf.Variable(undo_chumpy(dd['weights']),
                                   name='lbs_weights',
                                   dtype=self.dtype,
                                   trainable=False)

    def __call__(self, trans, beta, theta, name=''):
        """
        Obtain SMPL with shape (beta) & pose (theta) inputs.
        Theta includes the global rotation.
        Args:
          beta: N x num_betas
          theta: N x 3*num_joints (with 3-D axis-angle rep)

        Updates:
        self.J_transformed: N x num_joints x 3 joint location after shaping
                 & posing with beta and theta

        Returns:
          - Verts: N x num_verts x 3
        """

        with tf.name_scope(name, "smpl_main", [beta, theta]):
            num_batch = beta.shape[0].value

            # 1. Add shape blend shapes
            # (N x num_betas) x (num_betas x 3*num_verts) = N x num_verts x 3
            v_shaped = tf.reshape(tf.matmul(beta, self.shapedirs, name='shape_bs'),
                                  [-1, self.size[0], self.size[1]]) + self.v_template

            # 2. Infer shape-dependent joint locations.
            Jx = tf.matmul(v_shaped[:, :, 0], self.J_regressor)
            Jy = tf.matmul(v_shaped[:, :, 1], self.J_regressor)
            Jz = tf.matmul(v_shaped[:, :, 2], self.J_regressor)
            J = tf.stack([Jx, Jy, Jz], axis=2)

            # 3. Add pose blend shapes
            # N x num_joints x 3 x 3
            Rs = tf.reshape(batch_rodrigues(tf.reshape(theta, [-1, 3])), [-1, self.num_joints, 3, 3])
            with tf.name_scope("lrotmin"):
                # Ignore global rotation.
                pose_feature = tf.reshape(Rs[:, 1:, :, :] - tf.eye(3, dtype=self.dtype), [-1, 9*(self.num_joints-1)])

            # (N x 9*(num_joints-1))) x (9*(num_joints-1), 3*num_verts) -> N x num_verts x 3
            v_posed = tf.reshape(tf.matmul(pose_feature, self.posedirs),
                                 [-1, self.size[0], self.size[1]]) + v_shaped

            #4. Get the global joint location
            self.J_transformed, A = batch_global_rigid_transformation(Rs, J, self.parents)

            # 5. Do skinning:
            # W is N x num_verts x num_joints
            W = tf.reshape(tf.tile(self.weights, [num_batch, 1]), [num_batch, -1, self.num_joints])

            # (N x num_verts x num_joints) x (N x num_joints x 16)
            T = tf.reshape(
                tf.matmul(W, tf.reshape(A, [num_batch, self.num_joints, 16])),
                [num_batch, -1, 4, 4])
            v_posed_homo = tf.concat([v_posed, tf.ones([num_batch, v_posed.shape[1], 1], dtype=self.dtype)], 2)
            v_homo = tf.matmul(T, tf.expand_dims(v_posed_homo, -1))

            return tf.add(v_homo[:, :, :3, 0], trans)


import os
from typing import List, Dict
import numpy as np
import h5py as h5
from scipy.sparse import csr_matrix
from mpi4py import MPI

from .__init__ import (
	pprint,
	comm,
	rank,
	nproc,
)
from .conversion import (
	comoving_density,
	comoving_length,
	comoving_velocity,
	comoving_mass,
	comoving_kinetic_energy,
	comoving_momentum,
	comoving_ang_momentum,
	density_units,
	velocity_units,
	length_units,
	mass_units,
	momentum_units,
	energy_units,
)

def split(nfiles):
    nfiles=int(nfiles)
    nf=int(nfiles/nproc)
    rmd=nfiles % nproc
    st=rank*nf
    fh=(rank+1)*nf
    if rank < rmd:
        st+=rank
        fh+=(rank+1)
    else:
        st+=rmd
        fh+=rmd
    return st,fh

def commune(data):
    tmp=np.zeros(nproc,dtype=np.int)
    tmp[rank]=len(data)
    cnts=np.zeros(nproc,dtype=np.int)
    comm.Allreduce([tmp,MPI.INT],[cnts,MPI.INT],op=MPI.SUM)
    del tmp
    dspl=np.zeros(nproc,dtype=np.int)
    i=0
    for j in range(nproc):
        dspl[j]=i
        i+=cnts[j]
    rslt=np.zeros(i,dtype=data.dtype)
    comm.Allgatherv([data,cnts[rank]],[rslt,cnts,dspl,MPI._typedict[data.dtype.char]])
    del data,cnts,dspl
    return rslt

def compute_M(data):
    cols = np.arange(data.size)
    return csr_matrix((cols, (data.ravel(), cols)), shape=(data.max() + 1, data.size))

def get_indices_sparse(data):
    M = compute_M(data)
    return [np.unravel_index(row.data, data.shape) for row in M]

def find_files(redshift: str) -> list:
    z_value = ['z004p688', 'z004p061', 'z003p053', 'z003p078', 'z002p688', 'z002p349',
            'z002p053', 'z001p792', 'z001p561', 'z001p354', 'z001p168', 'z001p000', 'z000p846', 'z000p706',
            'z000p577', 'z000p457', 'z000p345', 'z000p240', 'z000p140', 'z000p046', 'z000p000']
    z_IDNumber = ['002', '003', '004', '005', '006', '007', '008', '009', '010', '011', '012', '013',
            '014', '015', '016', '017', '018', '019', '020', '021', '022']
    sn = dict(zip(z_value, z_IDNumber))[redshift]
    path='/cosma5/data/dp004/dc-hens1/macsis/macsis_gas'
    pprint(f"[+] Find simulation files {redshift:s}...")

    halos = [x for x in os.listdir(path) if x.startswith('halo_')]
	master_files = []
	for halo in halos:
		groups = []
		for x in os.listdir(os.path.join(path, halo, f'data/groups_{sn}')):
			if x.startswith('eagle_subfind_tab'):
				complete_path = os.path.join(path, halo, f'data/groups_{sn}/{x}')
				groups.append(complete_path)

		particle = []
		for x in os.listdir(os.path.join(path, halo, f'data/particledata_{sn}')):
			if x.startswith('eagle_subfind_particles'):
				complete_path = os.path.join(path, halo, f'data/particledata_{sn}/{x}')
				particle.append(complete_path)

		master_files.append([groups[0], particle[0]])

	return master_files

def fof_header(files: list):
	pprint(f"[+] Find header information...")
	header = {}
	with h5.File(files[0][1], 'r') as f:
		header['Hub']  = f['Header'].attrs['HubbleParam']
		header['aexp'] = f['Header'].attrs['ExpansionFactor']
		header['zred'] = f['Header'].attrs['Redshift']
		header['OmgL'] = f['Header'].attrs['OmegaLambda']
		header['OmgM'] = f['Header'].attrs['Omega0']
		header['OmgB'] = f['Header'].attrs['OmegaBaryon']
	return header

def fof_groups(files: list):
	pprint(f"[+] Find groups information...")
	group_files = [pair[0] for pair in files]
	st, fh = split(len(group_files))
	Mfof = np.empty(0, dtype=np.float32)
	M2500 = np.empty(0, dtype=np.float32)
	M500 = np.empty(0, dtype=np.float32)
	M200 = np.empty(0, dtype=np.float32)
	R2500 = np.empty(0, dtype=np.float32)
	R500 = np.empty(0, dtype=np.float32)
	R200 = np.empty(0, dtype=np.float32)
	COP = np.empty(0, dtype=np.float32)
	NSUB = np.empty(0, dtype=np.int)
	FSID = np.empty(0, dtype=np.int)
	SCOP = np.empty(0, dtype=np.float32)
	for x in range(st, fh, 1):
		with h5.File(group_files[x], 'r') as f:
			Mfof = np.append(Mfof, f['FOF/GroupMass'][:])
			M2500 = np.append(M2500, f['FOF/Group_M_Crit2500'][:])
			R2500 = np.append(R2500, f['FOF/Group_R_Crit2500'][:])
			M500 = np.append(M500, f['FOF/Group_M_Crit500'][:])
			R500 = np.append(R500, f['FOF/Group_R_Crit500'][:])
			M200 = np.append(M200, f['FOF/Group_M_Crit200'][:])
			R200 = np.append(R200, f['FOF/Group_R_Crit200'][:])
			COP = np.append(COP, f['FOF/GroupCentreOfPotential'][:])
			NSUB = np.append(NSUB, f['FOF/NumOfSubhalos'][:])
			FSID = np.append(FSID, f['FOF/FirstSubhaloID'][:])
			SCOP = np.append(SCOP, f['Subhalo/CentreOfPotential'][:])

	header = {}
	with h5.File(group_files[0], 'r') as f:
		header['Hub'] =  f['Header'].attrs['HubbleParam']
		header['aexp'] = f['Header'].attrs['ExpansionFactor']
		header['zred'] = f['Header'].attrs['Redshift']

	# Conversion
	Mfof = comoving_mass(header, Mfof * 1.0e10)
	M2500 = comoving_mass(header, M2500 * 1.0e10)
	M500 = comoving_mass(header, M500 * 1.0e10)
	M200 = comoving_mass(header, M200 * 1.0e10)
	R2500 = comoving_length(header, R2500)
	R500 = comoving_length(header, R500)
	R200 = comoving_length(header, R200)
	COP = comoving_length(header, COP)
	SCOP = comoving_length(header, SCOP)

	data = {}
	data['groupfiles'] = np.asarray(group_files)
	data['particlefiles'] = np.asarray([pair[1] for pair in files])
	data['Mfof'] = commune(Mfof)
	data['M2500'] = commune(M2500)
	data['R2500'] = commune(R2500)
	data['M500'] = commune(M500)
	data['R500'] = commune(R500)
	data['M200'] = commune(M200)
	data['R200'] = commune(R200)
	data['COP']  = commune(COP.reshape(-1, 1)).reshape(-1, 3)
	data['NSUB'] = commune(NSUB)
	data['FSID'] = commune(FSID)
	data['SCOP'] = commune(SCOP.reshape(-1, 1)).reshape(-1, 3)

	return data

def fof_group(clusterID: int, fofgroups: Dict[str, np.ndarray] = None):
	pprint(f"[+] Find group information for cluster {clusterID}")
	new_data = {}
	new_data['clusterID'] = clusterID
	new_data['Mfof']  = fofgroups['Mfof'][clusterID]
	new_data['M2500'] = fofgroups['M2500'][clusterID]
	new_data['R2500'] = fofgroups['R2500'][clusterID]
	new_data['M500']  = fofgroups['M500'][clusterID]
	new_data['R500']  = fofgroups['R500'][clusterID]
	new_data['M200']  = fofgroups['M200'][clusterID]
	new_data['R200']  = fofgroups['R200'][clusterID]
	new_data['COP']   = fofgroups['COP'][clusterID]
	new_data['NSUB']  = fofgroups['NSUB'][clusterID]
	new_data['FSID']  = fofgroups['FSID'][clusterID]
	new_data['SCOP']  = fofgroups['SCOP'][clusterID]
	new_data['groupfiles']  = fofgroups['groupfiles'][clusterID]
	new_data['particlefiles'] = fofgroups['particlefiles'][clusterID]
	return new_data

def cluster_partgroupnumbers(fofgroup: Dict[str, np.ndarray] = None):
	"""

	:param fofgroups:
	:return:
	"""
	pgn = []
	with h5.File(fofgroup['particlefiles'], 'r') as h5file:

		for pt in ['0', '1', '4']:
			Nparticles = h5file['Header'].attrs['NumPart_ThisFile'][int(pt)]
			st, fh = split(Nparticles)
			pprint(f"[+] Collecting particleType {pt} GroupNumber...")
			groupnumber = h5file[f'/PartType{pt}/GroupNumber'][st:fh]

			# Clip out negative values and exceeding values
			groupnumber = np.clip(groupnumber, 0, 6)
			pprint(f"\t Computing CSR indexing matrix...")
			groupnumber_csrm = get_indices_sparse(groupnumber)
			del groupnumber_csrm[0], groupnumber_csrm[-1]
			pgn.append(groupnumber_csrm)
			del groupnumber

	return pgn

def cluster_particles(fofgroup: Dict[str, np.ndarray] = None, groupNumbers: List[np.ndarray] = None):
	"""

	:param fofgroup:
	:param groupNumbers:
	:return:
	"""
	pprint(f"[+] Find particle information for cluster {fofgroup['clusterID']}")
	data_out = {}
	header = {}
	partTypes = ['0', '1', '4']
	with h5.File(fofgroup['particlefiles'], 'r') as h5file:

		header['Hub']  = h5file['Header'].attrs['HubbleParam']
		header['aexp'] = h5file['Header'].attrs['ExpansionFactor']
		header['zred'] = h5file['Header'].attrs['Redshift']

		for pt in partTypes:

			# Initialise particledata arrays
			pgn_core = np.empty(0, dtype=np.int)
			subgroup_number = np.empty(0, dtype=np.int)
			velocity = np.empty(0, dtype=np.float32)
			coordinates = np.empty(0, dtype=np.float32)
			mass = np.empty(0, dtype=np.float32)
			temperature = np.empty(0, dtype=np.float32)
			sphdensity = np.empty(0, dtype=np.float32)
			sphlength = np.empty(0, dtype=np.float32)

			# Let each CPU core import a portion of the pgn data
			pgn = groupNumbers[partTypes.index(pt)]
			st, fh = split(len(pgn))
			pgn_core = np.append(pgn_core, pgn[st:fh])
			del pgn

			# Filter particle data with collected groupNumber indexing
			subgroup_number = np.append(subgroup_number, h5file[f'/PartType{pt}/SubGroupNumber'][pgn_core])
			velocity        = np.append(velocity, h5file[f'/PartType{pt}/Velocity'][pgn_core])
			coordinates     = np.append(coordinates, h5file[f'/PartType{pt}/Coordinates'][pgn_core])
			if pt == '1':
				particle_mass_DM = h5file['Header'].attrs['MassTable'][1]
				mass = np.append(mass, np.ones(len(pgn_core), dtype=np.float32) * particle_mass_DM)
			else:
				mass = np.append(mass, h5file[f'/PartType{pt}/Mass'][pgn_core])
			if pt == '0':
				temperature = np.append(temperature, h5file[f'/PartType{pt}/Temperature'][pgn_core])
				sphdensity  = np.append(sphdensity, h5file[f'/PartType{pt}/Density'][pgn_core])
				sphlength   = np.append(sphlength, h5file[f'/PartType{pt}/SmoothingLength'][pgn_core])

			del pgn_core

			# Conversion from comoving units to physical units
			velocity = comoving_velocity(header, velocity)
			coordinates = comoving_length(header, coordinates)
			mass = comoving_mass(header, mass * 1.0e10)
			if pt == '0':
				den_conv = h5file[f'/PartType{pt}/Density'].attrs['CGSConversionFactor']
				sphdensity = comoving_density(header, sphdensity * den_conv)
				sphlength = comoving_length(header, sphlength)

			# Gather the imports across cores
			data_out[f'partType{pt}'] = {}
			data_out[f'partType{pt}']['subgroupnumber'] = commune(subgroup_number)
			data_out[f'partType{pt}']['velocity']        = commune(velocity.reshape(-1, 1)).reshape(-1, 3)
			data_out[f'partType{pt}']['coordinates']     = commune(coordinates.reshape(-1, 1)).reshape(-1, 3)
			data_out[f'partType{pt}']['mass']            = commune(mass)
			if pt == '0':
				data_out[f'partType{pt}']['temperature'] = commune(temperature)
				data_out[f'partType{pt}']['sphdensity']  = commune(sphdensity)
				data_out[f'partType{pt}']['sphlength']   = commune(sphlength)

			del subgroup_number, velocity, mass, coordinates, temperature, sphdensity, sphlength

	return data_out

def cluster_data(clusterID: int,
                 header: Dict[str, float] = None,
                 fofgroups: Dict[str, np.ndarray] = None):
	"""

	:param clusterID:
	:param header:
	:param fofgroups:
	:param groupNumbers:
	:return:
	"""

	group_data  = fof_group(clusterID, fofgroups = fofgroups)
	halo_partgn = cluster_partgroupnumbers(fofgroup=group_data)
	part_data   = cluster_particles(fofgroup=group_data, groupNumbers= halo_partgn)

	out = {}
	out['Header'] = {**header}
	out['FOF'] = {**group_data}
	for pt in ['0', '1', '4']:
		out[f'partType{pt}'] = {**part_data[f'partType{pt}']}
	return out

def glance_cluster(cluster_dict: dict, verbose: bool = False, indent: int = 1) -> None:
	"""

	:param cluster_dict:
	:param verbose:
	:param indent:
	:return:
	"""
	if not verbose:
		for key, value in cluster_dict.items():
			if isinstance(value, dict):
				pprint('\t'*indent + str(key))
				glance_cluster(value, indent=indent+1)
			elif (isinstance(value, np.ndarray) or isinstance(value, list)) and len(value) > 10:
				pprint('\t' * indent + str(key) + ' : ' + f"len({len(value):d})\t val({value[0]} ... {value[-1]})")
			else:
				pprint('\t' * indent + str(key) + ' : ' + str(value))

	if verbose:
		for key, value in cluster_dict.items():
			if isinstance(value, dict):
				pprint('\t'*indent + str(key))
				glance_cluster(value, indent=indent+1)
			else:
				pprint('\t'*indent + str(key) +' : '+ str(value))


'''
Michel Lam 2015
Loads a parameter file
'''
import decimal
import re
import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as unit

numre = re.compile('(\d+[.]\d+D[+]\d+)|(-?\d+[.]\d+)')
flagre = re.compile('-[a-zA-Z]')

c = 2.9979e8
PC_TO_M = 3.086e16
MAS_TO_RAD = np.pi/(180*60*60*1000)
YR_TO_S = 3.154e7

DECIMAL = decimal.Decimal

class Parameter(object):
    def __init__(self, name, value=None, fit=None, error=None, flag=None,
                 flagvalue=None, numwrap=float, usedecimal=False):
        if usedecimal:
            self.numwrap = DECIMAL
        else:
            self.numwrap = numwrap
        # Initialize all values just in case
        self.name = name
        self.value = value
        self.fit = fit
        self.error = error
        self.flag = flag
        self.flagvalue = flagvalue
        if value is None: #parse all arguments
            self.parstring = name
            if name[0] == "#":# this is a comment
                self.name = "#"
                self.value = name[1:]
                return
            splitstring = self.parstring.strip().split()
            if len(splitstring) == 0:
                return None #?
            self.name = splitstring[0]
            if self.name == "C": #this is also a comment:
                self.value = self.parstring[2:]
                return

            if len(splitstring) == 1: # sometimes true in PSRFITS PSREPHEM table
                return
            elif flagre.match(splitstring[1][:2]) and len(splitstring) >= 4: #flag present
                self.flag = splitstring[1]
                self.flagvalue = splitstring[2]
                self.value = self.numwrap(splitstring[3])
                if len(splitstring) >= 5:
                    self.error = self.numwrap(splitstring[-1])
                    if len(splitstring) == 6:
                        self.fit = int(splitstring[4])
            else: #no flag present
                if numre.match(splitstring[1]):
                    self.value = self.numwrap(splitstring[1].replace('D', 'e'))
                elif splitstring[1].isdigit():
                    self.value = int(splitstring[1])
                elif splitstring[1][1:].isdigit() and (splitstring[1][0] == "+" or splitstring[1][0] == "-"):
                    self.value = int(splitstring[1])
                else:
                    self.value = splitstring[1]

                if len(splitstring) == 3 or len(splitstring) == 4:
                    if numre.match(splitstring[-1]):
                        self.error = self.numwrap(splitstring[-1].replace('D', 'e'))
                    elif splitstring[1].isdigit():
                        if splitstring[-1] == "NaN":
                            self.error = np.nan
                        else:
                            self.error = int(splitstring[-1])
                    elif splitstring[1][1:].isdigit() and (splitstring[1][0] == "+" or splitstring[1][0] == "-"):
                        self.error = int(splitstring[-1])
                    else:
                        self.error = splitstring[-1]
                    # Fit flag
                    if len(splitstring) == 3 and not numre.match(splitstring[2]) and splitstring[2] != "NaN":
                        if splitstring[2].isdigit():
                            self.fit = int(splitstring[2])
                        else: #this is actually NAME VAL ERROR rather than NAME VAL FLAG
                            self.fit = 0
                            self.error = float(splitstring[2])
                    elif len(splitstring) == 4:
                        self.fit = int(splitstring[2])

    def getName(self):
        return self.name

    def getValue(self):
        return self.value

    def getFit(self):
        return self.fit

    def getError(self):
        return self.error

    def getFlag(self):
        return self.flag

    def getFlagValue(self):
        return self.flagvalue

    def hasFlag(self):
        if self.flagvalue is not None:
            return True
        return False

class Par(object):
    def __init__(self, filename, numwrap=float, usedecimal=False):
        self.filename = filename
        if isinstance(filename, (list, np.ndarray)):
            lines = filename
        elif isinstance(filename, (str, np.str_)):
            FILE = open(filename, 'r')
            lines = FILE.readlines()
        else:
            return None

        self.usedecimal = usedecimal
        if self.usedecimal:
            self.numwrap = DECIMAL
        else:
            self.numwrap = numwrap
        self.paramlist = list() #each unique parameter
        self.paramnames = list() #the names of each parameter
        for line in lines:
            if len(line) == 0:
                continue
            p = Parameter(line, numwrap=self.numwrap)
            self.paramlist.append(p)
            self.paramnames.append(p.getName())
        self.paramnames = np.array(self.paramnames, dtype=np.str_)
        if isinstance(filename, str):
            FILE.close()

    def __repr__(self):
        numwrapstr = repr(self.numwrap).split("'")[1]
        return "Par(%r, numwrap=%s, usedecimal=%r)" % (self.filename, numwrapstr, self.usedecimal)

    def __str__(self):
        if isinstance(self.filename, (list, np.ndarray)):
            return "\n".join(self.filename)
        return self.filename

    def save(self, filename):
        # Crude saving attempt
        output = ""
        for param in self.paramlist:
            output += param.parstring
        with open(filename, 'w') as FILE:
            FILE.write(output)

    def getInd(self, tag):
        return np.where(self.paramnames == tag)[0]#[0]

    def get(self, tag, flag=False, error=False):
        if tag in self.paramnames:
            ind = self.getInd(tag)
            retval = []
            for i in ind:
                if error:
                    val = self.paramlist[i].getError()
                elif flag:
                    val = self.paramlist[i].getFlagValue()
                else:
                    val = self.paramlist[i].getValue()
                try:
                    retval.append(self.numwrap(val))
                except (ValueError, TypeError, decimal.InvalidOperation):
                    retval.append(val)
            if len(retval) == 1:
                return retval[0]
            return np.array(retval)
        return None

    """
    def getParameterFlags(self, tag):
        #If a parameter has flags, return all of the possible values as a list
        if tag in self.paramnames:
            ind = self.getInd(tag)
            retval = []
            for i in ind:
                retval.append(self.paramlist[i].getFlagValue())
            if len(retval) == 1:
                return retval[0]
            return np.array(retval)
        return None
    """

    def getPeriod(self):
        if 'P0' in self.paramnames:
            return self.get('P0')
        if 'F0' in self.paramnames:
            F0 = self.get('F0')
        elif 'F' in self.paramnames:
            F0 = self.get('F')
        elif 'IF0' in self.paramnames:
            F0 = (self.get('IF0') + self.get('FF0'))/self.numwrap(1000.0)
        return self.numwrap(1.0)/F0

    def getPeriodDot(self, shklovskii=False):
        if 'P1' in self.paramnames:
            Pdot = self.get('P1')
        elif 'F1' in self.paramnames:
            Pdot = self.numwrap(-1.0)*self.get('F1') / (self.get('F0')**2)
        else:
            return None

        if shklovskii: #Correct for the shklovskii effect
            PM = self.getPM()
            if PM is None or "PX" not in self.paramnames:
                return Pdot

            P = self.getPeriod() #s
            PX = self.get("PX") #mas
            PXerr = self.get("PX", error=True)
            if PX <= 0 or (PXerr is not None and PXerr >= np.abs(PX)):
                return Pdot

            PM = PM * self.numwrap(MAS_TO_RAD/YR_TO_S) #mas/yr -> rad/s
            D = self.numwrap(1000*PC_TO_M)/PX #kpc -> m
            Pdot_pm = P*PM**2 *D/self.numwrap(c) #s/s
            return Pdot-Pdot_pm
        else:
            return Pdot

    def getFrequency(self):
        return self.numwrap(1.0)/self.getPeriod()

    def getFrequencyDot(self, shklovskii=False):
        return self.numwrap(-1.0) * self.getPeriodDot(shklovskii=shklovskii) / self.getPeriod()**2

    def getPM(self, error=False):
        keys = self.paramnames
        retval = None
        if "PMRA" in keys and "PMDEC" in keys:
            PMRA = self.get("PMRA")
            PMDEC = self.get("PMDEC")
            PM = np.sqrt(PMRA**2 + PMDEC**2) #mas/yr
            if error:
                PMRAerr = self.get("PMRA", error=True)
                PMDECerr = self.get("PMDEC", error=True)
                retval = np.sqrt((PMRAerr*PMRA/PM)**2 + (PMDECerr*PMDEC/PM)**2)
            else:
                retval = PM
        elif "PMRA" in keys:
            retval = abs(self.get("PMRA", error=error))
        elif "PMDEC" in keys:
            retval = abs(self.get("PMDEC", error=error))
        elif "PMLAMBDA" in keys and "PMBETA" in keys:
            PMLAMBDA = self.get("PMLAMBDA")
            PMBETA = self.get("PMBETA")
            PM = np.sqrt(PMLAMBDA**2 + PMBETA**2) #mas/yr
            if error:
                PMLAMBDAerr = self.get("PMLAMBDA", error=True)
                PMBETAerr = self.get("PMBETA", error=True)
                retval = np.sqrt((PMLAMBDAerr*PMLAMBDA/PM)**2 + (PMBETAerr*PMBETA/PM)**2)
            else:
                retval = PM
        elif "PMLAMBDA" in keys:
            retval = abs(self.get("PMLAMBDA", error=error))
        elif "PMBETA" in keys:
            retval = abs(self.get("PMBETA", error=error))
        return retval

    def getPX(self, error=False):
        return self.get('PX', error=error)

    def getDIST(self, error=False):
        PX = self.getPX()
        if error:
            PXerr = self.getPX(error=True)
            return PXerr/PX**2
        else:
            return self.numwrap(1.0)/PX

    def getVpperp(self, error=False):
        '''
        Get transverse velocity
        v = 4.74 km/s (D/kpc) (mu/ mas yr^-1)

        for errors: assuming PM and PX are
        '''
        PM = self.getPM() #mas yr^-1
        DIST = self.getDIST() #kpc
        retval = self.numwrap(4.74) * PM * DIST

        if error:
            PMerr = self.getPM(error=True)
            DISTerr = self.getDIST(error=True)
            return retval * np.sqrt((PMerr/PM)**2 + (DISTerr/DIST)**2)
        else:
            return retval

    def getCoord(self):
        '''
        Return an Astropy SkyCoord object based on the coordinates
        '''
        keys = self.paramnames
        if "RAJ" in keys: #return Equatorial coordinates
            RAJ = self.get("RAJ")
            DECJ = self.get("DECJ")
            if DECJ[0].isdigit(): #is positive
                DECJ = "+" + DECJ
            coord = SkyCoord(ra=RAJ, dec=DECJ, frame="icrs", unit=(unit.hourangle, unit.deg))
        elif "LAMBDA" in keys: #return Ecliptic coordinates
            LAMBDA = self.get("LAMBDA")
            BETA = self.get("BETA")
            coord = SkyCoord(lon=LAMBDA, lat=BETA, frame="barycentrictrueecliptic", unit=unit.deg)
        elif "ELAT" in keys: #return Ecliptic coordinates
            LAMBDA = self.get("ELONG")
            BETA = self.get("ELAT")
            coord = SkyCoord(lon=LAMBDA, lat=BETA, frame="barycentrictrueecliptic", unit=unit.deg)
        return coord
    getCoords = getCoord

    def getDM(self):
        return self.get('DM')

    def getDMX(self, full_output=False):
        keys = self.paramnames
        Ncomponents = 0
        for key in keys:
            if key[0:4] == 'DMX_':
                Ncomponents += 1
        if Ncomponents == 0:
            return None
        #DM = self.getDM()
        ts = np.zeros(Ncomponents)
        dmxs = np.zeros(Ncomponents)
        errs = np.zeros(Ncomponents)
        if full_output:
            R1s = np.zeros(Ncomponents)
            R2s = np.zeros(Ncomponents)
            F1s = np.zeros(Ncomponents)
            F2s = np.zeros(Ncomponents)
        for i in range(Ncomponents):
            ts[i] = self.get('DMXEP_%04i'%(i+1))
            if np.isnan(ts[i]):
                ts[i] = self.numwrap(0.5)*(self.get('DMXR1_%04i'%(i+1))+self.get('DMXR2_%04i'%(i+1)))
            dmxs[i] = self.get('DMX_%04i'%(i+1))
            errs[i] = self.get('DMX_%04i'%(i+1), error=True) #check to make sure this exists?
            if full_output:
                R1s[i] = self.get('DMXR1_%04i'%(i+1))
                R2s[i] = self.get('DMXR2_%04i'%(i+1))
                F1s[i] = self.get('DMXF1_%04i'%(i+1))
                F2s[i] = self.get('DMXF2_%04i'%(i+1))
                if np.isnan(ts[i]):
                    ts[i] = (R1s[i]+R2s[i])/2.0
        if full_output:
            return ts, dmxs, errs, R1s, R2s, F1s, F2s
        return ts, dmxs, errs

    def getXMX(self):#, full_output=False):
        keys = self.paramnames
        Ncomponents = 0
        for key in keys:
            if key[0:4] == 'XMX_':
                Ncomponents += 1
        if Ncomponents == 0:
            return None
        #DM = self.getDM()
        xmxs = np.zeros(Ncomponents)
        errs = np.zeros(Ncomponents)
        R1s = np.zeros(Ncomponents)
        R2s = np.zeros(Ncomponents)
        EXPs = np.zeros(Ncomponents)
        for i in range(Ncomponents):
            xmxs[i] = self.get('XMX_%04i'%(i+1))
            errs[i] = self.get('XMX_%04i'%(i+1), error=True) #check to make sure this exists?
            R1s[i] = self.get('XMXR1_%04i'%(i+1))
            R2s[i] = self.get('XMXR2_%04i'%(i+1))
            EXPs[i] = self.get('XMXEXP_%04i'%(i+1))
        return xmxs, errs, R1s, R2s, EXPs

    def getDMseries(self):
        ts, dmxs, errs = self.getDMX()
        DM = self.getDM()
        return ts, dmxs + float(DM), errs #float from possible decimal!

    def getFD(self):
        coeffs = []
        for param in self.paramnames:
            if "FD" in param:
                coeffs.append(self.get(param))
        if len(coeffs) == 0:
            return None
        return np.array(coeffs)

    def getFDfunc(self):
        """
        Returns a function that provides the timing delays as a function of observing frequency
        """
        FD = self.getFD()
        if FD is None:
            return None
        FD = FD[::-1]
        FD = np.concatenate((FD, [0]))
        f = lambda nu: 1e6*np.polyval(FD, np.log(nu)) #nu in GHz, returns values in microseconds
        return f

    def getName(self):
        """ Returns the pulsar name """
        if "PSR" in self.paramnames:
            return self.get("PSR")
        elif "PSRJ" in self.paramnames:
            return self.get("PSRJ")
        return None

    def getTspan(self, years=False):
        """ Returns the total time span the par file covers """
        start = self.get("START")
        finish = self.get("FINISH")
        if years:
            return (finish-start)/365.25
        return finish-start

import numpy,import astropy,from astropy


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

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from builtins import object

import numpy as np
import pkg_resources
import fsps
import os
import linecache
#grid: 2 files: line, cont
#columns: wavelengths
#rows: models
#header
#wavelength grid
#age, Z, logU
#flux1
#flux2
#flux3

###
# reads from ZAU***.out_lines, ZAU***.out_cont
# produces ZAU_**.lines, ZAU_**.cont
###
class writeFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        # load each model's parameters from prefix.pars
        self.loadModInfo()
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        # interp and print neb cont onto FSPS wavelenth arr
        self.doContOut()
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            ddata[key] = data[i]
            self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.logZ))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}\n".format(self.logZ[n-1], self.Age[n-1], self.logU[n-1])
        f.write(tstr)
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        efrac=self.efrac[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["logZ"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return

#------------------------------------------------
def zmet_to_nuZ(zmet):
    '''
    zmet = 1,2,3,4,5,6,7,8
        ...corresponding to...
    m1.00, m1.02, m1.04, m1.06, p0.00, p0.02, p0.04, p0.06
    '''
    zmets = np.array([1,2,3,4,5,6,7,8])
    zs = np.array([-1.00, -1.02, -1.04, -1.06, 0.00, 0.02, 0.04, 0.06])
    if type(zmet) is float:
        ii, = np.where(zmets == zmet)[0]
        out_val = zs[ii]
    else:
        out_val = np.zeros_like(zmet, dtype='float')
        for i in range(len(zmet)):
            ii, = np.where(zmets == zmet[i])[0]
            out_val[i] = zs[ii]
    return out_val

class writeAltFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        self.loadModInfo() # load each model's parameters from prefix.pars
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        self.doContOut() # interp and print neb cont onto FSPS wavelenth arr
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac", "zmet"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            if key == 'zmet':
                temp_arr = zmet_to_nuZ(data[i])
                ddata[key] = temp_arr
                self.__setattr__(key, temp_arr)
            else:
                ddata[key] = data[i]
                self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.zmet))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
            f.write(tstr)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(self.zmet[n-1], self.Age[n-1], self.logU[n-1])
            f.write(tstr+"\n")
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        zmet=self.zmet[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["zmet"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return