text:
"""
Open Nodes Crawler
Copyright (c) 2018 Opennodes / Blake Bjorn Anderson

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import datetime
import json
import re
import socket
import time
import requests
import sys
import os
import logging
import pandas as pd
import numpy as np
from concurrent.futures import ThreadPoolExecutor
from geoip2.database import Reader
from geoip2.errors import AddressNotFoundError
from sqlalchemy import create_engine, and_, or_, func, not_, case
from sqlalchemy.orm import sessionmaker
from protocol import ProtocolError, Connection, ConnectionError, Keepalive
from models import Node, NodeVisitation, Base, CrawlSummary
from config import load_config, DefaultFlaskConfig

try:
    from flask_user_config import SQLALCHEMY_DATABASE_URI
except ImportError:
    try:
        from flask_config import SQLALCHEMY_DATABASE_URI
    except ImportError:
        SQLALCHEMY_DATABASE_URI = DefaultFlaskConfig.SQLALCHEMY_DATABASE_URI

logging.basicConfig(level=logging.INFO)

CONF = load_config()

def connect(network, address, port, to_services, network_data, user_agent=None, explicit_p2p=False, p2p_nodes=True,
            from_services=None, keepalive=False, attempt=1):
    now = datetime.datetime.utcnow()
    results = {'network': network, 'address': address, 'port': port, 'timestamp': now, 'seen': 0, 'attempt': attempt}

    try:
        handshake_msgs = []
        new_addrs = []

        proxy = None
        if address.endswith(".onion"):
            proxy = CONF['tor_proxy']

        conn = Connection((address, port),
                          (CONF['source_address'], 0),
                          magic_number=network_data['magic_number'],
                          socket_timeout=CONF['socket_timeout'],
                          proxy=proxy,
                          protocol_version=int(network_data['protocol_version']),
                          min_protocol_version=network_data['min_protocol_version'],
                          to_services=int(to_services),
                          from_services=int(from_services or network_data['services']),
                          user_agent=user_agent or CONF['user_agent'],
                          height=int(network_data['height']),
                          relay=CONF['relay'])

        try:
            conn.open()
        except (ProtocolError, ConnectionError, socket.error) as err:
            results['error'] = str(err)
            logging.debug("connection failed %s %s", type(err), err)
        else:
            try:
                handshake_msgs = conn.handshake()
                assert handshake_msgs
                results['seen'] = 1
                results['height'] = int(handshake_msgs[0]['height'])
                results['version'] = int(handshake_msgs[0]['version'])
                results['user_agent'] = handshake_msgs[0]['user_agent'].decode()
                results['services'] = int(handshake_msgs[0]['services'])
            except (ProtocolError, ConnectionError, socket.error, AssertionError) as err:
                results['error'] = str(err)
                logging.debug("handshake failed %s", err)

            msgs = []
            if len(handshake_msgs) > 0 and (p2p_nodes or explicit_p2p):
                getaddr = True
                chance = CONF['getaddr_prop']
                if chance < 1.0 and p2p_nodes and not explicit_p2p and not "--seed" in sys.argv:
                    if np.random.rand() > chance:
                        getaddr = False

                if getaddr:
                    try:
                        conn.getaddr(block=False)
                        msgs = msgs + conn.get_messages(commands=[b"addr"])
                        time.sleep(5)
                        msgs = msgs + conn.get_messages(commands=[b"addr"])
                    except (ProtocolError, ConnectionError, socket.error) as err:
                        logging.debug("getaddr failed %s", err)
            if keepalive:
                Keepalive(conn, 10).keepalive(addr=True if p2p_nodes else False)
            for msg in msgs:
                if msg['count'] > 1:
                    ts = now.timestamp()
                    for addr in msg['addr_list']:
                        if ts - addr['timestamp'] < 12 * 60 * 60:  # within 6 hours
                            new_addrs.append(addr)
        conn.close()
        return results, new_addrs
    except Exception as err:
        logging.warning("unspecified connection error: %s", err)
        return {}, []

def get_seeds(port, dns_seeds, address_seeds, default_services=0):
    """
    Initializes a list of reachable nodes from DNS seeders and hardcoded nodes to bootstrap the crawler.
    """
    exportList = []
    for seeder in dns_seeds:
        nodes = []

        try:
            ipv4_nodes = socket.getaddrinfo(seeder, None, socket.AF_INET)
        except socket.gaierror:
            if CONF['ipv6']:
                try:
                    ipv6_nodes = socket.getaddrinfo(seeder, None, socket.AF_INET6)
                except socket.gaierror as err:
                    logging.warning("%s %s", seeder, err)
                else:
                    nodes.extend(ipv6_nodes)
        else:
            nodes.extend(ipv4_nodes)

        for node in nodes:
            address = node[-1][0]
            exportList.append((address, port, default_services))

    for address in address_seeds:
        exportList.append((address, port, default_services))

    return exportList

def init_geoip():
    asn = Reader("geoip/GeoLite2-ASN.mmdb")
    country = Reader("geoip/GeoLite2-Country.mmdb")
    city = Reader("geoip/GeoLite2-City.mmdb")
    return country, city, asn

def init_db():
    if "sqlite:/" in SQLALCHEMY_DATABASE_URI:
        engine = create_engine(SQLALCHEMY_DATABASE_URI, connect_args={'timeout': 15}, echo=False)
    else:
        engine = create_engine(SQLALCHEMY_DATABASE_URI, echo=False)
    Base.metadata.create_all(engine)
    Sess = sessionmaker(bind=engine, autoflush=False)
    return Sess()

def init_crawler(session, networks):
    # Populates list of all known node addresses and block heights
    db_networks = [x[0] for x in session.query(Node.network).distinct().all()]
    node_addresses = {}
    recent_heights = {}
    for network in set(db_networks + networks):
        node_addresses[network] = {"{};{}".format(y.address, y.port) for y in \
                                   session.query(Node.address, Node.port).filter(Node.network == network).all()}
        count = session.query(Node).filter(and_(Node.network == network, Node.last_height != None)).count()
        if count > 0:
            median = session.query(Node.last_height).filter(and_(Node.network == network, Node.last_height != None)) \
                .order_by(Node.last_height).limit(1).offset(count // 2).one()[0]
            if median:
                recent_heights[network] = [median]
        if network not in recent_heights:
            recent_heights[network] = [500000]
    return node_addresses, recent_heights

def check_dns(network_data, nodeAddresses):
    nodes = []
    for network in network_data:
        nc = network_data[network]
        dnsNodeAddrs = get_seeds(nc['port'], nc['dns_seeds'], nc['address_seeds'], default_services=nc['services'])
        for nodeAddr in dnsNodeAddrs:
            if nodeAddr[0] and nodeAddr[1]:
                if not "{};{}".format(nodeAddr[0], nodeAddr[1]) in nodeAddresses[network]:
                    nodeAddresses[network].add("{};{}".format(nodeAddr[0], nodeAddr[1]))
                    newNode = Node(network=network, address=nodeAddr[0], port=int(nodeAddr[1]),
                                   services=int(nodeAddr[2]))
                    nodes.append(newNode)
    return nodes

def prune_nodes(session):
    now = datetime.datetime.utcnow()
    q = session.query(Node)

    # prune old nodes that can't be reached
    conditions = [and_(Node.last_seen == None, Node.first_checked != None,
             Node.first_checked <= now - datetime.timedelta(days=CONF['min_pruning_age']))]

    # prune nodes that haven't been seen in a long time
    if CONF['max_pruning_age'] and CONF['max_pruning_age'] > CONF['min_pruning_age']:
        conditions.append(and_(Node.last_seen != None,
                               Node.last_seen <= now - datetime.timedelta(days=CONF['max_pruning_age'])))
    q = q.filter(or_(*conditions))

    pruned = q.delete()
    if pruned > 0:
        logging.info("Pruned {} nodes".format(pruned))

        # prune visitations that no longer have a parent node
        if CONF['prune_visitations']:
            nq = session.query(NodeVisitation.id) \
                .outerjoin(Node, Node.id == NodeVisitation.parent_id) \
                .filter(Node.address == None)
            deleted = 100000
            while deleted < 100000:
                nvq = session.query(NodeVisitation).filter(NodeVisitation.id.in_([x[0] for x in nq.limit(100000).all()]))
                deleted = nvq.delete(synchronize_session=False)
                logging.info("{} Visitations deleted".format(deleted))

        session.commit()

def calculate_pending_nodes(session, start_time):
    now = datetime.datetime.utcnow()
    # Get a list of all never checked nodes, and nodes that have been checked recently:
    q = session.query(Node)
    q = q.filter(or_(
        Node.first_checked == None,
        Node.last_checked == None,
        # If it hasn't been seen before, check every 60m
        and_(Node.last_seen == None, Node.last_checked != None,
             Node.last_checked < now - datetime.timedelta(minutes=CONF['crawl_interval'] * 24)),
        # If it has been seen in the last 6 hours, check it every 10 minutes
        and_(Node.last_seen != None, Node.last_seen > now - datetime.timedelta(hours=6),
             Node.last_checked < now - datetime.timedelta(minutes=CONF['crawl_interval'])),
        # Otherwise every 60 min
        and_(Node.last_seen != None,
             Node.last_checked < now - datetime.timedelta(minutes=CONF['crawl_interval'] * 6))
    )).filter(not_(and_(Node.last_checked != None, Node.last_checked > start_time)))

    if CONF['crawl_order']:
        case_order = []
        for i in range(len(CONF['crawl_order'])):
            case_order.append((Node.network == CONF['crawl_order'][i], str(i)))
        q = q.order_by(case(case_order, else_=Node.network), Node.seen.desc(), Node.last_checked)
    else:
        q = q.order_by(Node.seen.desc(), Node.last_checked)

    if CONF['max_queue'] > 0:
        count = q.count()
        q = q.limit(CONF['max_queue'])
    else:
        count = q.count()
    if count > CONF['max_queue']:
        logging.info("{} nodes pending".format(count))

    if CONF['database_concurrency']:
        nodes = q.with_for_update().all()
        session.bulk_update_mappings(Node, [
            {'id': x.id, 'last_checked': now} for x in nodes])
        session.commit()
        return nodes

    return q.all()

def process_pending_nodes(node_addresses, node_processing_queue, recent_heights, thread_pool, session, ASN,
                          COUNTRY, CITY):
    futures_dict = {}

    checked_nodes = 0
    seen_nodes = 0
    pending_nodes = 0
    skipped_nodes = 0
    retried_nodes = 0
    found_on_retry = 0
    new_nodes_to_add = []

    for net in recent_heights:
        CONF['networks'][net]['height'] = max(set(recent_heights[net]),
                                              key=recent_heights[net].count)
        recent_heights[net] = [CONF['networks'][net]['height']]

    # Get list of seen IPs and Ports so we don't send a bitcoin magic number to a bitcoin-cash node
    q = session.query(Node.network, Node.address, Node.port, Node.last_seen) \
        .filter(Node.last_seen > datetime.datetime.utcnow() - datetime.timedelta(days=3))

    activeIps = {}
    for x in q.all():
        key = x.address + "|" + str(x.port)
        if key not in activeIps:
            activeIps[key] = (x.network, x.last_seen)
        else:
            # Prioritize bitcoin cash nodes, as its the only client that bans when the wrong magic number is sent
            if x.network == "bitcoin-cash":
                activeIps[key] = (x.network, x.last_seen)
            elif x.last_seen > activeIps[key][1] and activeIps[key][0] != "bitcoin-cash":
                activeIps[key] = (x.network, x.last_seen)

    while node_processing_queue:
        node = node_processing_queue.pop(0)
        if "{}|{}".format(node.address, node.port) in activeIps and activeIps["{}|{}".format(node.address, node.port)][
            0] != node.network:
            node.last_checked = datetime.datetime.utcnow()
            session.add(node)
            skipped_nodes += 1
            continue
        future = thread_pool.submit(connect, node.network, node.address, node.port, node.services,
                                    CONF['networks'][node.network])
        futures_dict["{}|{}|{}".format(node.network, node.address, node.port)] = node, future
        time.sleep(0.001)

    total_to_complete = len(futures_dict)

    while len(futures_dict) > 0:
        time.sleep(1)
        for i in list(futures_dict.keys()):
            if not futures_dict[i][1].done():
                continue

            checked_nodes += 1
            if checked_nodes % 1000 == 0:
                logging.info(" {}%".format(round(checked_nodes / total_to_complete * 100.0, 1)))

            node, future = futures_dict.pop(i)
            result, new_addrs = future.result()
            if not result:
                continue

            if not result['seen']:
                if CONF['retry_threshold'] and CONF['retry_threshold'] > 0 and (not node.seen or (
                        node.last_seen and node.last_seen < datetime.datetime.utcnow() -
                        datetime.timedelta(hours=CONF['retry_threshold']))):
                    pass
                elif result['attempt'] < CONF['retries'] + 1:
                    future = thread_pool.submit(connect, node.network, node.address, node.port, node.services,
                                                CONF['networks'][node.network], attempt=result['attempt'] + 1)
                    futures_dict["{}|{}|{}".format(node.network, node.address, node.port)] = node, future
                    total_to_complete += 1
                    retried_nodes += 1
                    continue
            elif result['seen'] and result['attempt'] > 1:
                found_on_retry += 1

            x = result['timestamp']
            timestamp = datetime.datetime(x.year, x.month, x.day, x.hour, x.minute, x.second, x.microsecond)

            node.last_checked = timestamp
            if node.first_checked is None:
                node.first_checked = timestamp
            if result["seen"] and not any((x.match(result['user_agent']) for x in CONF['excluded_user_agents'])):
                node.version = result['version']
                node.last_seen = timestamp
                node.services = result['services']
                node.user_agent = result['user_agent']
                node.last_height = result['height']
                if node.first_seen is None:
                    node.first_seen = timestamp
                node.seen = True

                seen_nodes += 1
                recent_heights[result['network']].append(result['height'])
                if not result['address'].endswith(".onion"):
                    if node.asn is None or node.aso is None:
                        try:
                            node.aso = ASN.asn(node.address).autonomous_system_organization
                            node.asn = ASN.asn(node.address).autonomous_system_number
                        except AddressNotFoundError:
                            pass
                    if node.country is None or node.city is None:
                        try:
                            node.country = COUNTRY.country(node.address).country.name
                            node.city = CITY.city(node.address).city.name
                        except AddressNotFoundError:
                            pass

            session.add(node)

            if node.seen:
                if not node.id:
                    session.commit()
                vis = NodeVisitation(parent_id=node.id,
                                     user_agent=result['user_agent'] if 'user_agent' in result else None,
                                     success=result["seen"],
                                     timestamp=timestamp,
                                     height=result['height'] if result["seen"] else None)

                session.add(vis)

            if new_addrs:
                for n in new_addrs:
                    addr = n['ipv4'] or n['ipv6'] or n['onion']
                    if not "{};{}".format(addr, n['port']) in node_addresses[result['network']]:
                        pending_nodes += 1
                        node_addresses[result['network']].add("{};{}".format(addr, n['port']))
                        newNode = Node(network=str(result['network']), address=addr, port=int(n['port']),
                                       services=int(n['services']))
                        if CONF['database_concurrency']:
                            new_nodes_to_add.append(newNode)
                        else:
                            session.add(newNode)

    if CONF['database_concurrency']:
        # Get all unchecked nodes and nodes first seen in the past hour,
        # don't insert any new nodes that have already been inserted
        nn = session.query(Node.network, Node.address, Node.port) \
            .filter(or_(Node.first_checked == None,
                        Node.first_checked > datetime.datetime.utcnow() - datetime.timedelta(hours=1))) \
            .with_for_update().all()
        newSet = {'{};{};{}'.format(n.network, n.address, n.port) for n in nn}
        for i in reversed(range(len(new_nodes_to_add))):
            ni = '{};{};{}'.format(new_nodes_to_add[i].network, new_nodes_to_add[i].address, new_nodes_to_add[i].port)
            if ni in newSet:
                del new_nodes_to_add[i]

    session.commit()
    logging.info("Checked {} Nodes, {} Seen, {} More queued up. ({}/{} retry successes, {} skipped x-network nodes)".format(
        checked_nodes - retried_nodes, seen_nodes, pending_nodes, found_on_retry, retried_nodes, skipped_nodes))
    return node_processing_queue, node_addresses

def code_ip_type(inp):
    if ".onion" in inp:
        return "Onion"
    elif "." in inp:
        return "IPv4"
    elif ":" in inp:
        return "IPv6"
    else:
        return "Unknown"

def dump_summary(session):

    q = session.query(Node.id, Node.network, Node.address, Node.port, Node.user_agent, Node.version, Node.asn, Node.aso,
                      Node.country, Node.city, Node.last_seen, Node.last_height) \
        .filter(Node.seen == True) \
        .filter(Node.last_seen >= datetime.datetime.utcnow() - datetime.timedelta(days=7))

    nodes = pd.read_sql(q.statement, q.session.bind)
    nodes[['port', 'version', 'last_height']] = nodes[['port', 'version', 'last_height']].fillna(0)
    nodes = nodes.fillna("")

    if nodes.empty:
        logging.warning("Nodes table is empty, no results to dump")
        return

    ## Exclude user agents
    if CONF['excluded_user_agents']:
        for agent_re in CONF['excluded_user_agents']:
            agent_re = re.compile(agent_re)
            nodes = nodes[~nodes['user_agent'].str.match(agent_re)].copy()

    now = datetime.datetime.utcnow()
    labels = []
    for age, label in [(2, "2h"), (8, "8h"), (24, "24h"), (24 * 7, "7d"), (24 * 30, "30d")]:
        stt = time.time()
        q = session.query(Node.id,
                          func.sum(case([(NodeVisitation.success, 1)], else_=0)).label("success"),
                          func.count(NodeVisitation.parent_id).label("total")) \
            .join(NodeVisitation, Node.id == NodeVisitation.parent_id) \
            .group_by(Node.id) \
            .filter(Node.last_seen > now - datetime.timedelta(hours=age)) \
            .filter(NodeVisitation.timestamp >= now - datetime.timedelta(hours=age))
        df = pd.read_sql(q.statement, q.session.bind)
        df[label] = (df['success'] / df['total']).fillna(0.0)
        nodes = nodes.merge(df[['id', label]], how="left")
        labels.append(label)
        logging.info("done {} in {}".format(label, round(time.time() - stt, 3)))
    nodes = nodes.drop(['id'], 1)
    nodes[labels] = nodes[labels].fillna(0.0).round(3)
    nodes[['network', 'address']] = nodes[['network', 'address']].fillna("")
    nodes['address_type'] = nodes['address'].apply(code_ip_type)

    nodes['network'] = nodes[['network', 'user_agent']].apply(
        lambda x: "bitcoin-sv" if x['network'] == 'bitcoin-cash' and ' SV' in x['user_agent'] else x['network'], axis=1)

    networks = nodes['network'].unique()

    ## Calculate summaries
    summaries = {}
    for network in networks:
        summ_df = nodes[(nodes['network'] == network) &
                        (nodes['last_seen'] > datetime.datetime.utcnow() - datetime.timedelta(
                            hours=8))]
        if summ_df.empty:
            continue

        summaries[network] = {
            "min": int(summ_df['last_height'].fillna(np.inf).min()),
            "max": int(summ_df['last_height'].fillna(0.0).max()),
            "mean": float(summ_df['last_height'].mean()),
            "stdev": float(summ_df['last_height'].std()),
            "med": float(summ_df['last_height'].median()),
            "1q": float(np.percentile(summ_df['last_height'], 25)),
            "3q": float(np.percentile(summ_df['last_height'], 75)),
            "2.5pct": float(np.percentile(summ_df['last_height'], 1)),
            "97.5pct": float(np.percentile(summ_df['last_height'], 99)),
            "age_min": nodes[nodes['network'] == network]['last_seen'].min().timestamp(),
            "age_max": summ_df['last_seen'].max().timestamp()
        }
        summaries[network]['iqr'] = summaries[network]['3q'] - summaries[network]['1q']
        summaries[network]['95_range'] = summaries[network]['97.5pct'] - summaries[network]['2.5pct']

    summaries["_timestamp"] = datetime.datetime.utcnow().isoformat()
    with open("static/network_summaries.json", 'w') as f:
        json.dump(summaries, f)

    if CONF['inactive_use_iqr']:
        deviations = {network: summaries[network]['iqr'] * (
            CONF['inactive_threshold'][network] if network in CONF['inactive_threshold'] else
            CONF['inactive_threshold']['default']) for network in networks}
    else:
        deviations = {network: CONF['inactive_threshold'][network] if network in CONF['inactive_threshold'] else
        CONF['inactive_threshold']['default'] for network in networks}

    for i in deviations:
        deviations[i] = (deviations[i], summaries[i]['med'])

    def check_active(height, deviations):
        return True if (deviations[1] - deviations[0]) <= height <= (deviations[1] + deviations[0]) else False

    nodes['is_active'] = nodes[['network', 'last_height']] \
        .apply(lambda x: check_active(x['last_height'], deviations[x['network']]), axis=1)

    if not CONF['export_inactive_nodes']:
        nodes = nodes[nodes['is_active']].copy()

    nodes['last_seen'] = nodes['last_seen'].values.astype(np.int64) // 10 ** 9
    nodes.to_csv("static/data.csv", index=False)

    with open("static/data.txt", "w") as f:
        f.write(space_sep_df(nodes))

    for network in nodes['network'].unique():
        netDF = nodes[nodes['network'] == network].copy()
        netDF = netDF.drop(['network'], 1)

        out_dict = {}
        out_dict['data'] = netDF.to_dict(orient="records")

        netDF.to_csv("static/data_{}.csv".format(network), index=False)
        with open(os.path.join("static", "data_{}.json".format(network)), "w") as f:
            json.dump(out_dict, f)
        with open(os.path.join("static", "data_{}.txt".format(network)), "w") as f:
            f.write(space_sep_df(netDF))

    nodes = nodes.drop(['user_agent', 'version', 'last_height'], 1)
    out_dict = {}
    out_dict['data'] = nodes.to_dict(orient="records")
    with open(os.path.join("static", "data.json"), "w") as f:
        json.dump(out_dict, f)

    ## Write unique addresses only
    group_nets = lambda x: ", ".join(sorted(set(x)))
    nodes = nodes.groupby(by=['address', 'asn', 'aso', 'country', 'city', 'address_type'], as_index=False).agg(
        {"network": group_nets, "2h": "mean", "8h": "mean", "24h": "mean", "7d": "mean", "30d": "mean"})
    nodes.to_csv("static/data_unique.csv", index=False)

    out_dict = {}
    out_dict['data'] = nodes.to_dict(orient="records")
    with open(os.path.join("static", "data_unique.json"), "w") as f:
        json.dump(out_dict, f)
    with open(os.path.join("static", "data_unique.txt"), "w") as f:
        f.write(space_sep_df(nodes))

    for network in networks:
        netDF = nodes[nodes['network'].str.contains(network)]
        netDF = netDF.drop(['network'], 1)
        netDF.to_csv("static/data_{}_unique.csv".format(network), index=False)
        out_dict = {}
        out_dict['data'] = netDF.to_dict(orient="records")
        with open(os.path.join("static", "data_{}_unique.json".format(network)), "w") as f:
            json.dump(out_dict, f)
        with open(os.path.join("static", "data_{}_unique.txt".format(network)), "w") as f:
            f.write(space_sep_df(netDF))

def space_sep_df(df, spacing=3):
    df = df.copy()
    df = pd.DataFrame([df.columns], columns=df.columns).append(df)
    for col in df.columns:
        df[col] = df[col].astype(str)
        max_len = df[col].str.len().max() + spacing
        df[col] = df[col].str.pad(max_len, side="right")
    out_str = "\n".join(("".join((str(row[x + 1]) for x in range(len(df.columns)))) for row in df.itertuples()))
    return out_str

def main(session, seed=False):
    start_time = datetime.datetime.utcnow()
    COUNTRY, CITY, ASN = init_geoip()
    thread_pool = ThreadPoolExecutor(max_workers=CONF['threads'])
    networks = list(CONF['networks'].keys())
    prune_nodes(session)
    node_addresses, recent_heights = init_crawler(session, networks)

    if seed:
        seed_nodes = check_dns(CONF['networks'], node_addresses)
        if seed_nodes:
            for n in seed_nodes:
                session.add(n)
        session.commit()

    node_processing_queue = calculate_pending_nodes(session, start_time)
    while node_processing_queue:
        node_processing_queue, node_addresses = process_pending_nodes(node_addresses, node_processing_queue,
                                                                      recent_heights, thread_pool, session, ASN,
                                                                      COUNTRY, CITY)
        node_processing_queue = calculate_pending_nodes(session, start_time)
    logging.info("Crawling complete in {} seconds".format(round((datetime.datetime.utcnow() - start_time).seconds, 1)))

def dump(session):
    start_time = datetime.datetime.utcnow()
    dump_summary(session)
    generate_historic_data(session)
    logging.info("Results saved in {} seconds".format(round((datetime.datetime.utcnow() - start_time).seconds, 1)))

def generate_historic_data(session):
    networks = [x[0] for x in session.query(Node.network).distinct()]
    sd = session.query(func.min(Node.first_seen)).one()[0]
    start_date = datetime.datetime( sd.year, sd.month, sd.day, sd.hour // CONF['historic_interval'] * CONF['historic_interval'], 0, 0)
    end_date = session.query(func.max(Node.last_seen)).one()[0]

    historic_interval = datetime.timedelta(hours=CONF['historic_interval'])

    last_date = start_date
    while last_date < end_date:
        last_date += historic_interval

    interval_end = start_date + historic_interval
    session.query(CrawlSummary).filter(
        CrawlSummary.timestamp >= (last_date - datetime.timedelta(hours=CONF['historic_interval']*1.5))).delete()
    session.commit()
    while interval_end < end_date:
        if session.query(CrawlSummary).filter(CrawlSummary.timestamp == interval_end).count() >= 1:
            interval_end += historic_interval
            continue
        logging.info("Summarizing period starting with {}".format(interval_end - historic_interval))

        q = session.query(NodeVisitation.parent_id.label("id"),
                          case([(NodeVisitation.user_agent.ilike("% SV%"), 'bitcoin-sv')], else_=Node.network).label("network"),
                          func.max(NodeVisitation.height).label("height")) \
            .join(Node, Node.id == NodeVisitation.parent_id) \
            .filter(NodeVisitation.timestamp >= interval_end - historic_interval) \
            .filter(NodeVisitation.timestamp <= interval_end) \
            .filter(NodeVisitation.success == True) \
            .filter(Node.first_seen <= interval_end) \
            .filter(Node.last_seen >= interval_end - historic_interval) \
            .filter(NodeVisitation.network == None) \
            .group_by(NodeVisitation.parent_id,
                      case([(NodeVisitation.user_agent.ilike("% SV%"), 'bitcoin-sv')], else_=Node.network))
        df = pd.read_sql(q.statement, q.session.bind)

        df['height'] = df['height'].astype(int)
        if not df.empty:
            networks = df['network'].unique()
            def check_active(height, deviations):
                return True if (deviations[1] - deviations[0]) <= height <= (deviations[1] + deviations[0]) else False

            medians = df.groupby(by=['network']).agg({"height": "median"})
            deviations = {network: CONF['inactive_threshold'][network] if network in CONF['inactive_threshold'] else
            CONF['inactive_threshold']['default'] for network in networks}
            for i in list(deviations.keys()):
                if i in medians.index:
                    deviations[i] = (deviations[i], medians.loc[i]['height'])
                else:
                    deviations.pop(i)

            df['active'] = df[['network', 'height']].apply(
                lambda x: check_active(x['height'], deviations[x['network']]), axis=1)
            df = df[df['active']].drop(['active'], 1)

        for network in networks:
            netDF = df[df['network'] == network]
            cs = CrawlSummary(timestamp=interval_end,
                              network=network,
                              node_count=len(netDF),
                              lookback_hours=CONF['historic_interval'])

            session.add(cs)
            session.commit()

        interval_end += datetime.timedelta(hours=CONF['historic_interval'])

    q = session.query(CrawlSummary).order_by(CrawlSummary.timestamp)
    df = pd.read_sql(q.statement, q.session.bind)
    df['timestamp'] = df['timestamp'].values.astype(np.int64) // 10 ** 9

    for network in networks:
        df[df['network'] == network][['timestamp', 'node_count']] \
            .to_json("static/history_{}.json".format(network), orient='records')

if __name__ == "__main__":
    session = init_db()

    if "--crawl" in sys.argv:
        main(session, seed=True if "--seed" in sys.argv else False)

    if "--dump" in sys.argv:
        dump(session)

    if "--loop" in sys.argv:
        main(session, seed=True if "--seed" in sys.argv else False)
        while True:
            dump(session)
            main(session, seed=False)

max_stars_repo_name: okdefinet/oknodes

max_issues_repo_name: okdefinet/oknodes

max_forks_repo_name: okdefinet/oknodes


text:
#!/bin/py
#
# lets plot the kinetic energy against
# the derived gravitational potential energy
#
import numpy as np
import matplotlib.pyplot as plt

#
# parameters (see document)
#
rho    = 1.225 # [Kg/m^3]
u      = 5.0   # [m/s]
R      = 3.0   # [m]
delta  = 0.10  # [m]
#
dt     = 30.0  # [Kelvin]
T      = 313   # [Kelvin]
beta   = 1.0/T   # [1/Kelvin]
g      = 9.8   # [m/s^2]

#
# arrays
#
h  = np.linspace(0.0, 100.0)
ke = R*rho*u**3 * (h - 10*delta/11.)
pe = u*beta*rho*dt*g*((h**3)/3.0 - (7*delta**3)/30.)

# plot it
plt.plot(h,ke,'blue',label='Kinetic Energy')
plt.plot(h,pe,'r--',label='Gravitational Potential Energy')
plt.xlabel('Height')
plt.ylabel('Energy')
plt.grid(True)
plt.legend(loc='best')
plt.ylim([0,10000])
plt.xlim([0,20])
plt.savefig('cross.png')
#plt.show()

#
# nick
# 8/19/2015
#

ext: py

max_stars_repo_name: nicholasmalaya/paleologos

max_issues_repo_name: nicholasmalaya/paleologos

max_forks_repo_name: nicholasmalaya/paleologos

reason: import numpy


text:
import tensorflow as tf
import numpy as np
import h5py
import pickle
import sys
import io
import os

# training parameters
training_epochs = 30
batch_size = 64

# architecture parameters
n_hidden = 2048
n_labels = 39   # 39 phones
n_coeffs = 26
n_context_frames = 11   # 5 + 1 + 5
p = 0.5             # keep rate

try:
    nonlinearity_name = sys.argv[1]       # 'relu', 'elu', or 'gelu'
except:
    print('Defaulted to gelu since no nonlinearity specified through command line')
    nonlinearity_name = 'gelu'

try:
    learning_rate = float(sys.argv[2])       # 0.001, 0.0001, 0.00001
except:
    print('Defaulted to a learning rate of 0.001')
    learning_rate = 1e-3

def enumerate_context(i, sentence, num_frames):
    r = range(i-num_frames, i+num_frames+1)
    r = [x if x>=0 else 0 for x in r]
    r = [x if x<len(sentence) else len(sentence)-1 for x in r]
    return sentence[r]

def add_context(sentence, num_frames=11):
    # [sentence_length, coefficients] -> [sentence_length, num_frames, coefficients]

    assert num_frames % 2 == 1, "Number of frames must be odd (since left + 1 + right, left = right)"

    if num_frames == 1:
        return sentence

    context_sent = []

    for i in range(0, len(sentence)):
        context_sent.append([context for context in enumerate_context(i, sentence, (num_frames-1)//2)])

    return np.array(context_sent).reshape((-1, num_frames*n_coeffs))

print('Making graph')
graph = tf.Graph()
with graph.as_default():
    x = tf.placeholder(dtype=tf.float32, shape=[None, n_coeffs*n_context_frames])
    y = tf.placeholder(dtype=tf.int64, shape=[None])
    is_training = tf.placeholder(tf.bool)

    if nonlinearity_name == 'relu':
        f = tf.nn.relu
    elif nonlinearity_name == 'elu':
        f = tf.nn.elu
    elif nonlinearity_name == 'gelu':
        # def gelu(x):
        #     return tf.mul(x, tf.erfc(-x / tf.sqrt(2.)) / 2.)
        # f = gelu
        def gelu_fast(_x):
            return 0.5 * _x * (1 + tf.tanh(tf.sqrt(2 / np.pi) * (_x + 0.044715 * tf.pow(_x, 3))))
        f = gelu_fast
    else:
        raise NameError("Need 'relu', 'elu', 'gelu', for nonlinearity_name")

    W = {
        '1': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_context_frames*n_coeffs, n_hidden]), 0)),
        '2': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_hidden]), 0)),
        '3': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_hidden]), 0)),
        '4': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_hidden]), 0)),
        '5': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_hidden]), 0)),
        '6': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_hidden]), 0)),
        '7': tf.Variable(tf.nn.l2_normalize(tf.random_normal([n_hidden, n_labels]), 0)),
    }
    b = {
        '1': tf.Variable(tf.zeros([n_hidden])),
        '2': tf.Variable(tf.zeros([n_hidden])),
        '3': tf.Variable(tf.zeros([n_hidden])),
        '4': tf.Variable(tf.zeros([n_hidden])),
        '5': tf.Variable(tf.zeros([n_hidden])),
        '6': tf.Variable(tf.zeros([n_hidden])),
        '7': tf.Variable(tf.zeros([n_labels]))
    }

    def feedforward(x):
        h1 = f(tf.matmul(x, W['1']) + b['1'])
        h1 = tf.cond(is_training, lambda: tf.nn.dropout(h1, p), lambda: h1)
        h2 = f(tf.matmul(h1, W['2']) + b['2'])
        h2 = tf.cond(is_training, lambda: tf.nn.dropout(h2, p), lambda: h2)
        h3 = f(tf.matmul(h2, W['3']) + b['3'])
        h3 = tf.cond(is_training, lambda: tf.nn.dropout(h3, p), lambda: h3)
        h4 = f(tf.matmul(h3, W['4']) + b['4'])
        h4 = tf.cond(is_training, lambda: tf.nn.dropout(h4, p), lambda: h4)
        h5 = f(tf.matmul(h4, W['5']) + b['5'])
        h5 = tf.cond(is_training, lambda: tf.nn.dropout(h5, p), lambda: h5)

        return tf.matmul(h5, W['6']) + b['6']

    logits = feedforward(x)
    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits, y))

    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss)

    compute_error = tf.reduce_mean(tf.to_float(tf.not_equal(tf.argmax(logits, 1), y)))

# store future results with previous results
if not os.path.exists("./data/"):
    os.makedirs("./data/")

if os.path.exists("./data/timit_fcn_" + nonlinearity_name + ".p"):
    history = pickle.load(open("./data/timit_fcn_" + nonlinearity_name + ".p", "rb"))
    key_str = str(len(history)//7 + 1)
    history["lr" + key_str] = learning_rate
    history["train_loss" + key_str] = []
    history["val_loss" + key_str] = []
    history["test_loss" + key_str] = []
    history["train_err" + key_str] = []
    history["val_err" + key_str] = []
    history["test_err" + key_str] = []
else:
    history = {
        "lr1": learning_rate,
        'train_loss1': [], 'val_loss1': [], 'test_loss1': [],
        'train_err1': [], 'val_err1': [], 'test_err1': []
    }
    key_str = '1'

print('Loading Data')
data = h5py.File("./data/train.h5")
X_train = data['X'][()]
Y_train = data['y'][()]
train_idxs = data['start_idx'][()]

train_mean = np.mean(X_train, axis=(0,1))
train_std = np.std(X_train, axis=(0,1))
X_train -= train_mean
X_train /= (train_std + 1e-11)

data = h5py.File("./data/dev.h5")
X_val = data['X'][()] - train_mean
Y_val = data['y'][()]
val_idxs = data['start_idx'][()]
X_val -= train_mean
X_val /= (train_std + 1e-11)

data = h5py.File("./data/core.h5")
X_test = data['X'][()] - train_mean
Y_test = data['y'][()]
test_idxs = data['start_idx'][()]
X_test -= train_mean
X_test /= (train_std + 1e-11)
del data
print('Number of training examples', X_train.shape[0])
print('Number of validation examples', X_val.shape[0])
print('Number of testing examples', X_test.shape[0])

with tf.Session(graph=graph) as sess:
    sess.run(tf.initialize_all_variables())

    num_batches = X_train.shape[0] // batch_size
    save_every = num_batches//5        # save training information 5 times per epoch

    for epoch in range(training_epochs):
        # shuffle data
        indices = np.arange(X_train.shape[0])
        np.random.shuffle(indices)
        X_train = X_train[indices]
        Y_train = Y_train[indices]
        train_idxs = train_idxs[indices]

        for i in range(num_batches):
            offset = i * batch_size
            _bx, mask_x, _by = X_train[offset:offset+batch_size], train_idxs[offset:offset+batch_size], Y_train[offset:offset+batch_size]

            bx, by = [], []
            for j in range(_bx.shape[0]):
                sentence_frames = add_context(_bx[j][mask_x[j]:])
                bx.append(sentence_frames)
                by.append(_by[j][mask_x[j]:])

            bx, by = np.concatenate(bx), np.concatenate(by)

            if p < 1-1e-5:   # we want to know how the full network is being optimized instead of the reduced version
                l, err = sess.run([loss, compute_error], feed_dict={x: bx, y: by, is_training: False})

            _, l_drop, err_drop = sess.run([optimizer, loss, compute_error], feed_dict={x: bx, y: by,
                                                                                is_training: True})

            if p < 1-1e-5:
                history["train_loss" + key_str].append(l)
                history["train_err" + key_str].append(err)
            else:
                history["train_loss" + key_str].append(l_drop)
                history["train_err" + key_str].append(err_drop)

            if i % save_every == 0:
                err_total = 0
                loss_total = 0
                for j in range(X_test.shape[0]//batch_size):
                    offset = j * batch_size
                    _bx, mask_x, _by = X_test[offset:offset+batch_size], test_idxs[offset:offset+batch_size], Y_test[offset:offset+batch_size]

                    bx, by = [], []
                    for k in range(_bx.shape[0]):
                        sentence_frames = add_context(_bx[k][mask_x[k]:])
                        bx.append(sentence_frames)
                        by.append(_by[k][mask_x[k]:])

                    bx, by = np.concatenate(bx), np.concatenate(by)

                    l, err = sess.run([loss, compute_error], feed_dict={x: bx, y: by, is_training: False})
                    loss_total += l
                    err_total += err
                history["test_loss" + key_str].append(loss_total/(X_test.shape[0]//batch_size))
                history["test_err" + key_str].append(err_total/(X_test.shape[0]//batch_size))

                err_total = 0
                loss_total = 0
                for j in range(X_val.shape[0]//batch_size):
                    offset = j * batch_size
                    _bx, mask_x, _by = X_val[offset:offset+batch_size], val_idxs[offset:offset+batch_size], Y_val[offset:offset+batch_size]

                    bx, by = [], []
                    for k in range(_bx.shape[0]):
                        sentence_frames = add_context(_bx[k][mask_x[k]:])
                        bx.append(sentence_frames)
                        by.append(_by[k][mask_x[k]:])

                    bx, by = np.concatenate(bx), np.concatenate(by)

                    l, err = sess.run([loss, compute_error], feed_dict={x: bx, y: by, is_training: False})
                    loss_total += l
                    err_total += err
                history["val_loss" + key_str].append(loss_total/(X_val.shape[0]//batch_size))
                history["val_err" + key_str].append(err_total/(X_val.shape[0]//batch_size))

        print('Epoch', epoch+1, 'Complete.', 'Val loss', history["val_loss" + key_str][-1], 'Val error', history["val_err" + key_str][-1])

# save history
pickle.dump(history, open("./data/timit_fcn_" + nonlinearity_name + ".p", "wb"))

max_stars_repo_name: hendrycks/GELUs

max_issues_repo_name: tuozhanjun/GELUs

max_forks_repo_name: tuozhanjun/GELUs


text:
"""
common_cv.py

Define collective variable superclasses (also known as milestone
shapes) that might be used in SEEKR2 calculations.
"""

import numpy as np
from abserdes import Serializer

import seekr2.modules.mmvt_cv as mmvt_cv

class Spherical_cv_anchor(Serializer):
    """
    This object represents an anchor within the concentric spherical
    CV. Used for input purposes only.

    Attributes:
    -----------
    radius : float
        The radius of this spherical anchor in units of nanometers.

    lower_milestone_radius : float
        Optionally define the locations of the milestones for each
        anchor. This is the radius of the lower milestone.

    upper_milestone_radius : float
        Optionally define the locations of the milestones for each
        anchor. This is the radius of the lower milestone.

    starting_amber_params : Amber_params or None
        If Amber inputs are used for this anchor, this object contains
        the necessary inputs to start a new simulation.

    starting_forcefield_params : Forcefield_params or None
        If Forcefield XML inputs are used for this anchor, this object
        contains the necessary inputs to start a new simulation.

    bound_state : bool
        Whether this anchor represents the bound state of a ligand-
        receptor system.

    bulk_anchor : bool
        Whether this anchor acts as a bulk state of a ligand-receptor
        system.
    """

    def __init__(self):
        self.radius = 0.0
        self.lower_milestone_radius = None
        self.upper_milestone_radius = None
        self.starting_amber_params = None
        self.starting_forcefield_params = None
        self.bound_state = False
        self.bulk_anchor = False
        self.connection_flags = []

    def check(self, j, cv_input):
        if self.lower_milestone_radius is not None:
            assert j > 0, "lower_milestone_radius must be None for lowest "\
                "anchor in cv."
            assert self.lower_milestone_radius \
                == cv_input.input_anchors[j-1].upper_milestone_radius,\
                "If lower_milestone_radius is defined for anchor "\
                "{}, the anchor below (number {}).".format(j, j-1)\
                +" must have a corresponding upper_milestone_radius."

        if self.upper_milestone_radius is not None:
            assert j < len(cv_input.input_anchors), \
                "upper_milestone_radius must be None for highest anchor "\
                "in cv."
            assert self.upper_milestone_radius \
                == cv_input.input_anchors[j+1].lower_milestone_radius,\
                "If upper_milestone_radius is defined for anchor "\
                "{} at value {:.3f}, the anchor above ".format(j,
                self.upper_milestone_radius)\
                +"(number {}).".format(j+1)\
                +" must have a corresponding lower_milestone_radius, "\
                "current value: {:.3f}.".format(
                    cv_input.input_anchors[j+1].lower_milestone_radius)
        return

    def get_variable_value(self):
        return self.radius

class Spherical_cv_input(Serializer):
    """
    Inputs by the user resulting in concentric spherical anchors
    with milestones and the collective variable (CV).

    Attributes:
    -----------
    index : int
        The index of this CV input object in the Model_input object.

    group1 : list
        A list of ints representing atom indices whose center of mass
        is one end of the CV distance vector.

    group2 : list
        A list of ints representing atom indices whose center of mass
        is the other end of the CV distance vector.

    input_anchors : list
        A list of Spherical_cv_anchor objects which specify inputs for
        the spherical anchors.
    """

    def __init__(self):
        self.index = 0
        self.group1 = []
        self.group2 = []
        self.bd_group1 = []
        self.bd_group2 = []
        self.input_anchors = []
        return

    def read_plain_input(self, inputs):
        """
        Read a plain input file (as opposed to an XML)
        """

        raise Exception("Reading a plain text file is not yet implemented. "\
                        "Only an XML CV input may be read at this time.")
        return

    def check(self):
        """
        Check user inputs to ensure they have been entered properly.
        """

        last_radius = -1e9
        found_bulk_anchor = False
        for i, input_anchor in enumerate(self.input_anchors):
            radius = input_anchor.radius
            assert radius >= 0.0, "A radius must be greater than "\
                "or equal to zero."
            assert radius > last_radius, "Each subsequent radius "\
                "argument must be greater than the last (sorted)."

            if input_anchor.bound_state is None:
                input_anchor.bound_state = False

            assert input_anchor.bound_state in [True, False], \
                "bound_state must be a boolean"

            if input_anchor.bulk_anchor is None:
                input_anchor.bulk_anchor = False

            assert input_anchor.bulk_anchor in [True, False], \
                "bulk_anchor must be a boolean"

            if input_anchor.bulk_anchor:
                assert not found_bulk_anchor, "Only one bulk anchor allowed "\
                    "per set of anchors in a CV."
                found_bulk_anchor = False
            else:
                assert not found_bulk_anchor, "Only the outermost anchor "\
                    "should be the bulk anchor."

            if i > 0:
                assert not input_anchor.bound_state, "Only the lowest"\
                    "anchor can be the bound state."

            assert len(self.input_anchors) > 1, "A CV must contain "\
                "more than one anchor."
        return

    def make_mmvt_milestoning_objects(self, milestone_alias, milestone_index,
                                      input_anchor_index, anchor_index):
        milestones, milestone_alias, milestone_index = \
            mmvt_cv.make_mmvt_milestoning_objects_spherical(
            self, milestone_alias, milestone_index, input_anchor_index,
            anchor_index, self.input_anchors)
        return milestones, milestone_alias, milestone_index

class Tiwary_cv_anchor(Serializer):
    """
    This object represents an anchor within a Tiwary-style CV,
    which is composed of a linear sum of order parameters. Used
    for input purposes only.

    Attributes:
    -----------
    value : float
        The value of the anchor (center of Voronoi cell)

    lower_milestone_value : float
        Optionally define the locations of the milestones for each
        anchor. This is the value of the lower milestone.

    upper_milestone_value : float
        Optionally define the locations of the milestones for each
        anchor. This is the value of the lower milestone.

    starting_amber_params : Amber_params or None
        If Amber inputs are used for this anchor, this object contains
        the necessary inputs to start a new simulation.

    starting_forcefield_params : Forcefield_params or None
        If Forcefield XML inputs are used for this anchor, this object
        contains the necessary inputs to start a new simulation.

    bound_state : bool
        Whether this anchor represents the bound state of a ligand-
        receptor system.

    bulk_anchor : bool
        Whether this anchor acts as a bulk state of a ligand-receptor
        system.
    """

    def __init__(self):
        self.value = 0.0
        self.lower_milestone_value = None
        self.upper_milestone_value = None
        self.starting_amber_params = None
        self.starting_forcefield_params = None
        self.bound_state = False
        self.bulk_anchor = False
        self.connection_flags = []
        return

    def check(self, j, cv_input):
        if self.lower_milestone_value is not None:
            assert j > 0, "lower_milestone_value must be None for lowest "\
                "anchor in cv."
            assert self.lower_milestone_value \
                == cv_input.input_anchors[j-1].upper_milestone_value,\
                "If lower_milestone_value is defined for anchor "\
                "{}, the anchor below (number {}).".format(j, j-1)\
                +" must have a corresponding upper_milestone_value."

        if self.upper_milestone_value is not None:
            assert j < len(cv_input.input_anchors), \
                "upper_milestone_value must be None for highest anchor "\
                "in cv."
            assert self.upper_milestone_value \
                == cv_input.input_anchors[j+1].lower_milestone_value,\
                "If upper_milestone_value is defined for anchor "\
                "{} at value {:.3f}, the anchor above ".format(j,
                self.upper_milestone_value)\
                +"(number {}).".format(j+1)\
                +" must have a corresponding lower_milestone_value, "\
                "current value: {:.3f}.".format(
                    cv_input.input_anchors[j+1].lower_milestone_value)
        return

    def get_variable_value(self):
        return self.value

class Tiwary_cv_distance_order_parameter(Serializer):
    """
    An order parameter object for a Tiwary CV the represents the
    distance between two groups of atoms.
    """

    def __init__(self):
        self.index = 0
        self.group1 = []
        self.group2 = []
        return

    def get_openmm_expression(self, group_index):
        """

        """
        return "distance(g{0}, g{1})".format(group_index, group_index+1)

    def get_num_groups(self):
        """

        """
        return 2

    def get_group(self, index):
        """

        """
        if index == 0:
            return self.group1
        elif index == 1:
            return self.group2
        else:
            raise Exception("Invalid index: {}".format(index))
        return

    def get_value(self, coms):
        """

        """
        assert len(coms) == 2
        distance = np.linalg.norm(coms[1]-coms[0])
        return distance

class Tiwary_cv_angle_order_parameter(Serializer):
    """
    An order parameter object for a Tiwary CV the represents the
    distance between two groups of atoms.
    """

    def __init__(self):
        self.index = 0
        self.group1 = []
        self.group2 = []
        self.group3 = []
        return

    def get_openmm_expression(self, group_index):
        """

        """
        return "angle(g{0}, g{1}, g{2})".format(group_index, group_index+1,
                                                group_index+2)

    def get_num_groups(self):
        """

        """
        return 3

    def get_group(self, index):
        """

        """
        if index == 0:
            return self.group1
        elif index == 1:
            return self.group2
        elif index == 2:
            return self.group3
        else:
            raise Exception("Invalid index: {}".format(index))

    def get_value(self, coms):
        """

        """
        assert len(coms) == 3
        vec1 = coms[0] - coms[1]
        vec2 = coms[2] - coms[1]
        angle = np.arccos(np.dot(vec1/np.linalg.norm(vec1),
                                 vec2/np.linalg.norm(vec2)))
        return angle

class Tiwary_cv_torsion_order_parameter(Serializer):
    """
    An order parameter object for a Tiwary CV the represents the
    distance between two groups of atoms.
    """

    def __init__(self):
        self.index = 0
        self.group1 = []
        self.group2 = []
        self.group3 = []
        self.group4 = []
        return

    def get_openmm_expression(self, group_index):
        return "dihedral(g{0}, g{1}, g{2}, g{3})".format(
            group_index, group_index+1, group_index+2, group_index+3)

    def get_num_groups(self):
        return 4

    def get_group(self, index):
        if index == 0:
            return self.group1
        elif index == 1:
            return self.group2
        elif index == 2:
            return self.group3
        elif index == 3:
            return self.group4
        else:
            raise Exception("Invalid index: {}".format(index))

    def get_value(self, coms):
        """

        """
        assert len(coms) == 4
        x1 = coms[0][0]; y1 = coms[0][1]; z1 = coms[0][2]
        x2 = coms[1][0]; y2 = coms[1][1]; z2 = coms[1][2]
        x3 = coms[2][0]; y3 = coms[2][1]; z3 = coms[2][2]
        x4 = coms[3][0]; y4 = coms[3][1]; z4 = coms[3][2]
        vec1 = np.array([x2-x1,y2-y1,z2-z1])
        axis = np.array([x3-x2,y3-y2,z3-z2])
        vec2 = np.array([x4-x3,y4-y3,z4-z3])
        cross1 = np.cross(vec1,axis)
        cross2 = np.cross(axis,vec2)
        cross1 = cross1/np.linalg.norm(cross1)
        cross2 = cross2/np.linalg.norm(cross2)
        x = np.dot(cross1, cross2)
        y = np.dot(np.cross(cross1, axis/np.linalg.norm(axis)), cross2)
        phi = -np.arctan2(y,x) % (2.0*np.pi)
        return phi

class Tiwary_cv_input(Serializer):
    """
    Inputs by the user resulting in Tiwary-style anchors
    with milestones and the collective variable (CV).

    TODO: update docstring

    Attributes:
    -----------
    index : int
        The index of this CV input object in the Model_input object.

    group1 : list
        A list of ints representing atom indices whose center of mass
        is one end of the CV distance vector.

    group2 : list
        A list of ints representing atom indices whose center of mass
        is the other end of the CV distance vector.

    input_anchors : list
        A list of Spherical_cv_anchor objects which specify inputs for
        the spherical anchors.
    """

    def __init__(self):
        self.index = 0
        self.order_parameters = []
        self.order_parameter_weights = []
        self.input_anchors = []
        return

    def read_plain_input(self, inputs):
        """
        Read a plain input file (as opposed to an XML)
        """

        raise Exception("Reading a plain text file is not yet implemented. "\
                        "Only an XML CV input may be read at this time.")
        return

    def check(self):
        """
        Check user inputs to ensure they have been entered properly.
        """

        last_value = -1e9
        num_bulk_anchors = 0
        for i, input_anchor in enumerate(self.input_anchors):
            value = input_anchor.value
            assert value >= 0.0, "A value must be greater than "\
                "or equal to zero."
            assert value > last_value, "Each subsequent value "\
                "argument must be greater than the last (sorted)."

            if input_anchor.bound_state is None:
                input_anchor.bound_state = False

            assert input_anchor.bound_state in [True, False], \
                "bound_state must be a boolean"

            if input_anchor.bulk_anchor is None:
                input_anchor.bulk_anchor = False

            if input_anchor.bulk_anchor:
                num_bulk_anchors += 1

            if i > 0:
                assert not input_anchor.bound_state, "Only the lowest"\
                    "anchor can be the bound state."

            assert len(self.input_anchors) > 1, "A CV must contain "\
                "more than one anchor."

        assert num_bulk_anchors <= 1, "There cannot be more than one bulk "\
            "anchor."
        return

    def make_mmvt_milestoning_objects(self, milestone_alias, milestone_index,
                                      input_anchor_index, anchor_index):
        milestones, milestone_alias, milestone_index = \
            mmvt_cv.make_mmvt_milestoning_objects_tiwary(
            self, milestone_alias, milestone_index, input_anchor_index,
            anchor_index, self.input_anchors)
        return milestones, milestone_alias, milestone_index

class Planar_cv_anchor(Serializer):
    """
    This object represents an anchor within the planar incremental
    CV. Used for input purposes only.

    TODO: update docstring

    Attributes:
    -----------
    radius : float
        The radius of this spherical anchor in units of nanometers.

    lower_milestone_radius : float
        Optionally define the locations of the milestones for each
        anchor. This is the radius of the lower milestone.

    upper_milestone_radius : float
        Optionally define the locations of the milestones for each
        anchor. This is the radius of the lower milestone.

    starting_amber_params : Amber_params or None
        If Amber inputs are used for this anchor, this object contains
        the necessary inputs to start a new simulation.

    starting_forcefield_params : Forcefield_params or None
        If Forcefield XML inputs are used for this anchor, this object
        contains the necessary inputs to start a new simulation.

    bound_state : bool
        Whether this anchor represents the bound state of a ligand-
        receptor system.

    bulk_anchor : bool
        Whether this anchor acts as a bulk state of a ligand-receptor
        system.
    """

    def __init__(self):
        self.value = 0.0
        self.lower_milestone_value = None
        self.upper_milestone_value = None
        self.starting_amber_params = None
        self.starting_forcefield_params = None
        self.bound_state = False
        self.bulk_anchor = False
        self.connection_flags = []
        return

    def check(self, j, cv_input):
        if self.lower_milestone_value is not None:
            assert j > 0, "lower_milestone_value must be None for lowest "\
                "anchor in cv."
            assert self.lower_milestone_value \
                == cv_input.input_anchors[j-1].upper_milestone_value,\
                "If lower_milestone_value is defined for anchor "\
                "{}, the anchor below (number {}).".format(j, j-1)\
                +" must have a corresponding upper_milestone_value."

        if self.upper_milestone_value is not None:
            assert j < len(cv_input.input_anchors), \
                "upper_milestone_value must be None for highest anchor "\
                "in cv."
            assert self.upper_milestone_value \
                == cv_input.input_anchors[j+1].lower_milestone_value,\
                "If upper_milestone_value is defined for anchor "\
                "{} at value {:.3f}, the anchor above ".format(j,
                self.upper_milestone_value)\
                +"(number {}).".format(j+1)\
                +" must have a corresponding lower_milestone_value, "\
                "current value: {:.3f}.".format(
                    cv_input.input_anchors[j+1].lower_milestone_value)
        return

    def get_variable_value(self):
        return self.value

class Planar_cv_input(Serializer):
    """
    Inputs by the user resulting in planar incremental anchors
    with milestones and the collective variable (CV).

    Attributes:
    -----------
    index : int
        The index of this CV input object in the Model_input object.

    group1 : list
        A list of ints representing atom indices whose center of mass
        is one end of the CV distance vector.

    group2 : list
        A list of ints representing atom indices whose center of mass
        is the other end of the CV distance vector.

    input_anchors : list
        A list of Spherical_cv_anchor objects which specify inputs for
        the spherical anchors.
    """

    def __init__(self):
        self.index = 0
        self.start_group = []
        self.end_group = []
        self.mobile_group = []
        self.input_anchors = []
        return

    def read_plain_input(self, inputs):
        """
        Read a plain input file (as opposed to an XML)
        """

        raise Exception("Reading a plain text file is not yet implemented. "\
                        "Only an XML CV input may be read at this time.")
        return

    def check(self):
        """
        Check user inputs to ensure they have been entered properly.
        """

        last_value = -1e9
        num_bulk_anchors = 0
        for i, input_anchor in enumerate(self.input_anchors):
            value = input_anchor.value
            assert value > last_value, "Each subsequent value "\
                "argument must be greater than the last (sorted)."

            if input_anchor.bound_state is None:
                input_anchor.bound_state = False

            assert input_anchor.bound_state in [True, False], \
                "bound_state must be a boolean"

            if input_anchor.bulk_anchor is None:
                input_anchor.bulk_anchor = False

            if input_anchor.bulk_anchor:
                num_bulk_anchors += 1

            if i > 0:
                assert not input_anchor.bound_state, "Only the lowest"\
                    "anchor can be the bound state."

            assert len(self.input_anchors) > 1, "A CV must contain "\
                "more than one anchor."
        assert num_bulk_anchors <= 1, "There cannot be more than one bulk "\
            "anchor."

        return

    def make_mmvt_milestoning_objects(self, milestone_alias, milestone_index,
                                      input_anchor_index, anchor_index):
        milestones, milestone_alias, milestone_index = \
            mmvt_cv.make_mmvt_milestoning_objects_planar(
            self, milestone_alias, milestone_index, input_anchor_index,
            anchor_index, self.input_anchors)
        return milestones, milestone_alias, milestone_index

class Multidimensional_cv(Serializer):
    """
    An object for the input to define when input CVs should be combined to
    make a multidimensional anchor - an anchor with more than one dimension
    of milestones.
    """
    def __init__(self):
        self.cv_inputs = []

        return


text:
#!/usr/bin/env python
#################################################################################
# Copyright 2018 ROBOTIS CO., LTD.
#
# 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.
#################################################################################

# Authors: Gilbert #

import rospy
import os
import json
import numpy as np
import random
import time
import datetime
import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
from collections import deque
from std_msgs.msg import Float32MultiArray
# from src.turtlebot3_ddpg.environment_stage_4_torch_ddpg import Env
from src.turtlebot3_ddpg.ddpg_model import Actor,Critic
from src.turtlebot3_ddpg.data_logger import DataLogger
# from src.turtlebot3_ddpg.randomise_action import OrnsteinUhlenbeckProcess
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter

# set the random seed:
torch.manual_seed(1000)
np.random.seed(1000)
random.seed(1000)

# initialize self.tensorboard
current_time_global = datetime.datetime.now().strftime("%Y_%m_%d-%H:%M:%S")

class ReinforceAgent():
    def __init__(self, env, state_size, action_size, stage, method, mode, current_time = current_time_global):

        ############## Init Parameters ##############
        self.env = env
        self.stage = str(stage)
        self.method = method
        self.mode = mode
        self.current_time = current_time
        self.state_size = state_size
        self.action_size = action_size

        ############## Publisher ##############
        self.pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        self.result = Float32MultiArray()

        ############## Path Parameters ##############
        self.dirPath = os.path.dirname(os.path.realpath(__file__))
        self.dirPath = self.dirPath.replace('turtlebot3_ddpg/nodes', 'turtlebot3_ddpg/save_model_ddpg/')
        self.model_path = self.dirPath + self.current_time + '/' + 'stage_'+self.stage+'_'
        self.log_dir = self.dirPath + 'logs_ddpg/' + 'stage_'+self.stage+'_' + self.current_time
        self.tensorboard = SummaryWriter(log_dir=self.log_dir)

        ############## Train Parameters ##############
        self.n_iterations = 7000
        self.episode_step = 5000
        self.e = 0
        self.episode_step = 10000
        self.discount_factor = 0.99
        self.learning_rate = 0.00025
        self.batch_size = 64
        self.tau = 0.001
        self.warmup = 10
        self.steps = 1
        self.save_model_freq = 50
        self.buffer_memory = 1000000
        self.buffer = []
        self.batch = []

        ############## Epsilon Parameters ##############
        self.max_steps = 1000000
        self.annealing_steps = 100000
        self.start_epsilon = 1
        self.end_epsilon_1 = 0.1
        self.end_epsilon_2 = 0.01
        self.epsilon = self.start_epsilon
        self.slope1 = -(self.start_epsilon - self.end_epsilon_1)/self.annealing_steps
        self.constant1 = self.start_epsilon
        self.slope2 = -(self.end_epsilon_1 - self.end_epsilon_2)/(self.max_steps - self.annealing_steps)
        self.constant2 = self.end_epsilon_2 - self.slope2*self.max_steps

        ############## Random_Noise Parameters ##############
        # self.ou_theta = 0.15
        # self.ou_mu = 0.2
        # self.ou_sigma = 0.5
        # self.random_noise = OrnsteinUhlenbeckProcess(size=self.action_size, theta=self.ou_theta, mu=self.ou_mu, sigma=self.ou_sigma)

        # print(self.dirPath + self.current_time)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print('Device:',self.device)

        ############## Initialise Networks ##############
        self.actor = Actor(self.state_size,self.action_size).to(self.device)
        self.critic = Critic(self.state_size,self.action_size).to(self.device)

        self.target_actor = Actor(self.state_size,self.action_size).to(self.device)
        self.target_critic = Critic(self.state_size,self.action_size).to(self.device)

        self.target_actor.load_state_dict(self.actor.state_dict())
        self.target_critic.load_state_dict(self.critic.state_dict())

        self.criteria = nn.MSELoss()
        self.actor_optimiser = optim.Adam(self.actor.parameters(),self.learning_rate)
        self.critic_optimiser = optim.Adam(self.critic.parameters(),self.learning_rate)

        ############## Testing ##############

        if self.mode == "test":
            print("############## Testing model ##############")
            actor_weights = torch.load(self.model_path + '3200' + '_actor.pth', map_location = self.device)
            critic_weights = torch.load(self.model_path + '3200' + '_critic.pth', map_location = self.device)
            self.actor.load_state_dict(actor_weights)
            self.critic.load_state_dict(critic_weights)
            self.actor.eval()
            self.critic.eval()

        ############## Conitnuing ##############
        if self.mode == "cont":
            print("############## Resuming Tarining ##############")
            last_train_memory = torch.load(self.model_path+ 'last_train_memory.tar')
            self.epsilon = last_train_memory['epsilon']
            self.e = last_train_memory['e']
            self.current_time = last_train_memory['current_time']
            self.steps = last_train_memory['steps']
            self.buffer = last_train_memory['buffer']

            last_train_weights = torch.load(self.model_path+ 'last_train_weights.tar', map_location = self.device)
            actor_weights = last_train_weights['actor']
            critic_weights = last_train_weights['critic']
            target_actor_weights = last_train_weights['target_actor']
            target_critic_weights = last_train_weights['target_critic']
            actor_optimiser_weights = last_train_weights['actor_optimiser']
            critic_optimiser_weights = last_train_weights['critic_optimiser']
            self.actor.load_state_dict(actor_weights)
            self.critic.load_state_dict(critic_weights)
            self.target_actor.load_state_dict(target_actor_weights)
            self.target_critic.load_state_dict(target_critic_weights)
            self.actor_optimiser.load_state_dict(actor_optimiser_weights)
            self.critic_optimiser.load_state_dict(critic_optimiser_weights)
            self.actor.train()
            self.critic.train()

    def soft_update(self,model):
        if model == "Critic":
            for current,target in zip(self.critic.parameters(),self.target_critic.parameters()):
                target.data.copy_(target.data *(1 - self.tau) + current.data * self.tau)
        elif model == "Actor":
            for current,target in zip(self.actor.parameters(),self.target_actor.parameters()):
                target.data.copy_(target.data *(1 - self.tau) + current.data * self.tau)

    def getAction(self, state, test = False):
        # process state
        self.actor.eval()
        if not self.mode == "test":
            #self.epsilon-=self.epsilon_decay_step

            if self.steps <=  self.annealing_steps:
              self.epsilon = self.steps*self.slope1 + self.constant1

            elif self.steps > self.annealing_steps:
              self.epsilon = self.steps*self.slope2 + self.constant2

            p = random.random()
            if p < self.epsilon:
            	return self.random_action()

        action = self.actor(torch.from_numpy(np.array([state])).to(self.device)).squeeze(0).detach().cpu().numpy()
        # print(action)
        # action += (1-test)*max(self.epsilon,0)*self.random_noise.sample()
        action = np.clip(action,-1.,1.)
        # print(action)
        return action

    def random_action(self):
        action = np.random.uniform(-1,1,self.action_size)
        # print(action)
        return action

    def appendMemory(self, episode):
        if len(self.buffer) >= self.buffer_memory:
            self.buffer.pop(0)
        self.buffer.append(episode)
        # self.memory.append((state, action, reward, next_state, done))

    def replayBuffer(self):
         # sample from the replay buffer
        batch  = random.sample(self.buffer, self.batch_size)
        self.batch = list(zip(*batch))

    def learn(self):
        # update epsilon value: causing it to decay

        ## sample random minibatch from buffer
        self.replayBuffer()
        self.actor.train()

        ## process parameters
        states = torch.from_numpy(np.asarray(self.batch[0])).to(self.device)
        # print(self.batch[1])
        actions = torch.from_numpy(np.asarray(self.batch[1])).to(self.device)
        # print(actions.size())
        rewards = torch.from_numpy(np.asarray(self.batch[2])).to(self.device)
        # print(rewards.size())
        next_states = torch.from_numpy(np.asarray(self.batch[3])).to(self.device)
        dones = torch.from_numpy(np.asarray(self.batch[4])).to(self.device)

        predicted_q_values = self.critic(states,actions).squeeze(1)
        next_state_actions  = self.target_actor(next_states)
        next_state_q_values = self.target_critic(next_states, next_state_actions).squeeze(1).detach()
        next_state_q_values[dones] = 0

        Y = next_state_q_values*self.discount_factor + rewards
        # print(Y.size())

        critic_loss = self.criteria(predicted_q_values.double(),Y.double())
        self.critic_optimiser.zero_grad()
        critic_loss.backward()
        #torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 1)
        self.critic_optimiser.step()

        actor_loss = -self.critic(states,self.actor(states)).squeeze(1).mean()
        self.actor_optimiser.zero_grad()
        actor_loss.backward()
        # torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 1)
        self.actor_optimiser.step()

        self.soft_update("Critic")
        self.soft_update("Actor")

    def train_model(self):

        rospy.set_param('mode', self.mode)

        rospy.init_node('turtlebot3_ddpg_stage_'+ str(self.stage))
        # initialize result publisher
        pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        # initialize get_action publisher
        pub_get_action = rospy.Publisher('get_action', Float32MultiArray, queue_size=5)
        # set varaibles
        result = Float32MultiArray()
        get_action = Float32MultiArray()

        # set variables
        scores, episodes, total_reward = [], [], []
        global_step = 0
        # set start time
        start_time = time.time()
        prev_e = 0

        # main loop: for each episode
        for e in range(self.e + 1, self.n_iterations):
            done = False
            state = self.env.reset()
            score = 0

            for t in range(self.episode_step):

                if e < self.warmup:
                    action = self.random_action()
                else:
                    action = self.getAction(state)

                self.steps += 1
                self.e = e
                # take action and return state, reward, status
                next_state, reward, done = self.env.step(action)

                # append memory to memory buffer
                self.appendMemory((state, action, reward, next_state, done))

                ## check if replay buffer is ready:
                if e > self.warmup and len(self.buffer) > self.batch_size:
                    self.learn()

                # increment score and append reward
                score += reward
                total_reward.append(reward)

                # update state
                state = next_state

                # publish get_action
                get_action.data = [action[0], score, reward]
                pub_get_action.publish(get_action)

                # save to self.tensorboard
                num = 30
                self.tensorboard.add_scalar('step reward', reward, global_step)
                self.tensorboard.add_scalar('average step reward (over 30 steps)',
                                        sum(total_reward[-num:])/num, global_step)

                # save model after every N episodes
                if e % self.save_model_freq == 0 and e != prev_e:
                    print('Saving Model!')
                    self.save_model()
                    prev_e = e

                # timeout after 1200 steps (robot is just moving in circles or so)
                if t >= 2000: # changed this from 500 to 1200 steps
                    rospy.loginfo("Time out!!")
                    done = True

                if done:
                    result.data = [score, action]
                    pub_result.publish(result)
                    scores.append(score)
                    episodes.append(e)
                    m, s = divmod(int(time.time() - start_time), 60)
                    h, m = divmod(m, 60)

                    rospy.loginfo('Ep: %d | score: %.2f | memory: %d | epsilon: %.6f | time: %d:%02d:%02d',
                                  e, score, len(self.buffer), self.epsilon, h, m, s)
                    # add to self.tensorboard
                    k = 10
                    self.tensorboard.add_scalar('episode reward', score, e)
                    self.tensorboard.add_scalar('average episode reward (over 10 episodes)',
                                        sum(scores[-k:])/k, e)
                    break

                global_step += 1

    def test_model(self, episodes):

        rospy.set_param('mode', self.mode)

        rospy.init_node('turtlebot3_ddpg_stage_'+ str(self.stage))
        # initialize result publisher
        pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        # initialize get_action publisher
        pub_get_action = rospy.Publisher('get_action', Float32MultiArray, queue_size=5)
        # set varaibles
        result = Float32MultiArray()
        get_action = Float32MultiArray()

        logger = DataLogger(self.stage, self.method)

        for t in range(1,episodes+1):

            rospy.set_param('trial_number', t-1)

            goal_count = 0

            rospy.set_param('goal_count', goal_count)

            done = False
            self.env.initGoal = True
            state = self.env.reset()
            score = 0

            episode_done = False

            while not episode_done:
                action = self.getAction(state)

                next_state, reward, done = self.env.step(action)

                score += reward

                state = next_state

                get_action.data = [action[0], score, reward]
                pub_get_action.publish(get_action)

                num = 30
                # print(goal_count)
                logger.store_data()

                if self.env.goal_reached:

                 # goal_count += 1
                    goal_count = rospy.get_param('goal_count')
                    rospy.loginfo('Trial: %d | Goal [%d] completed', t, goal_count)
                    self.env.goal_reached = False
                    if goal_count == 5:
                        episode_done = True
                        logger.save_data(t)

                if done:
                    # if episode/trial fails (i.e. collides)
                    logger.save_data(t, done='fail')
                    break

    def save_model(self):
        # print(self.dirPath + self.current_time)
        if not os.path.isdir(self.dirPath + self.current_time):
            print('Creating Directory')
            os.makedirs(self.dirPath + self.current_time)

        torch.save(self.actor.state_dict(), self.model_path + str(self.e) + '_actor.pth')
        torch.save(self.critic.state_dict(), self.model_path + str(self.e) + '_critic.pth')
        torch.save({
            'epsilon':self.epsilon,
            'steps':self.steps,
            'buffer':self.buffer,
            'e':self.e,
            'current_time':self.current_time
            }
            ,self.model_path+ 'last_train_memory.tar')
        torch.save({
            'actor':self.actor.state_dict(),
            'critic':self.critic.state_dict(),
            'target_actor':self.target_actor.state_dict(),
            'target_critic':self.target_critic.state_dict(),
            'actor_optimiser':self.actor_optimiser.state_dict(),
            'critic_optimiser':self.critic_optimiser.state_dict()
            }
            ,self.model_path+ 'last_train_weights.tar')


text:
import numpy as np
import random
import os
import json

from transformers import AutoModel, AutoTokenizer, AutoConfig

class Config(object):
    def __init__(self, args):
        self.args = args

        # train hyper parameter
        self.multi_gpu = args.multi_gpu
        self.learning_rate = args.lr
        self.batch_size = args.batch_size
        self.max_epoch = args.max_epoch
        self.max_len = args.max_len
        self.model_name = args.model_name

        # dataset
        self.dataset = args.dataset
        self.dataset_lang = args.dataset_lang

        # path and name
        self.root = args.root_path
        self.project_dir = os.path.join(self.root, args.project_dir)
        self.data_path = os.path.join(args.root_path, 'data', self.dataset)
        self.bert_path = args.bert_cache_dir

        # autoconfig BERT
        self.bert_type = args.bert_type
        bert_cfg = AutoConfig.from_pretrained(self.bert_type, cache_dir=self.bert_path)
        self.auto_model = AutoModel.from_pretrained(self.bert_type, config=bert_cfg, cache_dir=self.bert_path)
        self.tokenizer = AutoTokenizer.from_pretrained(self.bert_type, cache_dir=self.bert_path)

        self.checkpoint_dir = self.project_dir + '/checkpoint/' + self.dataset
        self.log_dir = self.project_dir + '/log/' + self.dataset
        self.result_dir = self.project_dir + '/result/' + self.dataset
        self.train_prefix = args.train_prefix
        self.dev_prefix = args.dev_prefix
        self.test_prefix = args.test_prefix
#         self.add_scheduler = args.add_scheduler

        # relation num are read from file

        with open(os.path.join(self.data_path, "rel2id.json"), 'r', encoding='utf8') as wf:
            rel_num = json.load(wf)
            self.rel_num = len(rel_num[0])
            print(len(rel_num[0]))

        self.head_only = args.head_only
        # self.loop_type = args.loop_type
        self.step_dim = args.step_dim
        # self.method = args.method
        self.run_name = args.run_name
        self.model_save_name = args.model_name +"diff_name_" + args.run_name+ '_DATASET_N_' + self.dataset + "_LR_" + str(self.learning_rate) + "_BS_" + str(self.batch_size) + "_head_only_" + str(self.head_only)
        self.log_save_name = 'LOG_' + args.model_name+"diff_name_" + args.run_name + '_DATASET_N_' + self.dataset + "_LR_" + str(self.learning_rate) + "_BS_" + str(self.batch_size) + "_head_only_" + str(self.head_only)
        self.result_save_name = 'RESULT_' + args.model_name+"diff_name_" + args.run_name + '_DATASET_N_' + self.dataset + "_LR_" + str(self.learning_rate) + "_BS_" + str(self.batch_size) + "_head_only_" + str(self.head_only)+ ".json"

        random.seed(0)
        self.step_matrix = np.random.random((2 * self.max_len, self.step_dim))

        # log setting
        self.period = args.period
        self.test_epoch = args.test_epoch

        # debug
        self.debug = args.debug
        if self.debug:
            self.dev_prefix = self.train_prefix
            self.test_prefix = self.train_prefix

max_stars_repo_name: wbqhb/SEPC

max_issues_repo_name: wbqhb/SEPC

max_forks_repo_name: wbqhb/SEPC