'''
https://github.com/jeremite/channel-attribution-model/blob/master/process_data.py

'''

import numpy as np
import random
import tensorflow as tf
import tf_slim as slim
# from keras.utils import to_categorical
from keras.utils import np_utils
from keras.preprocessing.text import Tokenizer
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from pandas.api.types import is_string_dtype, is_numeric_dtype
import warnings
import pandas as pd

warnings.filterwarnings("ignore")

def cat(data, field, leng, s=','):
    data[field] = data[field].map(lambda x: x.split(s))
    data[leng] = data[field].map(lambda x: len(x))

def most_fre(lt):
    return max(lt, key=lt.count)

# most_fre('north')

def sca(col):
    col = [int(v) for v in col]
    return [(v - min(col)) / (max(col) - min(col)) for v in col]

def train_cats(df):
    """Change any columns of strings in a panda's dataframe to a column of
    categorical values. This applies the changes inplace.
    """
    for n, c in df.items():
        if is_string_dtype(c):
            df[n] = c.astype('category').cat.as_ordered()

def process_data(data, seq_length=20):
    # data_a = data_att.loc[:,['path_id','path','total_conversions','last_time_lapse','null_conversion']]
    data.dropna(axis=0, how='any',inplace=True)  # 删除了所有na的,axis=0删除包含缺失值的行，how='any'只要有缺失值出现就删除该行，inplace=True在原数据上进行操作
    # merge target variables

    # str -> list
    cat(data, 'path', "leng_path", s='>')
    cat(data, 'marketing_area', "leng_area", s=',')
    cat(data, 'tier', "leng_tier", s=',')
    cat(data, 'customer_type', "leng_type", s=',')
    cat(data, 'age', "leng_type", s=',')
    cat(data, 'gender', "leng_type", s=',')
    cat(data, 'country', "leng_type", s=',')

    # remove those path with channels less than 3

    data_new = data[(data.leng_path >= 3)]
    # 可以还原索引，重新变为默认的整型索引
    data_new = data_new.reset_index()
    # leave with the most common value
    data_new.marketing_area = data_new.marketing_area.map(lambda x: most_fre(x))
    data_new.tier = data_new.tier.map(lambda x: most_fre(x))
    data_new.customer_type = data_new.customer_type.map(lambda x: most_fre(x))
    data_new.age = data_new.age.map(lambda x: most_fre(x))
    data_new.gender = data_new.gender.map(lambda x: most_fre(x))
    data_new.country = data_new.country.map(lambda x: most_fre(x))
    data_new.replace('', 'NA', inplace=True)  # ？？？？？？？？前面不是删掉了缺失行吗

    # 输出12：所有y
    y = data_new.total_conversions

    # 抽取训练集、测试集
    # got train and test data indices
    idx = [x for x in range(data_new.shape[0])]  # 输出矩阵的行数,以[0, 1, 2, 3]形式输出,共17行，0~16
    np.random.seed(111)
    random.shuffle(idx)  # 打乱
    # print(idx)
    tr_idx = idx[0:int(0.8 * len(idx))]
    te_idx = idx[int(0.8 * len(idx)):]
    # print(te_idx)

    # got data for time decay
    cat(data_new, 'last_time_lapse', "leng_time_lapse", s=',')
    data_new.last_time_lapse = data_new.last_time_lapse.map(lambda x: sca(x))
    pad_sequence = tf.keras.preprocessing.sequence.pad_sequences
    time_decay = pad_sequence(data_new.last_time_lapse, maxlen=seq_length, padding='pre', truncating='pre',
                              dtype='float32')  # 17*20
    # 输出10：所有时间
    time_decay = time_decay.reshape(-1, 20, 1)  # 17*20*1
    # 输出1：训练集时间
    time_decay_tr = time_decay[tr_idx]  # 占0.8
    # 输出2：测试集时间
    time_decay_te = time_decay[te_idx]  # 占0.2

    # got data for attribution
    text = data_new.path  # (17,) object
    # encoding
    t = Tokenizer()  # Tokenizer实际上只是生成了一个字典，并且统计了词频等信息，并没有把文本转成需要的向量表示。
    # 首先用Tokenizer的 fit_on_texts 方法学习出文本的字典，然后word_index 就是对应的单词和数字的映射关系dict，
    # 通过这个dict可以将每个string的每个词转成数字，可以用texts_to_sequences，这是我们需要的，然后通过padding的方法补成同样长度，
    # 在用keras中自带的embedding层进行一个向量化，并输入到LSTM中。
    t.fit_on_texts(text)
    # vocabulary size
    vocab_size = len(t.word_index) + 3  # 8
    print(vocab_size)
    # integer encode the documents
    encoded_docs = t.texts_to_sequences(text)
    # padding and truncating path data
    # 输出11：newlines.shape,17*20,pad_sequence
    newlines = pad_sequence(encoded_docs, maxlen=seq_length, padding='pre', truncating='post')  # 17*20
    X_train = newlines[tr_idx]
    # 输出7
    Y_train = y[tr_idx]
    X_test = newlines[te_idx]
    # 输出8
    Y_test = y[te_idx]
    # 输出9：所有x
    all_X = np.array(list(map(lambda x: np_utils.to_categorical(x, num_classes=vocab_size), newlines)), ndmin=3)
    # 输出3：用户路径模块-训练集X
    X_tr = np.array(list(map(lambda x: np_utils.to_categorical(x, num_classes=vocab_size), X_train)), ndmin=3)
    # 输出4：用户路径模块-测试集X
    X_te = np.array(list(map(lambda x: np_utils.to_categorical(x, num_classes=vocab_size), X_test)), ndmin=3)
    # 输出14：所有路径
    paths = text[tr_idx].reset_index().path

    # got customer data (control data)
    data_lr = data_new.loc[:, ['marketing_area', 'tier', 'customer_type', 'age', 'gender', 'country']]
    train_cats(data_lr)
    data_lr['c_type' + '_na'] = [1 if v == 'NA' else 0 for v in data_lr['customer_type']]
    for col in data_lr.columns:
        if not is_numeric_dtype(data_lr[col]):
            data_lr[col] = data_lr[col].cat.codes + 1
    # 输出5：用户属性模块：训练集X
    X_tr_lr = data_lr.iloc[tr_idx, :]
    # print(X_tr_lr)
    # 输出6：用户属性模块：测试集X
    X_te_lr = data_lr.iloc[te_idx, :]  # iloc根据行号选
    # 利用 MinMaxScaler 放缩数据
    scaler = MinMaxScaler()
    scaler.fit(X_tr_lr)
    X_tr_lr[['marketing_area', 'tier', 'customer_type', 'age', 'gender', 'country']] = scaler.fit_transform(
        X_tr_lr[['marketing_area', 'tier', 'customer_type', 'age', 'gender', 'country']])
    X_te_lr[['marketing_area', 'tier', 'customer_type', 'age', 'gender', 'country']] = scaler.transform(
        X_te_lr[['marketing_area', 'tier', 'customer_type', 'age', 'gender', 'country']])
    # 输出13：用户属性模块的哪些字段
    categorical_vars = data_lr.columns[0:6]
    # print(categorical_vars)

    return [time_decay_tr, time_decay_te, X_tr, X_te, X_tr_lr, X_te_lr, Y_train, Y_test,
            all_X, time_decay, newlines, y, categorical_vars, paths]

# data = pd.read_excel('fake_data.xlsx')
# data['last_time_lapse'] = data['last_time_lapse'].map(eval)
if __name__ == '__main__':
    data = pd.read_excel('fake_data2.xlsx')
    seq_length = 20
    # 预处理后的数据
    data_all = process_data(data, seq_length=seq_length)
    # print('data_all[-1]:')
    # print(data_all[-1])
    channels = list(set([_p for p in data_all[-1] for _p in p]))
    # channels.sort()
    print(channels)


##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2019, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
#
# This software is distributed under the 3-clause BSD License.
##############################################################################
"""
A simple GUI viewer/editor for Pyomo models.
"""
from __future__ import division, print_function, absolute_import

__author__ = "John Eslick"

import os
import logging

_log = logging.getLogger(__name__)

from pyomo.contrib.viewer.qt import *
from pyomo.contrib.viewer.report import value_no_exception, get_residual
import pyomo.environ as pyo

mypath = os.path.dirname(__file__)
try:
    _ResidualTableUI, _ResidualTable = \
        uic.loadUiType(os.path.join(mypath, "residual_table.ui"))
except:
    class _ResidualTableUI(object):
        pass
    class _ResidualTable(object):
        pass

class ResidualTable(_ResidualTable, _ResidualTableUI):
    def __init__(self, ui_data, parent=None):
        super(ResidualTable, self).__init__(parent=parent)
        self.setupUi(self)
        self.ui_data = ui_data
        datmodel = ResidualDataModel(parent=self, ui_data=ui_data)
        self.datmodel = datmodel
        self.tableView.setModel(datmodel)
        self.ui_data.updated.connect(self.refresh)
        self.sortButton.clicked.connect(self.sort)
        self.calculateButton.clicked.connect(self.calculate)

    def sort(self):
        self.datmodel.sort()

    def refresh(self):
        self.datmodel.update_model()
        self.datmodel.sort()
        self.datmodel.layoutChanged.emit()

    def calculate(self):
        self.ui_data.calculate_constraints()
        self.refresh()

class ResidualDataModel(QAbstractTableModel):
    def __init__(self, parent, ui_data):
        super(ResidualDataModel, self).__init__(parent)
        self.column = ["name", "residual", "value", "ub", "lb", "active"]
        self.ui_data = ui_data
        self.include_inactive = True
        self.update_model()
        self.sort()

    def update_model(self):
        if self.include_inactive:
            ac = None
        else:
            ac = True
        self._items = list(self.ui_data.model.component_data_objects(
            pyo.Constraint, active=ac))

    def sort(self):
        self._items.sort(key=
            lambda o: (o is None, get_residual(self.ui_data, o)
                                  if get_residual(self.ui_data, o) is not None
                                  else -float("inf")), reverse=True)

    def rowCount(self, parent=QtCore.QModelIndex()):
        return len(self._items)

    def columnCount(self, parent=QtCore.QModelIndex()):
        return len(self.column)

    def data(self, index=QtCore.QModelIndex(), role=QtCore.Qt.DisplayRole):
        row = index.row()
        col = self.column[index.column()]
        if role == QtCore.Qt.DisplayRole:
            o = self._items[row]
            if col=="name":
                return str(o)
            elif col=="residual":
                return get_residual(self.ui_data, o)
            elif col=="active":
                return o.active
            elif col=="ub":
                return value_no_exception(o.upper)
            elif col=="lb":
                return value_no_exception(o.lower)
            elif col=="value":
                try:
                    return self.ui_data.value_cache[o]
                except KeyError:
                    return None
        else:
            return None

    def headerData(self, i, orientation, role=QtCore.Qt.DisplayRole):
        '''
            Return the column headings for the horizontal header and
            index numbers for the vertical header.
        '''
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return self.column[i]
        return None


#!/usr/bin/env python3
import functools
import os
from multiprocessing import Process
from baselines import logger
from mpi4py import MPI
import mpi_util
import tf_util
from cmd_util import make_atari_env, arg_parser
from policies.cnn_gru_policy_dynamics import CnnGruPolicy
from policies.cnn_policy_param_matched import CnnPolicy
from ppo_agent import PpoAgent
from utils import set_global_seeds
from vec_env import VecFrameStack
from copy import deepcopy
import datetime
from baselines.common.vec_env import VecNormalize
from baselines.common.cmd_util import make_vec_env

def make_mujoco_env(env_id, seed, reward_scale=1.0):
    env = make_vec_env(env_id, "mujoco", num_env=1, seed=seed, reward_scale=reward_scale,
                       flatten_dict_observations=True)
    env = VecNormalize(env)
    return env

def train(*, origin_paper, env_id, num_env, hps, num_timesteps, seed):
    # venv = VecFrameStack(
    #     make_atari_env(env_id, num_env, seed, wrapper_kwargs=dict(),
    #                    start_index=num_env * MPI.COMM_WORLD.Get_rank(),
    #                    max_episode_steps=hps.pop('max_episode_steps')),
    #     hps.pop('frame_stack'))
    # # venv.score_multiple = {'Mario': 500,
    # #                        'MontezumaRevengeNoFrameskip-v4': 100,
    # #                        'GravitarNoFrameskip-v4': 250,
    # #                        'PrivateEyeNoFrameskip-v4': 500,
    # #                        'SolarisNoFrameskip-v4': None,
    # #                        'VentureNoFrameskip-v4': 200,
    # #                        'PitfallNoFrameskip-v4': 100,
    # #                        }[env_id]
    # venv.score_multiple = 1
    # venv.record_obs = True if env_id == 'SolarisNoFrameskip-v4' else False
    venv = make_mujoco_env(env_id, seed, reward_scale=1.0)
    setattr(venv, "score_multiple", 1.0)
    setattr(venv, "record_obs", False)
    ob_space = venv.observation_space
    ac_space = venv.action_space
    gamma = hps.pop('gamma')
    policy = {'rnn': CnnGruPolicy,
              'cnn': CnnPolicy}[hps.pop('policy')]

    agent = PpoAgent(
        scope='ppo',
        ob_space=ob_space,
        ac_space=ac_space,
        stochpol_fn=functools.partial(
                policy,
                scope='pol',
                ob_space=ob_space,
                ac_space=ac_space,
                update_ob_stats_independently_per_gpu=hps.pop('update_ob_stats_independently_per_gpu'),
                proportion_of_exp_used_for_predictor_update=hps.pop('proportion_of_exp_used_for_predictor_update'),
                dynamics_bonus = hps.pop("dynamics_bonus")
            ),
        gamma=gamma,
        gamma_ext=hps.pop('gamma_ext'),
        lam=hps.pop('lam'),
        nepochs=hps.pop('nepochs'),
        nminibatches=hps.pop('nminibatches'),
        lr=hps.pop('lr'),
        cliprange=hps.pop('cliprange'),
        nsteps=hps.pop('nsteps'),
        ent_coef=hps.pop('ent_coef'),
        max_grad_norm=hps.pop('max_grad_norm'),
        use_news=hps.pop("use_news"),
        comm=MPI.COMM_WORLD if MPI.COMM_WORLD.Get_size() > 1 else None,
        update_ob_stats_every_step=hps.pop('update_ob_stats_every_step'),
        int_coeff=hps.pop('int_coeff'),
        ext_coeff=hps.pop('ext_coeff'),
        origin_paper=origin_paper,
    )
    agent.start_interaction([venv])
    if hps.pop('update_ob_stats_from_random_agent'):
        agent.collect_random_statistics(num_timesteps=128*10)
    if len(hps) != 0:
        logger.info("warning!Unused hyperparameters: %s" % list(hps.keys()))

    counter = 0
    while True:
        info = agent.step()
        if info['update']:
            logger.logkvs(info['update'])
            logger.dumpkvs()
            counter += 1
        if agent.I.stats['tcount'] > num_timesteps:
            break

    agent.stop_interaction()

def add_env_params(parser):
    parser.add_argument('--env', help='environment ID', default='HalfCheetah-v2')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--max_episode_steps', type=int, default=4500)

def main(args):
    server = args.server
    origin_paper = args.origin_paper
    if not origin_paper:
        args.nminibatches = 32
        args.nsteps = 2048
        args.max_grad_norm = 0.5
        args.lr = 3e-4
        args.ent_coef = 0
        args.cliprange = 0.2

    if not server:
        path = "logs"
    else:
        path = os.path.join("/home", "lizn", "openai", "rnd",
                            datetime.datetime.now().strftime("ppo-rnd-{}-%Y-%m-%d-%H-%M-%S-%f".format(args.env)))
    logger.configure(dir=path, format_strs=['stdout', 'log', 'csv'] if MPI.COMM_WORLD.Get_rank() == 0 else [])
    if MPI.COMM_WORLD.Get_rank() == 0:
        with open(os.path.join(logger.get_dir(), 'experiment_tag.txt'), 'w') as f:
            f.write(args.tag)

    if server:
        mpi_util.setup_mpi_gpus()

    seed = 10000 * args.seed + MPI.COMM_WORLD.Get_rank()
    set_global_seeds(seed)

    hps = dict(
        frame_stack=4,
        ent_coef=args.ent_coef,
        nminibatches=args.nminibatches,
        nepochs=4 if origin_paper else 10,
        lr=args.lr,
        max_grad_norm=args.max_grad_norm,
        use_news=args.use_news,
        gamma=args.gamma,
        gamma_ext=args.gamma_ext,
        max_episode_steps=args.max_episode_steps,
        lam=args.lam,
        update_ob_stats_every_step=args.update_ob_stats_every_step,
        update_ob_stats_independently_per_gpu=args.update_ob_stats_independently_per_gpu,
        update_ob_stats_from_random_agent=args.update_ob_stats_from_random_agent,
        proportion_of_exp_used_for_predictor_update=args.proportion_of_exp_used_for_predictor_update,
        policy=args.policy,
        int_coeff=args.int_coeff,
        ext_coeff=args.ext_coeff,
        dynamics_bonus=args.dynamics_bonus,
        cliprange=args.cliprange,
        nsteps=args.nsteps

    )
    logger.info("training parameters:")
    logger.info(str(hps))
    tf_util.make_session(make_default=True)
    train(origin_paper=origin_paper, env_id=args.env, num_env=args.num_env, seed=seed,
        num_timesteps=args.num_timesteps, hps=hps)

if __name__ == '__main__':
    from baselines.common.misc_util import boolean_flag

    parser = arg_parser()
    add_env_params(parser)
    parser.add_argument('--num-timesteps', type=int, default=int(1e6))
    parser.add_argument('--num_env', type=int, default=1)
    parser.add_argument('--use_news', type=int, default=0)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--gamma_ext', type=float, default=0.99)
    parser.add_argument('--lam', type=float, default=0.95)
    parser.add_argument('--update_ob_stats_every_step', type=int, default=0)
    parser.add_argument('--update_ob_stats_independently_per_gpu', type=int, default=0)
    parser.add_argument('--update_ob_stats_from_random_agent', type=int, default=1)
    parser.add_argument('--proportion_of_exp_used_for_predictor_update', type=float, default=1.)
    parser.add_argument('--tag', type=str, default='')
    parser.add_argument('--policy', type=str, default='cnn', choices=['cnn', 'rnn'])
    parser.add_argument('--int_coeff', type=float, default=1.)
    parser.add_argument('--ext_coeff', type=float, default=2.)
    parser.add_argument('--dynamics_bonus', type=int, default=0)
    parser.add_argument('--nminibatches', type=int, default=1)
    parser.add_argument('--nsteps', type=int, default=128)
    parser.add_argument('--max_grad_norm', type=float, default=0)
    parser.add_argument('--lr', type=float, default=0.0001)
    parser.add_argument('--ent_coef', type=float, default=0.001)
    parser.add_argument('--cliprange', type=float, default=0.1)
    boolean_flag(parser, 'origin_paper', default=False)
    boolean_flag(parser, 'server', default=False)
    boolean_flag(parser, 'all_env', default=False)
    # todo: change the following two parameters when you run your local machine or server

    args = parser.parse_args()

    list_p = []
    list_env = ["HalfCheetah-v2", "Hopper-v2", "Swimmer-v2", "InvertedDoublePendulum-v2", "Reacher-v2",
                "Walker2d-v2", "Pusher-v2"] if args.all_env else [args.env]
    for env in list_env:
        _args = deepcopy(args)
        _args.env = env
        p = Process(target=main, args=(_args, ))
        list_p.append((env, p))
        p.start()
        print("============================================")
        print("{} start!".format(env))
        print("============================================")
    for (env, p) in list_p:
        p.join()
        print("============================================")
        print("{} end!".format(env))
        print("============================================")


import json
from typing import Union

from aioredis import Redis
from app.crud.crud_redis_general import CrudRedisGeneral
from app.models import KLine

class CrudRedisKLine:

    # @staticmethod
    # async def get_key_value_pairs(redis_client: Redis) -> list:
    #     keys = []

    @staticmethod
    async def set_kline_intervals(redis_client, intervals: list):
        kline_interval_json = json.dumps(intervals)
        return await CrudRedisGeneral.set(redis_client, "kline-intervals", kline_interval_json)

    @staticmethod
    async def get_kline_intervals(redis_client) -> list:
        kline_intervals_bytes = await CrudRedisGeneral.get(redis_client, "kline-intervals")
        kline_intervals_str = kline_intervals_bytes.decode("utf-8")
        kline_intervals = json.loads(kline_intervals_str)
        return kline_intervals

    @staticmethod
    async def set_kline(redis_client, key: str, kline: dict):
        kline_json = json.dumps(kline)
        return await CrudRedisGeneral.set(redis_client, key, kline_json)

    @staticmethod
    async def get_kline(redis_client, key: str) -> dict:
        kline_bytes = await CrudRedisGeneral.get(redis_client, key)
        kline_str = kline_bytes.decode("utf-8")
        kline_dict = json.loads(kline_str)
        return kline_dict

    # @staticmethod
    # def create_kline_latest_key(interval_seconds: int) -> str:
    #     return f"kline-{interval_seconds}sec-latest"

    @staticmethod
    def create_kline_key(market: str, interval_seconds: int, start_time: Union[int, str]) -> str:
        if type(start_time) == int:
            return f"kline-{market}-{interval_seconds}-{start_time}"
        else:
            return f"kline-{market}-{interval_seconds}-latest"

    @staticmethod
    def get_market_from_kline_latest_key(kline_latest_key: str) -> str:
        market = ""
        key_parts = kline_latest_key.split("-")
        if len(key_parts) == 4:
            # interval_parts = key_parts[1].split("sec")
            # interval = interval_parts[0]
            market = key_parts[1]
        return market

    @staticmethod
    def get_interval_from_kline_latest_key(kline_latest_key: str) -> int:
        interval = 0
        key_parts = kline_latest_key.split("-")
        if len(key_parts) == 4:
            # interval_parts = key_parts[1].split("sec")
            # interval = interval_parts[0]
            interval = key_parts[2]
        return int(interval)


# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

import yaml
from api_base import PREFIX_TENSOR_NAME
from api_gen import (
    ForwardAPI,
    api_namespace,
    declare_extension_api,
    header_include,
    source_include,
)

######################
# Code Gen Templates #
######################

API_IMPL_TEMPLATE = """
PADDLE_API {} {}({}) {{
  // Kernel Key Construction{}
  // Kernel Dispatch Body{}
}}
"""
DIPATCH_END_GUARD_TEMPLATE = """
PADDLE_THROW(phi::errors::Unimplemented(
          "The kernel of ({}) for input tensors is unimplemented, please check the type of input tensors."));
"""

# TODO(chenweihang): add profile function code later
# TODO(chenweihang): add view support later
MAIN_DIST_BRANCH_TEMPLATE = """
  // Auto Parallel condition
  if ({}) {{
    // 1. Create API Output & Prepare Dist and Dense Output{}
    // 2. InferSPMD (Infer Global Shape and DistAttr of Inputs&Outputs){}
    // 3. Select Kernel{}
    // 4. Reshard Input{}
    // 5. PrepareData (DataTransform & Prepare Dist and Dense Input){}
    // 6. Infer Local DenseTensor Meta{}
    // 7. DenseTensor Kernel Call{}
    // 8. Reshard Output{}
    // 9. Return
    {}
  }}
"""

# Auto Parallel condition
AUTO_PARALLEL_COND_TEMPLATE = """AllInputsAreDistTensor({})"""

# 1. Create API Outputs
API_OUT_CREATION_TEMPLATE = """
    {} api_output{};
"""
INPLACE_API_OUT_CREATION_TEMPLATE = """
    {} api_output{{{}}};
"""
SINGLE_OUT_CREATION_TEMPLATE = """
    auto dist_out = SetKernelDistOutput(&api_output);
    auto dense_out = dist_out->mutable_value();
"""
MULTI_SINGLE_OUT_CREATION_TEMPLATE = """
    auto dist_out_{} = SetKernelDistOutput({});
    auto dense_out_{} = dist_out_{}->mutable_value();
"""

# TODO(chenweihang): support vector and tuple output later
VECTOR_OUT_CREATION_TEMPLATE = """
"""
MULTI_VECTOR_OUT_CREATION_TEMPLATE = """
    auto dist_out_{} = {}({}, {});
    auto dense_out_{} = dist_out_{}->mutable_value();
"""
TUPLE_OUT_CREATION_TEMPLATE = """
"""

# 2. InferSPMD
# Call InferMeta now, replace by InferSPMD function later
# TODO(chenweihang): InferSPMD function design
SINGLE_DIST_META_IN_TEMPLATE = """MakeMetaTensor(*{}.impl()), """
# TODO(chenweihang): support vector and optional args later
VECTOR_DIST_META_IN_TEMPLATE = """
"""
OPTIONAL_DIST_VECTOR_META_IN_TEMPLATE = """
"""
SINGLE_DIST_META_OUT_DECL_TEMPLATE = """
    phi::MetaTensor meta_{}({});"""
INFER_SPMD_TEMPLATE = """
    phi::{}({}{});
"""

# 3. Select Kernel
KERNEL_SELECTION_TEMPLATE = """
    VLOG(6) << "{} API dist branch: kernel key: [" << kernel_backend << ", " << kernel_layout << ", "<< kernel_data_type << "]";
    auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
        "{}", {{kernel_backend, kernel_layout, kernel_data_type}});
    const auto& kernel = kernel_result.kernel;
    VLOG(6) << "{} kernel: " << kernel;
    auto* dev_ctx = GetDeviceContextByBackend(kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
"""

# 4. Reshard Input
INPUT_RESHARD_TEMPLATE = """
"""

# 5. PrepareData
SINGLE_PREPARE_DATA_TEMPLATE = """
    auto dist_input_{} = PrepareDataForDistTensor({}, GetKernelInputArgDef(kernel.InputAt({}), kernel_backend), {}, kernel_result.is_stride_kernel);
    auto input_{} = dist_input_{}->mutable_value();
"""
INFER_META_SINGLE_INPUT_TEMPLATE = """
    auto dist_input_{} = {}.impl();
    auto input_{} = static_cast<phi::distributed::DistTensor*>(dist_input_{}.get())->mutable_value();
"""
INFER_META_OPTIONAL_INPUT_TEMPLATE = """
    paddle::optional<phi::TensorBase> input_{} = {} ? paddle::optional<phi::TensorBase>(*{}->impl()) : paddle::none;
"""
INFER_META_VECTOR_INPUT_TEMPLATE = """
    auto input_{}_uq_ptr = TensorToDenseTensor({});
    const auto& input_{} = *input_{}_uq_ptr;
"""

# 6. Infer Local DenseTensor Meta
SINGLE_META_IN_TEMPLATE = """MakeMetaTensor(*input_{}), """
# TODO(chenweihang): support vector and optional args later
VECTOR_META_IN_TEMPLATE = """
"""
OPTIONAL_VECTOR_META_IN_TEMPLATE = """
"""
SINGLE_META_OUT_DECL_TEMPLATE = """
    phi::MetaTensor meta_{}({});"""
INFER_META_TEMPLATE = """
    phi::{}({}{});
"""

# 7. DenseTensor Kernel Call
# TODO(chenweihang): support kernel fallback later
SINGLE_OUTPUT_NAME = """dense_out"""
# TODO(chenweihang): support vector and tuple output later
VECTOR_OUTPUT_NAME_TEMPLATE = """
"""
TUPLE_OUTPUT_NAME_TEMPLATE = """
"""
KERNEL_CALL_TEMPLATE = """
    using kernel_signature = {};
    auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
    (*kernel_fn)({}, {});
"""

# 8. Reshard Output
OUTPUT_RESHARD_TEMPLATE = """
"""

# BaseAPI members:
# inputs:
#     names : [], list of input names
#     input_info : {input_name : type}
# attrs:
#     names : [], list of attribute names
#     attr_info : { attr_name : (type, default_values)}
# outputs:
#     names : [], list of output names
#     types : [], list of output types
#     out_size_expr : [], expression for getting size of vector<Tensor>

class DistForwardAPI(ForwardAPI):
    def __init__(self, api_item_yaml):
        super().__init__(api_item_yaml)
        self.init_dist_api_members()

    def init_dist_api_members(self):
        self.gene_dist_input_func = {
            "const Tensor&": {
                "dense": self.generate_dense_input,
            },
            "const paddle::optional<Tensor>&": {
                "dense": self.generate_dense_input,
            },
        }

        self.inplace_flag = False
        self.dist_output_args = []
        self.dense_output_args = []

    def need_to_generate_code_for_inplace_impl(self, i):
        return (
            self.inplace_flag
            and self.inplace_map is not None
            and self.outputs['names'][i] in self.inplace_map
        )

    def need_to_generate_code_for_view_impl(self, i):
        return (
            not self.inplace_flag
            and self.view_map is not None
            and self.outputs['names'][i] in self.view_map
        )

    def is_inplace_output(self, i):
        return self.outputs['names'][i] in self.inplace_map

    def is_inplace_and_optional_output(self, i):
        return (
            self.outputs['names'][i] in self.inplace_map
            and self.inplace_map[self.outputs['names'][i]] in self.optional_vars
        )

    def vector_output_size_assertion_check(self):
        assert (
            self.outputs['out_size_expr'] is not None
        ), f"{self.api}: The out size expr : '{{expr}}' should be set when output has Tensor[]. You can refer 'split' api."

    def generate_if_condition_code(self) -> str:
        input_args = ""
        for input_name in self.inputs['names']:
            input_args = input_args + input_name + ", "
        if len(input_args) > 2:
            input_args = input_args[:-2]
        return AUTO_PARALLEL_COND_TEMPLATE.format(input_args)

    def generate_output_creation_code(self) -> str:
        # forward api need to generate api and kernel outputs
        output_num = len(self.outputs['types'])
        return_type = self.get_return_type_with_intermediate(self.inplace_flag)
        output_creation_code = ""
        if output_num == 1:
            # api output generate
            if self.need_to_generate_code_for_inplace_impl(0):
                inplace_assign_code = (
                    " = " + self.inplace_map[self.outputs['names'][0]]
                )
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, inplace_assign_code
                )
            else:
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, ""
                )
            # kernel output generate
            self.dist_output_args.append('dist_out')
            self.dense_output_args.append('dense_out')
            if self.outputs['types'][0] == 'Tensor':
                output_creation_code += SINGLE_OUT_CREATION_TEMPLATE
            else:
                self.vector_output_size_assertion_check()
        elif output_num > 1:
            # api output generate
            if self.inplace_flag:
                inplace_assign_code = ""
                for i, out_name in enumerate(self.outputs['names']):
                    if self.need_to_generate_code_for_inplace_impl(i):
                        inplace_assign_code += self.inplace_map[out_name] + ', '
                    else:
                        inplace_assign_code += 'Tensor(), '
                inplace_assign_code = inplace_assign_code[:-2]
                output_creation_code += (
                    INPLACE_API_OUT_CREATION_TEMPLATE.format(
                        return_type, inplace_assign_code
                    )
                )
            else:
                output_creation_code += API_OUT_CREATION_TEMPLATE.format(
                    return_type, ""
                )
            # kernel output generate
            for i, out_type in enumerate(self.outputs['types']):
                self.dist_output_args.append(f'dist_out_{i}')
                self.dense_output_args.append(f'dense_out_{i}')
                set_out_func = "SetKernelDistOutput"
                get_out_code = f"&std::get<{i}>(api_output)"
                if self.is_inplace_and_optional_output(i):
                    get_out_code = f"std::get<{i}>(api_output).get_ptr()"

                if out_type == 'std::vector<Tensor>':
                    self.vector_output_size_assertion_check(i)
                    # Special case for inplace vector and inplace optional<vector>
                    # TODO(chenweihang): support this branch later
                    if self.is_inplace_output():
                        set_out_func = "SetInplaceVectorKernelOutput"
                        if self.is_inplace_and_optional_output(i):
                            set_out_func = (
                                "SetInplaceOptionalVectorKernelOutput"
                            )
                            get_out_code = f"std::get<{i}>(api_output)"
                    output_creation_code += (
                        MULTI_VECTOR_OUT_CREATION_TEMPLATE.format(
                            i,
                            set_out_func,
                            self.outputs['out_size_expr'][i],
                            get_out_code,
                            i,
                            i,
                        )
                    )
                else:
                    output_creation_code += (
                        MULTI_SINGLE_OUT_CREATION_TEMPLATE.format(
                            i, get_out_code, i, i
                        )
                    )
        else:
            raise ValueError(
                "{} : Output error: the output should not be empty.".format(
                    self.api
                )
            )

        return output_creation_code

    def generate_infer_spmd_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        output_names = self.outputs['names']

        # 1. get infer meta func name
        infer_meta = self.infer_meta
        infer_meta_func_code = infer_meta['func']

        # 2. get meta tensor input args
        infer_meta_params = (
            infer_meta['param']
            if infer_meta['param'] is not None
            else input_names + attr_names
        )
        input_args_code = ""
        for param in infer_meta_params:
            if param in input_names:
                if self.inputs['input_info'][param] == "const Tensor&":
                    input_args_code += SINGLE_DIST_META_IN_TEMPLATE.format(
                        param
                    )
                else:
                    raise ValueError(
                        f"{self.api} : Param of infer_spmd error : {self.inputs['input_info'][param]} type is not supported."
                    )
            elif param in attr_names:
                input_args_code = input_args_code + param + ", "
            elif isinstance(param, str):
                input_args_code = input_args_code + "\"" + param + "\", "
            elif isinstance(param, bool):
                input_args_code = input_args_code + str(param).lower() + ", "
            else:
                input_args_code = input_args_code + str(param) + ", "

        # 3. get meta tensor output args
        output_decl_code = ""
        output_args_code = ""
        for i, out_name in enumerate(self.dist_output_args):
            if self.outputs['types'][i] == 'std::vector<Tensor>':
                # TODO(chenweihang): support vector output later
                pass
            else:
                output_decl_code += SINGLE_DIST_META_OUT_DECL_TEMPLATE.format(
                    out_name, out_name
                )
                if len(self.dense_output_args) == 1:
                    output_args_code += f"&meta_{out_name}, "
                else:
                    output_args_code += (
                        f"{out_name} ? &meta_{out_name} : nullptr, "
                    )
        output_args_code = output_args_code[:-2]

        return output_decl_code + INFER_SPMD_TEMPLATE.format(
            infer_meta_func_code, input_args_code, output_args_code
        )

    def generate_kernel_selection_code(self) -> str:
        return KERNEL_SELECTION_TEMPLATE.format(
            self.api, self.kernel['func'][0], self.kernel['func'][0]
        )

    def generate_reshard_input_code(self) -> str:
        return INPUT_RESHARD_TEMPLATE.format()

    # override BaseAPI's method
    def generate_dense_input(
        self,
        input_name,
    ):
        input_tensor_code = ""
        trans_flag = self.gene_trans_flag(input_name)
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        kernel_param = self.kernel['param']
        if kernel_param is None:
            kernel_param = input_names + attr_names

        input_tensor_code += SINGLE_PREPARE_DATA_TEMPLATE.format(
            input_name,
            input_name,
            kernel_param.index(input_name),
            trans_flag,
            input_name,
            input_name,
        )

        return input_tensor_code

    def generate_prepare_data_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        kernel_param = self.kernel['param']
        if kernel_param is None:
            kernel_param = input_names + attr_names
        input_tensor_code = ""
        for i, input_name in enumerate(input_names):
            # set input code
            if input_name in kernel_param:
                # onlu support dense tensor
                api_tensor_type = self.inputs['input_info'][input_name]
                phi_tensor_type = 'dense'
                if api_tensor_type in self.gene_dist_input_func.keys():
                    input_tensor_code += self.gene_dist_input_func[
                        api_tensor_type
                    ][phi_tensor_type](input_name)
                else:
                    # do nothing
                    pass
            else:
                if input_name in self.infer_meta['param']:
                    if input_name in self.optional_vars:
                        input_tensor_code += (
                            INFER_META_OPTIONAL_INPUT_TEMPLATE.format(
                                input_name, input_name, input_name, input_name
                            )
                        )
                    else:
                        if (
                            self.inputs['input_info'][input_name]
                            == "const std::vector<Tensor>&"
                        ):
                            input_tensor_code += (
                                INFER_META_VECTOR_INPUT_TEMPLATE.format(
                                    input_name, input_name, input_name
                                )
                            )
                        else:
                            input_tensor_code += (
                                INFER_META_SINGLE_INPUT_TEMPLATE.format(
                                    input_name,
                                    input_name,
                                    input_name,
                                    input_name,
                                )
                            )

        return input_tensor_code

    def generate_infer_meta_code(self) -> str:
        input_names = self.inputs['names']
        attr_names = self.attrs['names']
        output_names = self.outputs['names']

        # 1. get infer meta func name
        infer_meta = self.infer_meta
        infer_meta_func_code = infer_meta['func']

        # 2. get meta tensor input args
        infer_meta_params = (
            infer_meta['param']
            if infer_meta['param'] is not None
            else input_names + attr_names
        )
        input_args_code = ""
        for param in infer_meta_params:
            if param in input_names:
                if self.inputs['input_info'][param] == "const Tensor&":
                    input_args_code += SINGLE_META_IN_TEMPLATE.format(param)
                else:
                    raise ValueError(
                        f"{self.api} : Param of infer_meta error : {self.inputs['input_info'][param]} type is not supported."
                    )
            elif param in attr_names:
                input_args_code = input_args_code + param + ", "
            elif isinstance(param, str):
                input_args_code = input_args_code + "\"" + param + "\", "
            elif isinstance(param, bool):
                input_args_code = input_args_code + str(param).lower() + ", "
            else:
                input_args_code = input_args_code + str(param) + ", "

        # 3. get meta tensor output args
        output_decl_code = ""
        output_args_code = ""
        for i, out_name in enumerate(self.dense_output_args):
            if self.outputs['types'][i] == 'std::vector<Tensor>':
                # TODO(chenweihang): support vector output later
                pass
            else:
                output_decl_code += SINGLE_META_OUT_DECL_TEMPLATE.format(
                    out_name, out_name
                )
                if len(self.dense_output_args) == 1:
                    output_args_code += f"&meta_{out_name}, "
                else:
                    output_args_code += (
                        f"{out_name} ? &meta_{out_name} : nullptr, "
                    )
        output_args_code = output_args_code[:-2]

        return output_decl_code + INFER_META_TEMPLATE.format(
            infer_meta_func_code, input_args_code, output_args_code
        )

    def generate_kernel_call_code(self) -> str:
        dense_input_trans_map = {
            'const Tensor&': 'const phi::DenseTensor&',
            'const std::vector<Tensor>&': 'const std::vector<const phi::DenseTensor*>&',
            'const paddle::optional<Tensor&>': 'paddle::optional<const phi::DenseTensor&>',
            'const paddle::optional<Tensor>&': 'const paddle::optional<phi::DenseTensor>&',
            'const paddle::optional<std::vector<Tensor>>&': 'const paddle::optional<std::vector<const phi::DenseTensor*>>&',
        }
        dense_output_trans_map = {
            'Tensor': 'phi::DenseTensor*',
            'std::vector<Tensor>': 'std::vector<phi::DenseTensor*>',
        }

        input_names = self.inputs['names']
        input_infos = self.inputs['input_info']
        kernel_args_type_list = ['const phi::DeviceContext&']

        attr_names = self.attrs['names']
        kernel_args = self.kernel['param']
        if kernel_args is None:
            kernel_args = input_names + attr_names

        # 1. generate input args list
        input_args = ["*dev_ctx"]
        for arg in kernel_args:
            if arg in input_names:
                if arg in self.optional_vars:
                    input_args.append(PREFIX_TENSOR_NAME + arg)
                else:
                    if input_infos[arg] == "const Tensor&":
                        input_args.append("*" + PREFIX_TENSOR_NAME + arg)
                    elif input_infos[arg] == "const std::vector<Tensor>&":
                        input_args.append(PREFIX_TENSOR_NAME + arg)
                    else:
                        # do nothing
                        pass
                kernel_args_type_list.append(
                    dense_input_trans_map[input_infos[arg]]
                )
            elif arg in attr_names:
                if 'IntArray' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append('const phi::IntArray&')
                    arg = 'phi::IntArray(' + arg + ')'
                elif 'vector<phi::Scalar>' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append(
                        'const std::vector<phi::Scalar>&'
                    )
                elif 'Scalar' in self.attrs['attr_info'][arg][0]:
                    kernel_args_type_list.append('const phi::Scalar&')
                    arg = 'phi::Scalar(' + arg + ')'
                else:
                    kernel_args_type_list.append(
                        self.attrs['attr_info'][arg][0]
                    )
                input_args.append(arg)
            elif isinstance(arg, bool):
                input_args.append(str(arg).lower())
            else:
                input_args.append(str(arg))

        # 2. generate output args list
        # record into `self.dense_output_args` in `generate_output_creation_code` function

        # 3. generate kernel signature
        for i, out_type in enumerate(self.outputs['types']):
            kernel_args_type_list.append(dense_output_trans_map[out_type])
        kernel_signature = "void(*)(" + ", ".join(kernel_args_type_list) + ")"

        return KERNEL_CALL_TEMPLATE.format(
            kernel_signature,
            ", ".join(input_args),
            ", ".join(self.dense_output_args),
        )

    def generate_reshard_output_code(self) -> str:
        return OUTPUT_RESHARD_TEMPLATE.format()

    def generate_return_code(self) -> str:
        return self.gene_return_code()

    def generate_auto_paralel_branch(self) -> str:
        # if no tensor input, do not genetate auto parallel branch
        if len(self.inputs['names']) == 0:
            return ""
        return MAIN_DIST_BRANCH_TEMPLATE.format(
            self.generate_if_condition_code(),
            self.generate_output_creation_code(),
            self.generate_infer_spmd_code(),
            self.generate_kernel_selection_code(),
            self.generate_reshard_input_code(),
            self.generate_prepare_data_code(),
            self.generate_infer_meta_code(),
            self.generate_kernel_call_code(),
            self.generate_reshard_output_code(),
            self.generate_return_code(),
        )

    def check_argument_whether_support_auto_parallel(self):
        for name in self.inputs['names']:
            if self.inputs['input_info'][name] != "const Tensor&":
                return False
        for out_type in self.outputs['types']:
            if out_type != "Tensor":
                return False
        return True

    # override BaseAPI's method
    def gene_base_api_code(self, inplace_flag=False):
        # init status
        self.inplace_flag = inplace_flag
        self.dist_output_args = []
        self.dense_output_args = []
        # generate api body
        api_func_name = self.get_api_func_name()
        if inplace_flag and api_func_name[-1] != '_':
            api_func_name += '_'

        if len(self.kernel['func']) > 1:
            kernel_dispatch_code = ''
            for kernel_name in self.kernel['func']:
                kernel_dispatch_code += self.gene_dispatch_code(
                    kernel_name, inplace_flag
                )
            return API_IMPL_TEMPLATE.format(
                self.get_return_type(inplace_flag),
                api_func_name,
                self.get_define_args(inplace_flag),
                self.gene_kernel_select(),
                kernel_dispatch_code
                + DIPATCH_END_GUARD_TEMPLATE.format(self.api),
            )
        else:
            # auto parallel branch, all apis contains this branch default
            # 1. only works for the ops contains single kernel
            # 2. doesn't support initialize ops now
            # 3. doesn't support view api
            # 4. only for general forward and backward
            # 5. only support single tensor input and output
            dist_branch_code = ""
            if (
                len(self.inputs['names']) > 0
                and len(self.view_map) == 0
                and self.check_argument_whether_support_auto_parallel()
            ):
                dist_branch_code = self.generate_auto_paralel_branch()

            return API_IMPL_TEMPLATE.format(
                self.get_return_type(inplace_flag),
                api_func_name,
                self.get_define_args(inplace_flag),
                self.gene_kernel_select(),
                dist_branch_code
                + self.gen_kernel_code(
                    self.kernel['func'][0], '', inplace_flag
                ),
            )

def generate_api(
    api_yaml_path, is_fused_ops_yaml, header_file_path, source_file_path
):
    apis = []

    for each_api_yaml in api_yaml_path:
        with open(each_api_yaml, 'r') as f:
            api_list = yaml.load(f, Loader=yaml.FullLoader)
            if api_list:
                apis.extend(api_list)

    header_file = open(header_file_path, 'w')
    source_file = open(source_file_path, 'w')

    namespace = api_namespace()

    header_file.write("#pragma once\n")
    header_file.write(header_include())
    header_file.write(namespace[0])

    include_header_file = (
        "paddle/phi/api/include/fused_api.h"
        if is_fused_ops_yaml is True
        else "paddle/phi/api/include/api.h"
    )
    # not all fused ops supoort dygraph
    if is_fused_ops_yaml is True:
        new_apis = [
            api
            for api in apis
            if "support_dygraph_mode" in api
            and api["support_dygraph_mode"] is True
        ]
        apis = new_apis

    source_file.write(source_include(include_header_file))
    source_file.write(namespace[0])

    for api in apis:
        dist_foward_api = DistForwardAPI(api)
        if dist_foward_api.is_dygraph_api:
            dist_foward_api.is_dygraph_api = False

        header_file.write(dist_foward_api.gene_api_declaration())
        if is_fused_ops_yaml is True:
            source_file.write(dist_foward_api.gene_api_code())
        else:
            source_file.write(dist_foward_api.gene_api_code())

    header_file.write(namespace[1])
    source_file.write(namespace[1])

    source_file.write(declare_extension_api())

    header_file.close()
    source_file.close()

def main():
    parser = argparse.ArgumentParser(
        description='Generate PaddlePaddle C++ API files'
    )
    parser.add_argument(
        '--api_yaml_path',
        help='path to api yaml file',
        nargs='+',
        default=['paddle/phi/api/yaml/ops.yaml'],
    )

    parser.add_argument(
        '--is_fused_ops_yaml',
        help='flag of fused ops yaml',
        action='store_true',
    )

    parser.add_argument(
        '--api_header_path',
        help='output of generated api header code file',
        default='paddle/phi/api/include/api.h',
    )

    parser.add_argument(
        '--api_source_path',
        help='output of generated api source code file',
        default='paddle/phi/api/lib/api.cc',
    )

    options = parser.parse_args()

    api_yaml_path = options.api_yaml_path
    is_fused_ops_yaml = options.is_fused_ops_yaml
    header_file_path = options.api_header_path
    source_file_path = options.api_source_path

    generate_api(
        api_yaml_path, is_fused_ops_yaml, header_file_path, source_file_path
    )

if __name__ == '__main__':
    main()


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import math
import os
from preprocessing import *

class Data :

    def __init__(self, X=None, y=None, Xnames=None, yname=None, weights=None) :
        """
        Data class.

        Attributes
        --------------------
            X       -- numpy array of shape (n,d), features
            y       -- numpy array of shape (n,), targets
        """

        # n = number of examples, d = dimensionality
        self.X = X
        self.y = y
        self.Xnames = Xnames
        self.yname = yname
        self.weights = weights

    def load(self, filenameX, filenamey, header=0) :
        """
        Load csv file into X array of features and y array of labels.

        Parameters
        --------------------
            filename   -- string, filename
        """

        # determine filename
        dir = os.path.dirname(__file__)
        fX = os.path.join(dir, 'Dataset', filenameX)
        fy = os.path.join(dir, 'Dataset', filenamey)

        # load data
        with open(fX, 'r') as fid :
            dataX = np.loadtxt(fid, delimiter=",", skiprows=header)

        with open(fy, 'r') as fid :
            datay = np.loadtxt(fid, delimiter=",", skiprows=header)

        # separate features and labels
        self.X = dataX
        self.y = datay

        # load feature and label names
        if header != 0:
            with open(fX, 'r') as fid :
                headerX = fid.readline().rstrip().split(",")

            with open(fy, 'r') as fid :
                headery = fid.readline().rstrip().split(",")

            self.Xnames = headerX
            self.yname = headery

        else:
            self.Xnames = None
            self.yname = None

    def plot(self) :
        """Plot features and labels."""
        pos = np.nonzero(self.y > 0)  # matlab: find(y > 0)
        neg = np.nonzero(self.y < 0)  # matlab: find(y < 0)
        plt.plot(self.X[pos,0], self.X[pos,1], 'b+', markersize=5)
        plt.plot(self.X[neg,0], self.X[neg,1], 'ro', markersize=5)
        plt.show()

    def plot_histogram(self) :
        """
        Plots histogram of values in X grouped by y.

        Parameters
        --------------------
            X     -- numpy array of shape (n,), feature values
            y     -- numpy array of shape (n,), target classes
            Xnames -- string, name of feature
            yname -- string, name of target
        """
        for i in xrange(len(self.Xnames)):
            # set up data for plotting
            targets = sorted(set(self.y))
            data = []; labels = []
            for target in targets :
                features = [self.X[:,i][j] for j in xrange(len(self.y)) if self.y[j] == target]
                data.append(features)
                labels.append('%s = %s' % (self.yname[0], target))

            # set up histogram bins
            features = set(self.X[:,i])
            nfeatures = len(features)
            test_range = range(int(math.floor(min(features))), int(math.ceil(max(features)))+1)
            if nfeatures < 33 and sorted(features) == test_range:
                bins = test_range + [test_range[-1] + 1] # add last bin
                align = 'left'
            else :
                bins = 33
                align = 'mid'

            # plot
            plt.figure()
            n, bins, patches = plt.hist(data, bins=bins, align=align, alpha=0.5, label=labels, log=True, density=False)
            # n, bins, patches = plt.hist(data, density=True)
            plt.xlabel(self.Xnames[i])
            plt.ylabel('Frequency')
            plt.legend()
            plt.legend(loc='upper left')
            plt.show()

def save_data(filename, df):
    df.to_csv(filename)

# helper functions
def load_data(filenameX, filenamey, header=0) :
    """Load csv file into Data class."""
    data = Data()
    data.load(filenameX, filenamey, header=header)
    return data

def update_data():
    npX, npy, Xnames, X, y = main()
    X.to_csv('Dataset/X.csv', index=False)
    y.to_csv('Dataset/y.csv', index=False, header=True)


from django.db import models

class DGNode(models.Model):
    id = models.AutoField(primary_key=True)

    creation_date = models.DateTimeField(auto_now=True)

    def get_owner(self):
      if hasattr(self, 'event_owner'):
        return self.event_owner
      if hasattr(self, 'user_profile_owner'):
        return self.user_profile_owner.user
      return None

    class Meta(object):
        app_label = 'wisimi_server'


import sys,os,re
from multiprocessing import Process, Queue, Pool

usage = """
<folder of alignments>
(N)ucleotide or (P)rotein?>
<Num bootstraps (0 for no bootstrapping)>
<output folder>
<cores to use>
<cores per tree>\n\n"""
argv = sys.argv
junk = argv.pop(0)

if len(sys.argv) == 6:
	in_folder = argv.pop(0)
	type_flag = argv.pop(0)
	s_flag = int(argv.pop(0))
	out_folder = argv.pop(0)
	cores=int(argv.pop(0))
	coresPerTree = int(argv.pop(0))
else:
	sys.exit(usage)

def multi_run(cmd):
	os.system(cmd)

rxml = '~/tools_and_software/RAxML-8.1.24/raxmlHPC-PTHREADS-AVX2'

os.system('mkdir -p {0}'.format(out_folder))
aligns = os.listdir(in_folder)

cmd = []
NAMES = {}

for f in aligns:
	f_name = f
	f_name = f_name.replace('.fna','')
	f_name = f_name.replace('.faa','')
	f_name = f_name.replace('.mafft','')
	NAMES[f_name] = ''
	if type_flag == 'N':
		if s_flag > 0:
			data = (rxml,in_folder,f,f_name,coresPerTree,s_flag)
			cmd.append('{0} -s {1}/{2} -n {3} -x 3524627 -p 3524627 -f a -T {4} -N {5} -m GTRCAT -c 6'.format(*data))
		else:
			data = (rxml,in_folder,f,f_name,coresPerTree,s_flag)
			cmd.append('{0} -s {1}/{2} -n {3} -T {4} -m GTRCAT -c 6'.format(*data))
	elif type_flag == 'P':
		if s_flag > 0:
			data = (rxml,in_folder,f,f_name,coresPerTree,s_flag)
			cmd.append('{0} -s {1}/{2} -n {3} -x 3524627 -f a -T {4} -N {5} -m PROTGAMMABLOSUM62'.format(*data))
		else:
			data = (rxml,in_folder,f,f_name,coresPerTree,s_flag)
			cmd.append('{0} -s {1}/{2} -n {3} -T {4} -m PROTGAMMABLOSUM62'.format(*data))

numRun = cores / coresPerTree
p=Pool(numRun)
for c in cmd:
	p.apply_async(multi_run, args=(c,))
	pass
p.close()
p.join()

outputs = os.listdir('./')
for f in outputs:
	front = f.split('.')[0]
	tail = '.'.join(f.split('.')[1:])
	if s_flag == 0:
		if front == 'RAxML_bestTree' and NAMES.has_key(tail):
			os.system('mv {0} {1}/{2}.nwk'.format(f,out_folder,tail))
	else:
		if front == 'RAxML_bipartitionsBranchLabels' and NAMES.has_key(tail):
			os.system('mv {0} {1}/{2}.nwk'.format(f,out_folder,tail))


#!/usr/bin/env python
"""Bayesian neural network using mean-field variational inference
(see, e.g., Blundell et al. (2015); Kucukelbir et al. (2016)).

Inspired by autograd's Bayesian neural network example.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import edward as ed
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from edward.models import Normal
from edward.stats import norm
from edward.util import rbf

class BayesianNN:
  """
  Bayesian neural network for regressing outputs y on inputs x.

  p((x,y), z) = Normal(y | NN(x; z), lik_std) *
                Normal(z | 0, prior_std),

  where z are neural network weights, and with known likelihood and
  prior standard deviations.

  Parameters
  ----------
  layer_sizes : list
    The size of each layer, ordered from input to output.
  nonlinearity : function, optional
    Non-linearity after each linear transformation in the neural
    network; aka activation function.
  lik_std : float, optional
    Standard deviation of the normal likelihood; aka noise parameter,
    homoscedastic noise, scale parameter.
  prior_std : float, optional
    Standard deviation of the normal prior on weights; aka L2
    regularization parameter, ridge penalty, scale parameter.
  """
  def __init__(self, layer_sizes, nonlinearity=tf.tanh,
               lik_std=0.1, prior_std=1.0):
    self.layer_sizes = layer_sizes
    self.nonlinearity = nonlinearity
    self.lik_std = lik_std
    self.prior_std = prior_std

    self.n_layers = len(layer_sizes)
    self.weight_dims = list(zip(layer_sizes[:-1], layer_sizes[1:]))
    self.n_vars = sum((m + 1) * n for m, n in self.weight_dims)

  def unpack_weights(self, zs):
    """Unpack weight matrices and biases from a flattened vector."""
    for m, n in self.weight_dims:
      yield tf.reshape(zs[:(m * n)], [m, n]), \
          tf.reshape(zs[(m * n):(m * n + n)], [n])
      zs = zs[(m + 1) * n:]

  def neural_network(self, x, zs):
    """Forward pass of the neural net, outputting a vector of
    `n_minibatch` elements."""
    h = x
    for W, b in self.unpack_weights(zs):
      h = self.nonlinearity(tf.matmul(h, W) + b)

    return tf.squeeze(h)  # n_minibatch x 1 to n_minibatch

  def log_prob(self, xs, zs):
    """Return scalar, the log joint density log p(xs, zs)."""
    x, y = xs['x'], xs['y']
    zs = zs['z']
    log_prior = tf.reduce_sum(norm.logpdf(zs, 0.0, self.prior_std))
    mu = self.neural_network(x, zs)
    log_lik = tf.reduce_sum(norm.logpdf(y, mu, self.lik_std))
    return log_lik + log_prior

def build_toy_dataset(N=40, noise_std=0.1):
  D = 1
  x = np.concatenate([np.linspace(0, 2, num=N / 2),
                      np.linspace(6, 8, num=N / 2)])
  y = np.cos(x) + np.random.normal(0, noise_std, size=N)
  x = (x - 4.0) / 4.0
  x = x.reshape((N, D))
  return x, y

ed.set_seed(42)
x_train, y_train = build_toy_dataset()

model = BayesianNN(layer_sizes=[1, 10, 10, 1], nonlinearity=rbf)

qz_mu = tf.Variable(tf.random_normal([model.n_vars]))
qz_sigma = tf.nn.softplus(tf.Variable(tf.random_normal([model.n_vars])))
qz = Normal(mu=qz_mu, sigma=qz_sigma)

# Set up figure
fig = plt.figure(figsize=(8, 8), facecolor='white')
ax = fig.add_subplot(111, frameon=False)
plt.ion()
plt.show(block=False)

sess = ed.get_session()
data = {'x': x_train, 'y': y_train}
inference = ed.KLqp({'z': qz}, data, model)
inference.initialize()

init = tf.initialize_all_variables()
init.run()

for t in range(inference.n_iter):
  info_dict = inference.update()
  inference.print_progress(info_dict)

  if t % inference.n_print == 0:
    # Sample functions from variational model
    mean, std = sess.run([qz.mu, qz.sigma])
    rs = np.random.RandomState(0)
    zs = rs.randn(10, model.n_vars) * std + mean
    zs = tf.convert_to_tensor(zs, dtype=tf.float32)
    inputs = np.linspace(-8, 8, num=400, dtype=np.float32)
    x = tf.expand_dims(inputs, 1)
    mus = []
    for z in tf.unpack(zs):
      mus += [model.neural_network(x, z)]

    outputs = tf.pack(mus).eval()

    # Get data
    x, y = data['x'], data['y']

    # Plot data and functions
    plt.cla()
    ax.plot(x, y, 'bx')
    ax.plot(inputs, outputs.T)
    ax.set_xlim([-8, 8])
    ax.set_ylim([-2, 3])
    plt.draw()
    plt.pause(1.0 / 60.0)


import argparse
import click
import pandas as pd
import re
from datetime import datetime
from datetime import timedelta
from os import chdir
from os import chdir
from os import getcwd
from os import makedirs
from os import mkdir
from os import remove
from os import rename
from os import system
from os.path import basename
from os.path import dirname
from os.path import expanduser
from os.path import getmtime
from os.path import isdir
from os.path import isdir
from os.path import isfile
from os.path import join
from os.path import splitext
from pydoni.db import progrun_update
from pydoni.os import listfiles
from pydoni.os import macos_notify
from pydoni.sh import syscmd
from pydoni.vb import echo
from pydoni.vb import macos_notify
from pydoni.vb import program_complete
from pydoni.vb import verbose_header
from send2trash import send2trash
from shutil import copy2
from tqdm import tqdm

def parse_date(date):
    """
    Get date in format YYYY-MM-DD if entered as 'today' or 'yesterday'.

    Arguments:
        date {str} -- date in format YYYYMMDD or 'today' or 'yesterday'

    Returns:
        {str}
    """
    if date == 'today':
        date = datetime.strftime(datetime.now(), '%Y-%m-%d')
    elif date == 'yesterday':
        date = datetime.strftime(datetime.now() - timedelta(days=1), '%Y-%m-%d')
    assert re.search(r'^\d{4}-\d{2}-\d{2}$', date)
    return date

def parse_target_subcollection(fname):
    """
    Given a filename, get the target subcollection as one of ['Photo', 'Video', 'Gopro', 'Drone'].

    Arguments:
        fname {str} -- name of file to parse

    Returns:
        {str}
    """
    c = Convention()
    e = Extension()
    media_type = 'Photo' if splitext(fname)[1].lower() in e.photo else 'Video'
    conv = c.photo if media_type == 'Photo' else c.video
    camera = re.match(conv, fname).group('camera')
    if re.search(r'gopro|hero', camera, flags=re.IGNORECASE):
        return 'Gopro'
    elif 'FC' in camera or 'L1D' in camera:
        if media_type == 'Photo':
            return join('Drone', 'Photo')
        elif media_type == 'Video':
            return join('Drone', 'Video')
    else:
        return media_type

class Extension(object):
    """
    Hold valid file extension lists of strings.
    """

    def __init__(self):
        self.photo = ['.jpg', '.jpeg', '.dng', '.arw', '.cr2']
        self.video = ['.mov', '.mp4', '.mts', '.m4v', '.avi']

class Convention(object):
    """
    Hold naming convention regex strings.
    """

    def __init__(self):
        self.photo = ''
        self.video = ''
        conv_base = '' +\
            r'(?P<year>\d{4})' +\
            r'(?P<month>\d{2})' +\
            r'(?P<day>\d{2})' +\
            r'(?P<initials>[A-Za-z]{2})' +\
            r'(?P<hours>\d{2})' +\
            r'(?P<minutes>\d{2})' +\
            r'(?P<seconds>\d{2})'
        self.photo = conv_base +\
            '_' +\
            r'(?P<camera>.*?)' +\
            '_' +\
            r'(?P<seqnum>(\d{4,}|\d+-\d|-\d|))' +\
            r'(?P<affix>(-HDR-*\d*|-Pano-*\d*|-Edit-*\d*|-Stack-*\d*|)*)' +\
            r'(?P<ext>\.[a-z0-9]{3})'
        self.video = conv_base +\
            '_' +\
            r'(?P<camera>.*?)' +\
            '_' +\
            r'(?P<seqnum>.*(\d+))' +\
            '_' +\
            r'(?P<quality>Q\d+(K|P))' +\
            r'(?P<fps>\d{2,3}FPS)' +\
            r'(?P<ext>\.[a-z0-9]{3})'

class MediaFile(Convention, Extension):
    """
    Gather information for a photo or video file.

    Arguments:
        fpath {str} -- absolute path to file to initialize MediaFile class on

    Keyword Arguments:
        run_exiftool {bool} -- if True, run ExifTool automatically and save output dictionary as `self.exif` attribute (default: {False})
    """

    def __init__(self, fpath, run_exiftool=False):
        from os.path import isfile, dirname, basename, splitext, abspath
        from pydoni.classes import Attribute

        assert isinstance(fpath, str)
        assert isfile(fpath)
        assert isinstance(run_exiftool, bool)

        # Inherit naming conventions and extensions
        Convention.__init__(self)
        Extension.__init__(self)

        # Assign filename and directory name attributes to `self`
        self.fpath_abs = abspath(fpath)
        self.fname = basename(self.fpath_abs)
        self.dname = dirname(self.fpath_abs)
        assert self.fname != self.fpath_abs
        assert self.dname > ''

        # Define extensions
        self.ext = splitext(self.fname)[1].lower()
        self.valid_ext = Extension()
        assert self.ext in self.valid_ext.photo + self.valid_ext.video

        # Parse media type from extension
        self.mtype = self.parse_media_type()
        self.remove_flag = True if self.mtype == 'remove' else False

        # Run ExifTool if specified
        if run_exiftool is True:
            from pydoni.sh import EXIF
            self.exif = EXIF(self.fpath_abs).run(wrapper='doni')

    def parse_media_type(self):
        """
        Given a file extension, get the type of media (one of 'photo' or 'video')

        Returns:
            {str} -- type of media
        """

        if self.ext in self.valid_ext.photo:
            return 'photo'
        elif self.ext in self.valid_ext.video:
            return 'video'
        elif self.ext in self.valid_ext.rm:
            return 'remove'
        else:
            from pydoni.vb import echo
            echo('Invalid file extension slipped through the cracks!', fn_name='MediaFile.parse_media_type', abort=True)

    def build_fname(self, initials, tz_adjust=0):
        """
        Build new filename from EXIF metadata.

        Arguments:
            initials {str} -- two character initials string

        Keyword Arguments:
            tz_adjust {int} -- alter 'hours' by set number to account for timezone adjust (default: {0})

        Returns:
            {str} -- new filename according to convention (basename)
        """

        from os.path import join

        assert isinstance(initials, str)
        assert len(initials) == 2
        req_attrs = ['exif', 'mtype', 'dname', 'fname']
        for ra in req_attrs:
            assert hasattr(self, ra)
        assert self.mtype in ['photo', 'video']
        assert isinstance(tz_adjust, int)
        assert tz_adjust in range(-24, 24)

        def search_exif(exif, attrs, def_str):
            """
            Search exif dictionary for name attributes in the order specified. For example
            if `attrs` argument is ['image_width', 'video_size'], first search `exif` for
            the attribute 'image_width'. If it exists and is a valid value, return this
            value. If not, search for 'video_size' and return that value if it exists and
            is valid. If neither value exists, return an arbitrary string.

            Arguments:
                attrs {str} or {list} -- attribute name(s) to search exif dictionary for
                def_str {str} -- return this string if attributes not found
            """
            assert isinstance(exif, dict)
            assert isinstance(attrs, str) or isinstance(attrs, list)
            assert isinstance(def_str, str)

            attrs = [attrs] if isinstance(attrs, str) else attrs
            bogus = ['0000:00:00 00:00:00']

            for attr in attrs:
                if attr in exif.keys():
                    val = exif[attr]
                    if val not in bogus:
                        return val

            return def_str

        def extract_video_quality(exif, def_str):
            """
            Given EXIF dictionary, extract video quality.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- video quality string
            """
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)

            image_width = search_exif(
                exif    = exif,
                attrs   = ['image_width', 'video_size'],
                def_str = def_str)
            if image_width == def_str:
                return def_str
            elif image_width == '3840':
                return '4K'
            elif image_width == '2704':
                return '1520P'
            elif image_width == '1920':
                return '1080P'
            elif image_width == '1280':
                return '720P'
            else:
                return image_width

        def extract_video_framerate(exif, def_str):
            """
            Given EXIF dictionary, extract video frame rate.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- video framerate string
            """
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)

            framerate = search_exif(
                exif    = exif,
                attrs=['video_avg_frame_rate', 'video_frame_rate'],
                def_str = def_str)
            if framerate == def_str:
                return def_str
            else:
                try:
                    return round(float(framerate))
                except:
                    return framerate

        def extract_capture_date_and_time(exif, def_str, mtype, tz_adjust=None):
            """
            Given EXIF dictionary, extract file capture date and time.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found
                mtype {str} -- `MediaFile.mtype` (media type string)

            Keyword Arguments:
                tz_adjust {int} -- alter 'hours' by set number to account for timezone adjust

            Returns:
                {tuple} -- (capture date string, capture time string)
            """
            import re
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)
            assert mtype in ['photo', 'video']

            if mtype == 'photo':
                attrs = ['create_date', 'file_modification_date/time']
            elif mtype == 'video':
                attrs = ['file_modification_date/time', 'create_date']

            # Get capture date
            raw_dt = search_exif(
                exif    = exif,
                attrs   = attrs,
                def_str = def_str
            )
            capture_date = re.sub(r'(.*?)(\s+)(.*)', r'\1', raw_dt)
            capture_date = capture_date.replace(':', '')

            # Get capture time
            capture_time = raw_dt.replace(':', '')
            capture_time = re.sub(r'(.*?)(\s+)(.*)', r'\3', capture_time)[0:6]

            # Adjust hours by timezone if specified. May also need to alter day if timezone
            # adjust makes the day spill over into the next/previous day
            if tz_adjust is not None:
                import datetime
                dt = datetime.datetime.strptime(capture_date + capture_time, '%Y%m%d%H%M%S')
                dt = dt + datetime.timedelta(hours=tz_adjust)
                capture_date = dt.strftime('%Y%m%d')
                capture_time = dt.strftime('%H%M%S')

            return (capture_date, capture_time)

        def extract_sequence_number(fname, def_str):
            """
            Given filename, extract file sequence number as the final occurrence of 4 or 5
            digits in a filename.

            Arguments:
                fname {str} -- filename to extract sequence number from
                def_str {str} -- return this string if value is not not found

            Returns:
                {str} -- sequence number
            """
            import re
            from os.path import basename
            assert isinstance(fname, str)
            assert isinstance(def_str, str)

            # Extract sequence number. Search first for 5 digits at the end of the filename.
            # If not found, search for 4 digits.

            # Clean up searchable area of filename by removing YYYYMMDDASHHMMSS and the string
            # for video files QXXXXFPS
            fname_ = basename(fname)
            fname_ = re.sub('_Q.*?FPS', '', fname_)
            fname_ = re.sub(r'^\d{8}\w{2}\d{6}_', '', fname_)

            # Search for 5 or 4 digits
            ptn = r'^(.*)(\d{5})(.*)$'
            if re.search(ptn, fname_):
                return re.sub(ptn, r'\2', fname_)
            else:
                ptn = r'^(.*)(\d{4})(.*)$'
                if re.search(ptn, fname_):
                    return re.sub(ptn, r'\2', fname_)
                else:
                    return def_str

        def extract_camera_model(exif, def_str, dname):
            """
            Given EXIF dictionary, extract camera model.

            Arguments:
                exif {dict} -- EXIF dictionary
                def_str {str} -- return this string if value is not not found
                dname {str} -- directory name of photo, only used in parsing camera model if all else fails

            Returns:
                {str} -- camera model string
            """
            from os.path import basename, isdir
            assert isinstance(exif, dict)
            assert isinstance(def_str, str)
            assert isinstance(dname, str)
            assert isdir(dname)

            # Get camera model from raw exif
            cm_raw = search_exif(
                exif    = exif,
                attrs   = ['camera_model_name', 'device_model_name', 'model'],
                def_str = def_str
            )

            # If no camera model found, apply manual corrections
            if cm_raw == def_str:

                # Get make
                make_raw = search_exif(
                    exif    = exif,
                    attrs   = ['make', 'device_manufacturer', 'compressor_name'],
                    def_str = '{}'
                )
                valid_makes = ['Sony', 'Canon', 'DJI', 'Gopro', 'Gopro AVC Encoder', '{}']
                assert any([make_raw.lower() in [x.lower() for x in valid_makes]])
                for vm in valid_makes:
                    if vm.lower() in make_raw.lower():
                        make = vm

                # Manual corrections
                if make.lower() == 'dji' or basename(dname).lower() == 'drone':
                    return 'L1D-20c'  # DJI Mavic 2 Pro
                elif make.lower() == 'gopro' or basename(dname).lower() == 'gopro':
                    return 'HERO5 Black'
                else:
                    return def_str

            else:
                return cm_raw

        exif = self.exif
        fname = self.fname
        self.initials = initials

        # Default string to be returned by any of the extract* functions if that
        # attribute is unable to be found
        def_str = '{}'

        # Extract attributes from EXIF used regardless of photo or video
        self.cd, self.ct = extract_capture_date_and_time(exif, def_str, self.mtype, tz_adjust)
        self.sn = extract_sequence_number(fname, def_str)
        self.cm = extract_camera_model(exif, def_str, self.dname)

        # Build new filename
        if self.mtype == 'photo':
            newfname = "{}{}{}_{}_{}{}".format(
                self.cd, self.initials, self.ct, self.cm, self.sn, self.ext.lower())

        elif self.mtype == 'video':
            self.vq = extract_video_quality(exif, def_str)
            self.vfr = extract_video_framerate(exif, def_str)
            newfname = "{}{}{}_{}_{}_Q{}{}FPS{}".format(
                self.cd, self.initials, self.ct, self.cm, self.sn, self.vq, self.vfr, self.ext.lower())

        return newfname

    def convert_dng(self, remove_original=False):
        """
        Convert a photo file to dng.

        Keyword Arguments:
            remove_original {bool} -- if True, remove original raw file (default: {True})
        """
        from pydoni.sh import adobe_dng_converter
        assert self.ext in self.valid_ext.photo
        adobe_dng_converter(self.fpath_abs)
        if remove_original:
            from send2trash import send2trash
            send2trash(self.fpath_abs)

class DumpSD(object):
    """
    Dump photo and video files from a source volume to a target directory. Rename and Convert
    RAW files to DNG as necessary.

    Arguments:
        vol_source {str} -- source volume name
        coll_target {str} -- absolute directory path of destination collection
        date {str} -- dump photos with this creation date, format YYYY-MM-DD
        rename {bool} -- rename media files according to defined convention
        convert {bool} -- convert RAW photos to DNG
        verbose {bool} -- print messages to STDOUT
        notify {bool} -- send macOS notification upon program completion
        test_mode {bool} -- run program in test mode
        rc_initials {str} -- initials parameter for RenameConvert. Only must be specified if 'rename' or 'convert' is True (default: {AS})
        rc_tz_adjust {int} -- tz_adjust parameter for RenameConvert. Only must be specified if 'rename' or 'convert' is True (default: {0})
    """

    def __init__(
        self,
        vol_source,
        coll_target,
        date,
        rename,
        convert,
        verbose,
        notify,
        test_mode,
        rc_initials,
        rc_tz_adjust
        ):
        self.vol_source = join('/', 'Volumes', vol_source)
        self.coll_target = coll_target
        self.vol_target = '/'.join(self.coll_target.split('/')[0:3])
        self.date = parse_date(date)
        self.rename = rename
        self.convert = convert
        self.verbose = verbose
        self.notify = notify
        self.test_mode = test_mode
        self.rc_initials = rc_initials
        self.rc_tz_adjust = rc_tz_adjust

        assert isdir(self.vol_source)
        assert isdir(self.coll_target)
        assert isdir(self.vol_target)

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):

        prog_start = datetime.now()
        res = {'source': [], 'target': []}

        if self.test_mode:
            echo("Initializing Program: 'Dump SD' in Test Mode...", sleep=1)
            echo('')
            echo('Begin Operational Code')
            echo('')

        if self.test_mode:
            echo('Checking source and target directories...', sleep=.25, fg='blue')
            echo('Source volume name     : {} {}'.format(
                self.vol_source, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('Source volume type     : {} {}'.format(
                self.vol_source_type, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('Target collection      : {} {}'.format(
                self.coll_target, emojize(':white_check_mark:', use_aliases=True)), sleep=.25)
            echo('')

        # List all files in source volume that match given date
        chdir(self.vol_source)
        ext = Extension()
        fnames = [f for f in listfiles(recursive=True, ext=ext.photo+ext.video) \
            if datetime.fromtimestamp(getmtime(f)).strftime('%Y-%m-%d') == self.date \
            and not re.search(r'T\d{2,3}\.', f)  # Eliminate Sony movie thumbnail files
        ]
        if not len(fnames):
            echo("No photo or video files found at '{}' for day '{}'!".format(
                self.vol_source, self.date), abort=True)

        # Copy files to temporary directory
        tmpdir = join(self.vol_target, '.tmp.dumpsd')
        if isdir(tmpdir):
            send2trash(tmpdir)
        mkdir(tmpdir)
        if self.verbose:
            verbose_header('Dumping %s media files' % str(len(fnames)))
            with tqdm(total=len(fnames), unit='mediafile') as pbar:
                for f in fnames:
                    pbar.set_postfix(mediafile=f[-10:])
                    copy2(f, join(tmpdir, basename(f)))
                    res['source'].append(f)
                    pbar.update(1)
        else:
            for f in fnames:
                res['source'].append(f)
                copy2(f, join(tmpdir, basename(f)))

        # Run RenameConvert on directory
        if self.rename or self.convert:
            ignore_rename = not self.rename
            ignore_convert = not self.convert
            RenameConvert(
                dname=tmpdir,
                initials=self.rc_initials,
                tz_adjust=self.rc_tz_adjust,
                ignore_rename=ignore_rename,
                ignore_convert=ignore_convert,
                recursive=False,
                notify=self.notify,
                verbose=self.verbose
            ).run()

        # Copy files to target collection
        chdir(tmpdir)
        media_files = [x for x in listfiles() if splitext(x)[1].lower() != '.arw']
        dumped_subcolls = []
        for f in media_files:
            subcoll = parse_target_subcollection(f)
            if subcoll == basename(self.coll_target):
                destf = join(self.coll_target, f)
            else:
                destf = join(self.coll_target, subcoll, f)
            if not isdir(dirname(destf)):
                makedirs(dirname(destf))
            rename(f, destf)
            dumped_subcolls.append(subcoll)
            res['target'].append(f)

        if isdir(tmpdir):
            send2trash(tmpdir)

        if self.verbose:
            echo('')
            echo('Dump Summary', underline=True, fg='white', bold=True)
            dumped_subcolls = pd.Series(dumped_subcolls).value_counts()
            echo('Collection: %s' % self.coll_target, indent=2)
            for subcol, count in dumped_subcolls.iteritems():
                echo("Subcollection '{}': {} media files dumped".format(
                    click.style(subcol, fg='cyan'),
                    str(count)),
                    indent=2)

        if self.notify:
            content_image = '/Users/Andoni/GDrive/Code/git-doni/pydoni/pydoni/scripts/dump_sd/img/Dump-SD-Icon.png'
            content_image = content_image if isfile(content_image) else None
            macos_notify(title='Dump SD', message='Dump completed successfully!',
                content_image=content_image)

        progrun_update(
            name='Photo.dump_sd',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res=res)

class RenameConvert(object):
    """
    Rename/Convert program master class.

    Arguments:
        dname {str} -- directory name to run program on
        initials {str} -- two-character initial string
        tz_adjust {int} -- subtract this many hours from photo created date to use in file renaming (default: {0})
        ignore_rename {bool} -- if True, ignore the Rename portion of program (default: {False})
        ignore_convert {bool} -- if True, ignore the Convert portion of program (default: {False})
        recursive {bool} -- if True, operate on all subdirectories found in `dname` (default: {False})
        verbose {bool} -- if True, print messages to STDOUT (default: {False})
    """

    def __init__(
        self,
        dname,
        initials='AS',
        tz_adjust=0,
        ignore_rename=False,
        ignore_convert=False,
        recursive=False,
        verbose=False,
        notify=False
        ):
        self.dname = dname
        self.initials = initials
        self.tz_adjust = tz_adjust
        self.ignore_rename = ignore_rename
        self.ignore_convert = ignore_convert
        self.recursive = recursive
        self.verbose = verbose
        self.notify = notify

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):

        prog_start = datetime.now()
        res = {'rename': [], 'convert': []}

        assert isdir(self.dname)
        assert isinstance(self.initials, str)
        assert len(self.initials) == 2
        assert isinstance(self.tz_adjust, int)
        assert isinstance(self.ignore_rename, bool)
        assert isinstance(self.ignore_convert, bool)
        assert isinstance(self.recursive, bool)
        assert isinstance(self.verbose, bool)

        chdir(self.dname)
        c = Convention()

        # Rename
        if not self.ignore_rename:
            media_files_r = listfiles(full_names=True, recursive=self.recursive)
            media_files_r = [x for x in media_files_r \
                if not re.match(c.photo, basename(x)) \
                and not re.match(c.video, basename(x))]

            if len(media_files_r):

                if self.verbose:
                    print()
                    verbose_header('Renaming %s media files' % str(len(media_files_r)))
                    pbar = tqdm(total=len(media_files_r), unit='mediafile')

                for f in media_files_r:
                    if self.verbose:
                        pbar.set_postfix(mediafile=f[-10:])

                    mf = MediaFile(f, run_exiftool=True)

                    if mf.remove_flag is True:
                        send2trash(f)
                        if self.verbose:
                            pbar.update(1)
                        continue

                    newfname = mf.build_fname(initials=self.initials, tz_adjust=self.tz_adjust)
                    if basename(f) != basename(newfname):
                        rename(f, newfname)
                        res['rename'].append(f)

                    if self.verbose:
                        pbar.update(1)

            else:
                if self.verbose:
                    echo('No files to rename!', fg='green')

        # Convert to DNG
        if not self.ignore_convert:
            media_files_c = listfiles(full_names=True, ext=['arw', 'cr2'], recursive=self.recursive)

            if len(media_files_c):
                if self.verbose:
                    print()
                    verbose_header('Converting %s raw files' % str(len(media_files_c)))
                    pbar = tqdm(total=len(media_files_c), unit='rawfile')

                for f in media_files_c:
                    if self.verbose:
                        pbar.set_postfix(mediafile=f[-10:])

                    mf = MediaFile(f, run_exiftool=False).convert_dng(remove_original=True)
                    res['convert'].append(f)

                    if self.verbose:
                        pbar.update(1)

            else:
                if self.verbose:
                    echo('No files to convert!', fg='green')

        if self.verbose:
            echo('')
            echo('Rename/Convert Summary', underline=True, fg='white', bold=True)
            echo('Renamed media files: %s' % str(len(media_files_r)), indent=2)
            echo('Converted RAW files: %s' % str(len(media_files_c)), indent=2)

        if self.notify:
            macos_notify(title='Rename/Convert', message='Rename/Convert completed successfully!')

        progrun_update(
            name='Photo.rename_convert',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res=res)

# FIXME: res={}
class PhotoDatabaseRefresh(object):
    """-------------------------------------------------------------------------
    Script Name: Photo Database Refresh
    Script Author: Andoni Sooklaris
    Date: 2019-02-27

    Maintain a PostgreSQL database with EXIF metadata from photo library.
    Insert new photo/video records, update media files that have changed,
    and delete media files out of database if removed on filesystem.

    Keyword Arguments:
        rscript {str} -- Path to "photo_database_refresh.R" script (default: {None})
        dir {str} -- Path to photos dir to refresh database with (default: {None})
        pg_user {str} -- Postgres username (default: {None})
        pg_db {str} -- Postgres database name (default: {None})
        pg_schema {str} -- Postgres schema name (default: {None})
        logfile {str} -- SQL file to write results to. If None (default) do not write SQL log (default: {None})
        skip_insert {bool} -- Do not insert any records to Postgres database (default: {None})
        skip_update {bool} -- Do not update any records in Postgres database (default: {None})
        skip_delete {bool} -- Do not delete any records from Postgres database (default: {None})
        verbose {bool} -- Print messages to console (default: {None})
    -------------------------------------------------------------------------"""

    def __init__(
        self,
        rscript=None,
        dir=None,
        pg_user=None,
        pg_db=None,
        pg_schema=None,
        logfile=None,
        skip_insert=None,
        skip_update=None,
        skip_delete=None,
        verbose=None,
        ):
        self.rscript = rscript
        self.dir = dir
        self.pg_user = pg_user
        self.pg_db = pg_db
        self.pg_schema = pg_schema
        self.logfile = logfile
        self.skip_insert = skip_insert
        self.skip_update = skip_update
        self.skip_delete = skip_delete
        self.verbose = verbose

        assert isfile(self.rscript)
        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def run(self):
        # If parameter(s) unspecified, default to default in R script
        prog_start = datetime.now()
        cmdargs = []
        for param, value in self.params.items():
            if param == 'rscript':
                continue
            if value is not None:
                if type(value) in [int, float]:
                    cmdargs.append('--{}={}'.format(param, value))
                elif type(value) == bool:
                    if value:
                        cmdargs.append('--{}'.format(param))
                else:
                    cmdargs.append('--{}="{}"'.format(param, value))

        cmdargs = ' '.join(cmdargs)
        cmd = 'Rscript "{}" {}'.format(self.rscript, cmdargs)
        system(cmd)

        progrun_update(
            name='Photo.photo_database_refresh',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res={})  # FIXME: Add 'res' from Rscript, likely save as CSV then read with Python

class OrganizeWebsiteExportedPlacesPhotos(object):
    """
    Rename and organize photos exported from Lightroom for andonisooklaris.com.

    Arguments:
        exported_dir {str} -- Path to exported photo directory
        places_dir {str} -- Path to andonisooklaris.com Media/Places folder
    """

    def __init__(self, exported_dir, places_dir):
        self.exported_dir = exported_dir
        self.places_dir = places_dir

        assert isdir(self.exported_dir)
        assert isdir(self.places_dir)

        self.params = {k: v for k, v in locals().items() if k != 'self'}

    def split_photo_fname(self, fname):
        """
        Extract collection, date, title and location from filename.

        Arguments:
            fname {str} -- filename to split

        Returns:
            {tuple}
        """

        fbasename, ext = splitext(fname)
        collection, date, title, location = fbasename.split('_')

        # Get appended number by Lightroom upon export in case of identical image names
        if re.search(r'-\d+$', location):
            suffix = re.search(r'-\d+$', location).group(0)
            location = location.replace(suffix, '')
        else:
            suffix = ''

        return (collection, date, title, location, suffix, ext)

    def run(self):

        prog_start = datetime.now()

        chdir(self.exported_dir)
        fnames = listfiles()

        if not len(fnames):
            echo('No files to process!', abort=True)

        for fname in fnames:
            chdir(self.places_dir)
            collection, date, title, location, suffix, ext = self.split_photo_fname(fname)

            # Determine directory to move photo to in self.places_dir,
            # based on `collection`
            collection = collection.replace('Places-', '').split('-')
            collection = collection[0]
            dest_dir = join(self.places_dir, collection)
            if not isdir(dest_dir):
                mkdir(dest_dir)

            if re.search(r'^\d{8}$', date):
                date = date[0:4] + '-' + date[4:6] + '-' + date[6:]

            if title > '':
                newfname = date + ' ' + title + ' ' + location + suffix + ext
            else:
                newfname = date + ' ' + location + suffix + ext

            newfname = join(dest_dir, newfname)

            chdir(self.exported_dir)
            if not isfile(newfname):
                rename(fname, newfname)

        progrun_update(
            name='Photo.organize_website_exported_places_photos',
            started=prog_start,
            ended=datetime.now(),
            args=self.params,
            res={'fnames': fnames})

{'isfile': 'os.path.isfile', 'dirname': 'os.path.dirname', 'basename': 'os.path.basename', 'splitext': 'os.path.splitext', 'abspath': 'os.path.abspath', 'Attribute': 'pydoni.classes.Attribute', 'EXIF': 'pydoni.sh.EXIF', 'echo': 'pydoni.vb.echo', 'join': 'os.path.join', 're': 're', 'datetime': 'datetime', 'isdir': 'os.path.isdir', 'adobe_dng_converter': 'pydoni.sh.adobe_dng_converter', 'send2trash': 'send2trash.send2trash'}


# Always prefer setuptools over distutils
from setuptools import setup, find_packages
# To use a consistent encoding
from codecs import open
import os, sys, ast
from scapyshark.version import version

here = os.path.abspath(os.path.dirname(__file__))

#with open(path.join(here, 'README.md'), encoding='utf-8') as f:
#    long_description = f.read()
long_description = "See website for more info."

setup(
    name='scapyshark',
    version=version,
    description='Wireshark, but with Scapy.',
    long_description=long_description,
    url='https://github.com/bannsec/scapyshark',
    author='Michael Bann',
    author_email='self@bannsecurity.com',
    license='MIT',
    classifiers=[
        'Development Status :: 3 - Alpha',
        'Intended Audience :: End Users/Desktop',
        'License :: OSI Approved :: MIT License',
        'Programming Language :: Python :: 3',
        'Operating System :: POSIX :: Linux',
        'Environment :: Console'
    ],
    keywords='wireshark scapy python tcpdump',
    packages=find_packages(exclude=['contrib', 'docs', 'tests']),
    install_requires=['scapy', 'prettytable', 'urwid', 'requests', 'pexpect'],
    entry_points={
        'console_scripts': [
            'scapyshark = scapyshark.scapyshark:main',
        ],
    },
)