import functools
import itertools
import random

from .qtensor.qtree import operators as ops
from .qtensor.qtree.optimizer import Tensor
from .qtensor.optimisation.TensorNet import QtreeTensorNet
from .Optimizer import ParallelVar

random.seed(0)

class ParallelTensorNet(QtreeTensorNet):

    def __init__(self, buckets, data_dict, bra_vars, ket_vars, **kwargs):
        self.measurement_circ = None
        self.measurement_op = None
        super().__init__(buckets, data_dict, bra_vars, ket_vars, **kwargs)

    def slice(self, slice_dict):
        sliced_buckets = self.backend.get_sliced_buckets(
            self.buckets, self.data_dict, slice_dict)
        self.buckets = sliced_buckets
        return self.buckets

    @classmethod
    def circ2buckets(cls, qubit_count, circuit, measurement_circ=None, measurement_op=None, pdict={}, max_depth=None):
        # import pdb
        # pdb.set_trace()
        n_batch = circuit[0][0].gen_tensor().shape[0]
        device = circuit[0][0].gen_tensor().device

        if max_depth is None:
            max_depth = len(circuit)

        data_dict = {}

        # Let's build buckets for bucket elimination algorithm.
        # The circuit is built from left to right, as it operates
        # on the ket ( |0> ) from the left. We thus first place
        # the bra ( <x| ) and then put gates in the reverse order

        # Fill the variable `frame`
        layer_variables = [ParallelVar(qubit, name=f'o_{qubit}')
                        for qubit in range(qubit_count)]
        current_var_idx = qubit_count

        # Save variables of the bra
        bra_variables = [var for var in layer_variables]

        # Initialize buckets
        for qubit in range(qubit_count):
            buckets = [[] for qubit in range(qubit_count)]
        # Place safeguard measurement circuits before and after
        # the circuit
        if measurement_circ == None:
            measurement_circ = [[ops.M(qubit, n_batch=n_batch, device=device, is_placeholder=False) for qubit in range(qubit_count)]]
        else:
            for op in measurement_circ[0]:
                op._parameters['n_batch'] = n_batch

        combined_circ = functools.reduce(
            lambda x, y: itertools.chain(x, y),
            [measurement_circ, reversed(circuit[:max_depth])])

        # Start building the graph in reverse order
        for layer in combined_circ:
            for op in layer:
                # CUSTOM
                # build the indices of the gate. If gate
                # changes the basis of a qubit, a new variable
                # has to be introduced and current_var_idx is increased.
                # The order of indices
                # is always (a_new, a, b_new, b, ...), as
                # this is how gate tensors are chosen to be stored
                variables = []
                current_var_idx_copy = current_var_idx
                min_var_idx = current_var_idx
                for qubit in op.qubits:
                    if qubit in op.changed_qubits:
                        variables.extend(
                            [layer_variables[qubit],
                            ParallelVar(current_var_idx_copy)])
                        current_var_idx_copy += 1
                    else:
                        variables.extend([layer_variables[qubit]])
                    min_var_idx = min(min_var_idx,
                                    int(layer_variables[qubit]))

                # fill placeholders in parameters if any
                for par, value in op.parameters.items():
                    if isinstance(value, ops.placeholder):
                        op._parameters[par] = pdict[value]

                data_key = (op.name,
                            hash((op.name, tuple(op.parameters.items()))))
                # Build a tensor
                t = Tensor(op.name, variables,
                        data_key=data_key)

                # Insert tensor data into data dict
                data_dict[data_key] = op.gen_tensor()

                # Append tensor to buckets
                # first_qubit_var = layer_variables[op.qubits[0]]
                buckets[min_var_idx].append(t)

                # Create new buckets and update current variable frame
                for qubit in op.changed_qubits:
                    layer_variables[qubit] = ParallelVar(current_var_idx)
                    buckets.append(
                        []
                    )
                    current_var_idx += 1

        # Finally go over the qubits, append measurement gates
        # and collect ket variables
        ket_variables = []

        if measurement_op==None:
            measurement_op = ops.M(0, n_batch=n_batch, device=device, is_placeholder=False)  # create a single measurement gate object
        else:
            measurement_op._parameters['n_batch'] = n_batch
        data_key = (measurement_op.name, hash((measurement_op.name, tuple(measurement_op.parameters.items()))))
        data_dict.update({data_key: measurement_op.gen_tensor()})

        for qubit in range(qubit_count):
            var = layer_variables[qubit]
            new_var = ParallelVar(current_var_idx, name=f'i_{qubit}', size=2)
            ket_variables.append(new_var)
            # update buckets and variable `frame`
            buckets[int(var)].append(
                Tensor(measurement_op.name,
                    indices=[var, new_var],
                    data_key=data_key)
            )
            buckets.append([])
            layer_variables[qubit] = new_var
            current_var_idx += 1

        return buckets, data_dict, bra_variables, ket_variables, measurement_circ, measurement_op

    @classmethod
    def from_qtree_gates(cls, qc, measurement_circ=None, measurement_op=None, **kwargs):
        all_gates = qc
        n_qubits = len(set(sum([g.qubits for g in all_gates], tuple())))
        qtree_circuit = [[g] for g in qc]
        buckets, data_dict, bra_vars, ket_vars, measurement_circ, measurement_op = cls.circ2buckets(n_qubits, qtree_circuit, measurement_circ, measurement_op)
        tn = cls(buckets, data_dict, bra_vars, ket_vars, **kwargs)
        return tn, measurement_circ, measurement_op


import pipes
import platform
import re
import socket

import eventlet

from libcloud.compute.base import NodeState
from libcloud.compute.providers import get_driver

from settings.libcloud import LibCloudSettings
from util.log import Logged
from verification.pubkey_registrar import PubkeyRegistrar
from virtual_machine import VirtualMachine

class OpenstackClient(object):
	if platform.system() == 'Darwin':
		import libcloud.security
		libcloud.security.VERIFY_SSL_CERT_STRICT = False

	@classmethod
	def get_client(cls):
		credentials = LibCloudSettings.extra_credentials
		credentials['key'] = LibCloudSettings.key
		credentials['secret'] = LibCloudSettings.secret
		return cls.connect(LibCloudSettings.cloud_provider, credentials)

	@classmethod
	def validate_credentials(cls, cloud_provider, credentials):
		try:
			cls.connect(cloud_provider, credentials).list_nodes()
		except:
			return False
		else:
			return True

	@classmethod
	def connect(cls, cloud_provider, credentials):
		return get_driver(cloud_provider)(**credentials)

@Logged()
class OpenstackVm(VirtualMachine):
	VM_USERNAME = 'lt3'

	CloudClient = OpenstackClient.get_client
	Settings = LibCloudSettings

	def __init__(self, vm_id, instance, vm_username=VM_USERNAME):
		super(OpenstackVm, self).__init__(vm_id, instance, vm_username)
		self.store_vm_info()

	@classmethod
	def from_id_or_construct(cls, vm_id, name=None, image=None, instance_type=None, vm_username=VM_USERNAME):
		return cls.from_vm_id(vm_id) or cls.construct(vm_id, name, image, instance_type, vm_username)

	@classmethod
	def construct(cls, vm_id, name=None, image=None, instance_type=None, vm_username=VM_USERNAME):
		if not name:
			name = "koality:%s:%s" % (socket.gethostname(), vm_id)
		if not image:
			image = cls.get_newest_image()
		instance_type = instance_type or cls.Settings.instance_type
		size = cls._get_instance_size(instance_type)

		security_group_name = cls.Settings.security_group
		security_group = cls._validate_security_group(security_group_name)

		cls._delete_instances_with_name(name)

		instance = cls.CloudClient().create_node(name=name, image=image, size=size,
			ex_userdata=cls._default_user_data(vm_username),
			ex_security_groups=[security_group])
		return cls(vm_id, instance, vm_username)

	@classmethod
	def _delete_instances_with_name(cls, instance_name):
		'''Openstack gets mad if you try to create two instances with the same name, so delete any offenders first'''
		while True:
			instances_with_name = filter(lambda instance: instance.name == instance_name, cls.CloudClient().list_nodes())
			if not instances_with_name:
				return
			for instance in instances_with_name:
				cls._safe_terminate(instance)
			eventlet.sleep(3)

	@classmethod
	def _default_user_data(cls, vm_username=VM_USERNAME):
		'''This utilizes Ubuntu cloud-init, which runs this script at "rc.local-like" time
		when it finishes first boot.
		This will fail if we use an image which doesn't utilitize EC2 user_data
		'''
		return '\n'.join(("#!/bin/sh",
			"useradd --create-home %s" % vm_username,
			"mkdir ~%s/.ssh" % vm_username,
			"echo '%s' >> ~%s/.ssh/authorized_keys" % (PubkeyRegistrar().get_ssh_pubkey(), vm_username),
			"chown -R %s:%s ~%s/.ssh" % (vm_username, vm_username, vm_username)))

	@classmethod
	def _validate_security_group(cls, security_group):
		cidr_ip = '%s/32' % socket.gethostbyname(socket.gethostname())
		# TODO: make this more generalized, but ec2metadata is usually provided for openstack clouds
		if platform.system() == 'Darwin':
			own_security_groups = []
		else:
			own_security_groups = cls._call(['ec2metadata', '--security-groups']).output.split('\n')
		group = cls._get_or_create_security_group(security_group)
		for rule in group.rules:
			if (rule.ip_protocol == 'tcp' and
				int(rule.from_port) <= 22 and
				int(rule.to_port) >= 22):
				if rule.ip_range == cidr_ip:
					cls.logger.debug('Found ssh authorization rule on security group "%s" for ip "%s"' % (security_group, cidr_ip))
					return group
				elif rule.group and rule.group['name'] in own_security_groups and rule.group['tenant_id'] == group.tenant_id:
					cls.logger.debug('Found ssh authorization rule on security group "%s" for security group "%s"' % (security_group, rule.group['name']))
					return group
		cls.logger.info('Adding ssh authorization rule to security group "%s" for ip "%s"' % (security_group, cidr_ip))
		cls.CloudClient().ex_create_security_group_rule(group, 'tcp', 22, 22, cidr_ip, None)
		return group

	@classmethod
	def _get_or_create_security_group(cls, security_group):
		client = cls.CloudClient()
		groups = filter(lambda group: group.name == security_group, client.ex_list_security_groups())
		if groups:
			cls.logger.debug('Found existing security group "%s"' % security_group)
			group = groups[0]
		else:
			cls.logger.info('Creating new security group "%s"' % security_group)
			group = client.ex_create_security_group(security_group, 'Auto-generated Koality verification security group')
		return group

	@classmethod
	def from_vm_id(cls, vm_id):
		try:
			vm_info = cls.load_vm_info(vm_id)
		except:
			return None
		else:
			return cls._from_instance_id(vm_id, vm_info['instance_id'], vm_info['username'])

	@classmethod
	def _from_instance_id(cls, vm_id, instance_id, vm_username=VM_USERNAME):
		try:
			client = cls.CloudClient()
			instance = filter(lambda instance: str(instance.id) == str(instance_id), client.list_nodes())[0]
			vm = cls(vm_id, instance, vm_username)
			if vm.instance.state == NodeState.REBOOTING:
				cls.logger.warn("Found VM %s in REBOOTING state" % (vm, instance_id))
				vm.delete()
				return None
			elif vm.instance.state == NodeState.TERMINATED:
				return None
			elif vm.instance.state not in (NodeState.RUNNING, NodeState.PENDING):
				cls.logger.critical("Found VM %s in unexpected %s state.\nState map: %s" % (vm, vm.instance.state, client.NODE_STATE_MAP))
				vm.delete()
				return None
			return vm
		except:
			return None

	@classmethod
	def _get_instance_size(cls, instance_type, matching_attribute='id'):
		return filter(lambda size: getattr(size, matching_attribute) == instance_type, cls.CloudClient().list_sizes())[0]

	def wait_until_ready(self):
		try:
			instance, ip = self.instance.driver.wait_until_running([self.instance], timeout=240, ssh_interface='private_ips')[0]
		except:
			self.logger.warn("VM %s in error state while waiting for startup" % self, exc_info=True)
			self.rebuild()
		else:
			self.instance = instance
			for remaining_attempts in reversed(range(20)):
				if remaining_attempts <= 2:
					self.logger.info("Checking VM %s for ssh access, %s attempts remaining" % (self, remaining_attempts))
				if self.ssh_call("true", timeout=10).returncode == 0:
					return
				eventlet.sleep(3)
			# Failed to ssh into machine, try again
			self.logger.warn("Unable to ssh into VM %s" % self)
			self.rebuild()

	def provision(self, repo_name, private_key, output_handler=None):
		return self.ssh_call("PYTHONUNBUFFERED=true koality-provision %s %s" % (pipes.quote(repo_name), pipes.quote(private_key)),
			timeout=3600, output_handler=output_handler)

	def ssh_args(self):
		options = {
			'LogLevel': 'error',
			'StrictHostKeyChecking': 'no',
			'UserKnownHostsFile': '/dev/null',
			'ServerAliveInterval': '20'
		}
		return self.SshArgs(self.vm_username, self.instance.private_ips[-1], options=options)

	def reboot(self, force=False):
		self.instance.reboot()

	@classmethod
	def get_base_image(cls):
		images = cls.CloudClient().list_images()
		image_id = LibCloudSettings.vm_image_id
		if image_id:
			try:
				return filter(lambda image: image.id == cls.get_image_id(image), images)[0]
			except:
				cls.logger.exception('Invalid image id specified, using default instead')
		return filter(lambda image: image.name == 'koality_verification_0.3', images)[0]

	@classmethod
	def get_available_base_images(cls):
		return cls.CloudClient().list_images()  # This can be REALLY shitty

	@classmethod
	def get_snapshots(cls, base_image):
		return filter(lambda image: re.match(cls.format_snapshot_name(base_image, '.+'), cls.get_image_name(image)), cls.CloudClient().list_images())

	@classmethod
	def get_image_id(cls, image):
		return image.id

	@classmethod
	def get_image_name(cls, image):
		return image.name

	def create_image(self, name, description=None):
		self.instance.driver.ex_save_image(self.instance, name, metadata={'description': description})

	def rebuild(self):
		image = self.get_newest_image()
		size = self.instance.size
		if size is None:
			if 'flavorId' in self.instance.extra:
				flavorId = self.instance.extra['flavorId']
				size = self._get_instance_size(flavorId, 'id')

		instance_name = self.instance.name

		security_group_name = self.Settings.security_group
		security_group = self._validate_security_group(security_group_name)

		self.delete()

		# Use a new connection to avoid connections becoming stale and unauthorized
		client = self.CloudClient()
		old_instance = self.instance
		instance_id = old_instance.id
		while old_instance is not None and old_instance.state != NodeState.TERMINATED:
			instances = filter(lambda instance: instance.id == instance_id, client.list_nodes())
			old_instance = instances[0] if instances else None

		for attempt in xrange(5):
			try:
				self.instance = client.create_node(name=instance_name, image=image, size=size,
					ex_userdata=self._default_user_data(self.vm_username),
					ex_security_groups=[security_group])
			except:
				if attempt < 4:
					eventlet.sleep(3)
				else:
					raise
			else:
				break

		self.store_vm_info()
		self.wait_until_ready()

	def delete(self):
		if self.instance.name:
			self._delete_instances_with_name(self.instance.name)
		else:
			self._safe_terminate(self.instance)
		self.remove_vm_info()
		super(OpenstackVm, self).delete()

	@classmethod
	def _safe_terminate(cls, instance):
		try:
			instance.destroy()
		except:
			cls.logger.info("Failed to terminate instance %s" % instance, exc_info=True)

class SecurityGroups(object):
	CloudClient = OpenstackClient.get_client

	@classmethod
	def get_security_group_names(cls):
		try:
			existing_groups = map(lambda group: group.name, cls.CloudClient().ex_list_security_groups())
		except:
			existing_groups = []
		existing_groups_plus_selected = list(set(existing_groups).union([str(LibCloudSettings.security_group), str(LibCloudSettings._default_security_group)]))
		return sorted(existing_groups_plus_selected)

class InstanceTypes(object):
	CloudClient = OpenstackClient.get_client

	@classmethod
	def set_largest_instance_type(cls, largest_instance_type):
		current_instance_type = LibCloudSettings.instance_type
		LibCloudSettings.largest_instance_type = largest_instance_type

		if current_instance_type not in cls.get_allowed_instance_types():
			LibCloudSettings.instance_type = largest_instance_type

	@classmethod
	def get_allowed_instance_types(cls):
		largest_instance_type = LibCloudSettings.largest_instance_type
		try:
			ordered_types = map(lambda size: size.id, sorted(cls.CloudClient().list_sizes(), key=lambda size: size.ram))
		except:
			ordered_types = []
		if largest_instance_type in ordered_types:
			return ordered_types[:ordered_types.index(largest_instance_type) + 1]
		else:
			return ordered_types


import random
import pygame
import os

from bird import Bird
from pipe import Pipe
from base import Base

def draw_win(win,bird,pipe,base,score):
    win.blit(pygame.transform.scale2x(pygame.image.load(os.path.join("gallery","back.png"))),(0,0))
    for i in pipe:
        i.draw(win)
    base.draw(win)
    #win.blit(base_image,(0,575))
    bird.draw(win)
    pygame.font.init()
    score_font = pygame.font.SysFont("Forte",50)
    score_text = score_font.render("Score: "+str(score),1,(0,0,0))
    win.blit(score_text,(10,10))
    pygame.display.update()

def main():
    bird = Bird(230,350)
    base = Base(730)
    pipes = [Pipe(600)]
    pygame.init()
    win = pygame.display.set_mode((500,800))
    pygame.display.set_caption('Flappy Birds!')
    done = True
    clk = pygame.time.Clock()
    score = 0
    while done:
        clk.tick(30)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                done = False
            if event.type ==pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    bird.jump()
        bird.move()
        add_pipe = False
        del_pipe = []
        for i in pipes:
            if i.collide(bird):
                pygame.quit()
                break
            if i.x+i.pipe_top.get_width()<0:
                del_pipe.append(i)
            if not i.passed and i.x<bird.x:
                i.passed = True
                add_pipe = True
            i.move()
        if add_pipe:
            score += 1
            pipes.append(Pipe(700))
        for j in del_pipe:
            pipes.remove(j)
        if bird.y+bird.img.get_height()>=730 or bird.y<0:
            done = False
            break
        base.move()
        draw_win(win,bird,pipes,base,score)
    pygame.quit()
    quit()
main()


from cx_Freeze import setup, Executable

# Get an error without this
excludes = ["PyQt5.uic"]

# Set requirements
packages = ["idna", "requests", "queue", "PyQt5.QtCore", "PyQt5.QtGui", "PyQt5.QtWidgets", "bs4", "os", "configparser",
            "ast", "threading", "time", "shutil", "sys"]

# Include project files
includes = ["AddResDialog.py", "ConfigEditor.py", "Downloader.py", "SettingsWindow.py", "icon.svg"]

setup(
        name = "Charts Finder",
        version = "1.1",
        description = "Charts Finder",
        author = "Abdullah Radwan",
        options = {"build_exe" : {"packages": packages, "excludes": excludes, "include_files": includes, "include_msvcr": True}},
        executables = [Executable("ChartsFinder.py", base = "Win32GUI")])


# Append current path to path.
import sys
import os
sys.path.append(os.path.dirname(__file__))
sys.path.append(os.path.dirname(__file__) + '/IOPi')
import ConfigParser
import importlib
import json
import logging.config
import logging
import testtimeclock
import testscoreboard
import testshotclock
from bottle import route, request, abort
import bottle

configFileName = os.path.dirname(__file__) + '/digits.cfg'

# ... build or import your bottle application here ...
# Do NOT use bottle.run() with mod_wsgi

@route('/scoreboard/score')
def scoreboardscore():

    try:
        logging.info('scoreboardscore() start')
        score = combi.getJSONScore()
        logging.info('scoreboardscore() end')
        return score
    except:
        logging.exception('Uncaught exception')
        raise

@route('/scoreboard/score-as-array')
def scoreboardscoreasarray():

    try:
        logging.info('scoreboardscoreasarray() start')
        score = combi.getJSONScore()
        data = json.loads(score)
        list = [data]
        score = json.dumps(list)
        logging.info('scoreboardscoreasarray() end')
        return score
    except:
        logging.exception('Uncaught exception')
        raise

@route('/scoreboard/test/inError', method='PUT')
def testtimeclockinError():
    data = request.body.readline()
    if not data:
        abort(400, 'No data received')
    entity = json.loads(data)
    if not entity.has_key('inError'):
        abort(400, 'No inError specified')
    error = entity['inError']
    tsb.inError(error)
    ttc.inError(error)

@route('/scoreboard/test/score', method='PUT')
def testscoreboardsetscore():

    data = request.body.readline()
    if not data:
        abort(400, 'No data received')
    entity = json.loads(data)
    if not entity.has_key('home'):
        abort(400, 'No home specified')
    if not entity.has_key('guest'):
        abort(400, 'No guest specified')
    home = entity['home']
    guest = entity['guest']

    tsb.setScore(home, guest)

@route('/scoreboard/test/score')
def testscoreboardscore():

    return tsb.getJSONScore()

@route('/scoreboard/test/score-as-array')
def testscoreboardscoreasarray():

    score = tsb.getJSONScore()
    data = json.loads(score)
    list = [data]
    score = json.dumps(list)

    return score

@route('/timeclock/time')
def timeclocktime():

    try:
        logging.info('timeclocktime() start')
        time = combi.getJSONTime()
        logging.info('timeclocktime() end')
        return time
    except:
        logging.exception('Uncaught exception')
        raise

@route('/timeclock/time-as-array')
def timeclocktimeasarray():

    try:
        logging.info('timeclocktimeasarray) start')
        time = combi.getJSONTime()
        data = json.loads(time)
        data['second'] = "%02d" % (data['second'],)
        data['minute'] = str(data['minute'])
        list = [data]
        time = json.dumps(list)
        logging.info('timeclocktimeasarray() end')
        return time
    except:
        logging.exception('Uncaught exception')
        raise

@route('/timeclock/test/start', method='PUT')
def testtimeclockStart():

    ttc.start()
    return "Started"

@route('/timeclock/test/stop', method='PUT')
def testtimeclockStop():

    ttc.stop()
    return "Stopped"

@route('/timeclock/test/reset', method='PUT')
def testtimeclockReset():

    ttc.reset()
    return "Reset executed"

@route('/timeclock/test/startstop', method='PUT')
def testtimeclockStartStop():

    ttc.startstop()
    if ttc.getRunning():
        return "Started"
    else:
        return "Stopped"

@route('/timeclock/test/time')
def testtimeclocktime():

    return ttc.getJSONTime()

@route('/timeclock/test/time-as-array')
def testtimeclocktimeasarray():

    time = ttc.getJSONTime()
    data = json.loads(time)
    data['second'] = "%02d" % (data['second'],)
    data['minute'] = str(data['minute'])
    list = [data]
    time = json.dumps(list)

    return time

@route('/timeclock/test/time', method='PUT')
def testscoreboardsetscore():

    data = request.body.readline()
    if not data:
        abort(400, 'No data received')
    entity = json.loads(data)
    if not entity.has_key('minute'):
        abort(400, 'No minue specified')
    if not entity.has_key('second'):
        abort(400, 'No second specified')
    minute = entity['minute']
    second = entity['second']

    ttc.setTime(minute, second)

@route('/shotclock/time')
def shotclocktime():

	try:
		logging.info('shotclocktime() start')
		time = combi.getJSONShotlockTime()
		logging.info('shotclocktime() end')
		return time
	except:
		logging.exception('Uncaught exception')
		raise

@route('/shotclock/time-as-array')
def shotclocktimeasarray():

	try:
		logging.info('shotclockyimeasarray() start')
		time = combi.getJSONShotlockTime()
		data = json.loads(time)
		list= [data]
		time = json.dumps(list)
		logging.info('shotclocktimeasarray() end')
		return time
	except:
		logging.exception('Uncaught exception')
		raise

@route('/shotclock/test/start', method='PUT')
def testshotclockStart():

	ts.start()
	return "Started"

@route('/shotclock/test/stop', method='PUT')
def testshotclockStop():

	ts.stop()
	return "Stopped"

@route('/shotclock/test/reset', method='PUT')
def testshotclockReset():

	ts.reset()
	return "Stopped"

@route('/shotclock/test/inError', method='PUT')
def testshotclockinError():

	data = request.body.readline()
	if not data:
		abort(400, 'No data received')
	entity = json.loads(data)
	if not entity.has_key('inError'):
		abort(400, 'No inError specified')
	error = entity['inError']
	ts.inError(error)

@route('/shotclock/test/time')
def testshotclocktime():

	return ts.getJSONTime()

@route('/shotclock/test/time-as-array')
def testshotclocktimeasarray():

	time = ts.getJSONTime()
	data = json.loads(time)
	list= [data]
	time = json.dumps(list)
	return time

logging.config.fileConfig(os.path.dirname(
    __file__) + '/logger.cfg')  # logfile config

# Read combi from config file
config = ConfigParser.RawConfigParser()
config.read(configFileName)
type = config.get('Combi', 'type')

# Start the configured combi class
module = importlib.import_module(type)
class_ = getattr(module, type.capitalize())

combi = class_(configFileName, True)

tsb = testscoreboard.Testscoreboard()
ttc = testtimeclock.Testtimeclock()
ts = testshotclock.Testshotclock()

logging.info('Started')
application = bottle.default_app()


from django.shortcuts import render, get_object_or_404, redirect
from .forms import SAForm
from .models import SA
from .charts import donut_chart
from django.http import JsonResponse, HttpResponse

from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
from chatterbot import ChatBot
from chatterbot.trainers import ChatterBotCorpusTrainer
from chatterbot.trainers import ListTrainer

class SAVariable:
    sa_nirvantika_view = None
    sa_image_view = None

    sa_positive = None
    sa_negative = None
    sa_neutral = None
    sa_compound = None

    sa_output_positive = None
    sa_output_negative = None
    sa_output_neutral = None
    sa_output_compound = None

    sa_list = []
    sa_list2 = []
    def __init__(self, end_result):
        self.end_result = end_result

bot = ChatBot('Buddy')
bot = ChatBot(
    'Buddy',
    storage_adapter='chatterbot.storage.SQLStorageAdapter',
    database_uri='sqlite:///database.sqlite3'
)

trainer1 = ChatterBotCorpusTrainer(bot)

# Train based on the english corpus
# trainer1.train("chatterbot.corpus.english")

# Train based on english greetings corpus
# trainer1.train("chatterbot.corpus.english.greetings")

# Train based on the english conversations corpus
# trainer1.train("chatterbot.corpus.english.conversations")

trainer = ListTrainer(bot)

trainer.train([
'My name is Nirvantika',
'Hi',
'Hello',
'I need your assistance regarding my order',
'Please, Provide me with your order id',
'I have a complaint.',
'Please elaborate, your concern',
'How long it will take to receive an order ?',
'An order takes 3-5 Business days to get delivered.',
'Okay Thanks',
'No Problem! Have a Good Day!',
'chatterbot.corpus.english',
'chatterbot.corpus.english.greetings',
'chatterbot.corpus.english.conversations',
])

sid = SentimentIntensityAnalyzer()

def sa_view(request):
    form = SAForm(request.POST or None)
    if form.is_valid():
        form.save()
        form = SAForm()

    try:
        SAVariable.sa_positive = None
        SAVariable.sa_negative = None
        SAVariable.sa_neutral = None
        SAVariable.sa_compound = None

        input_value = None
        output_value = None
        image_url_info = None
        nirvantika_response = None

        field_name = 'input_sentence'
        input_id = SA.objects.latest('id')

        field_object = SA._meta.get_field(field_name)
        input_value = str(field_object.value_from_object(input_id))

        output_value  = sid.polarity_scores(input_value)

        labels = ['Positive', 'Negative', 'Neutral', 'Compound']
        size1 = [abs(output_value['pos']), abs(output_value['neg']), abs(output_value['neu']), abs(output_value['compound'])]
        SAVariable.sa_positive = abs(output_value['pos'])
        SAVariable.sa_negative = abs(output_value['neg'])
        SAVariable.sa_neutral = abs(output_value['neu'])
        SAVariable.sa_compound = abs(output_value['compound'])

        SAVariable.sa_list = [SAVariable.sa_positive * 100, SAVariable.sa_negative* 100, SAVariable.sa_neutral * 100]
        colors = ['#ff9999','#66b3ff','#99ff99','#ffcc99']
        print(SAVariable.sa_list)
        import json
        SAVariable.sa_list = json.dumps(SAVariable.sa_list)
        print(SAVariable.sa_list)

        nirvantika_response = bot.get_response(input_value)
        output_value  = sid.polarity_scores(str(nirvantika_response))

        SAVariable.sa_output_positive = abs(output_value['pos'])
        SAVariable.sa_output_negative = abs(output_value['neg'])
        SAVariable.sa_output_neutral = abs(output_value['neu'])
        SAVariable.sa_output_compound = abs(output_value['compound'])

        SAVariable.sa_list2 = [SAVariable.sa_output_positive * 100, SAVariable.sa_output_negative * 100, SAVariable.sa_output_neutral * 100]

        size2 = [abs(output_value['pos']), abs(output_value['neg']), abs(output_value['neu']), abs(output_value['compound'])]

        sizes = [size1, size2]
        donut_chart(sizes, colors, labels)
        image_url_info = "/../../../static/sa/images/sa.png"

        SAVariable.sa_nirvantika_view = nirvantika_response
        SAVariable.sa_image_view = "/../../../static/search/bfs/maze.png"
        # BFS.objects.all().delete()

        SA.objects.latest('id').delete()
    except:
        print("Sentiment Analysis: Error in the try statement.")

    context = {
        'form': form, 'input_text': input_value, 'output_text': output_value,
        'image_url': image_url_info, 'nirvantika': nirvantika_response,
    }
    return render(request, "sentiment_analysis/sa.html", context)

def post_sa_view(request):
    if request.method == "POST" and request.is_ajax():
        sa_view(request)
        return JsonResponse({"success":True}, status=200)
    return JsonResponse({"success":False}, status=400)

def get_sa_view(request):
    if request.method == "GET" and request.is_ajax():
        try:
            final_view = None
            final_view = SAVariable.sa_nirvantika_view
        except:
            print("BFS ERROR: Error in the try session of BFS in view.py")
    return HttpResponse(final_view)

def get_sa_input_view(request):
    if request.method == "GET" and request.is_ajax():
        try:
            image_url_info = None
            image_url_info = SAVariable.sa_list
        except:
            print("BFS ERROR: Error in the try session of BFS in view.py")
    return HttpResponse(image_url_info)

def get_sa_output_view(request):
    if request.method == "GET" and request.is_ajax():
        try:
            image_url_info = None
            image_url_info = SAVariable.sa_list2
        except:
            print("BFS ERROR: Error in the try session of BFS in view.py")
    return HttpResponse(image_url_info)

def get_sa_state_view(request):
    if request.method == "GET" and request.is_ajax():
        try:
            states_view = None
            states_view = SAVariable.sa_list
        except:
            print("BFS ERROR: Error in the try session of BFS in view.py")
    return HttpResponse(states_view)

def get_sa_hash_view(request):
    if request.method == "GET" and request.is_ajax():
        try:
            hash_view = None
            hash_view = SAVariable.hash_value
        except:
            print("BFS ERROR: Error in the try session of BFS in view.py")
    return HttpResponse(hash_view)


import logging
from datetime import datetime
from typing import Optional

class Logger:
    def __init__(self, date: Optional[datetime] = None) -> None:
        self.date = date if date else datetime.now().strftime("%y%m%d")
        self._setup_logging()

    def _setup_logging(self) -> None:
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s %(levelname)s %(name)s %(message)s",
            datefmt="%m-%d %H:%M",
            filename=f"{self.date}.log",
        )

    def __call__(self, msg: str) -> None:
        print(logging.info(msg))


"""
File:           utils.py
Author:         Pavlina Koutecka
Date:           25/02/2020
Description:    This file holds some basic utilities to perform simple features.
                Currently, these utilities are implemented:

                    > time format function - prints time nicely
                    > progress bar visualization function - prints training progress nicely
                    > rename function - renames files in specified folder
                    > find neares function - finds array element with a best match
                    > parsing funcion - parses passing arguments.
"""

import sys
import time
import os
from os import path
import argparse
import numpy as np

last_time = time.time()
begin_time = last_time

def format_time(seconds):
    """
    This function converts given number of seconds into nice format.

    :param seconds: number of seconds to be formatted
    :return: formatted time

    Inspired by https://github.com/JooHyun-Lee/Camelyon17/blob/master/utils.py.
    """

    seconds = seconds
    hours = int(seconds / 3600)
    seconds = seconds - hours * 3600
    minutes = int(seconds / 60)
    seconds = seconds - minutes * 60
    secondsf = int(seconds)
    seconds = seconds - secondsf
    millis = int(seconds * 1000)

    f = ''
    i = 1

    if hours > 0 and i <= 2:
        f += str(hours) + 'h'
        i += 1
    if minutes > 0 and i <= 2:
        f += str(minutes) + 'm'
        i += 1
    if secondsf > 0 and i <= 2:
        f += str(secondsf) + 's'
        i += 1
    if millis > 0 and i <= 2:
        f += str(millis) + 'ms'
        i += 1
    if f == '':
        f = '0ms'
    return f

def visualize_progress(current, total, estimated=None, text=None):
    """
    This function prints progress bar of some running process to standard output.

    :param current: current number
    :param total: total number
    :param estimated: estimated time till the end of process
    :param text: text to be printed with progress bar
    :return:

    Inspired by https://github.com/JooHyun-Lee/Camelyon17/blob/master/utils.py.
    """

    global last_time, begin_time
    if current == 0:
        begin_time = time.time()

    # bar's parameters
    total_len = 60.
    _, term_width = os.popen('stty size', 'r').read().split()
    term_width = int(term_width)

    cur_len = int(total_len * current / total)
    rest_len = int(total_len - cur_len) - 1

    sys.stdout.write(' [')
    for i in range(cur_len):
        sys.stdout.write('=')
    sys.stdout.write('>')
    for i in range(rest_len):
        sys.stdout.write('.')
    sys.stdout.write(']')

    cur_time = time.time()
    step_time = cur_time - last_time
    last_time = cur_time
    total_time = cur_time - begin_time

    L = []
    if text:
        L.append(str(text))
    L.append(' | Step: %s' % format_time(step_time))
    L.append(' | Total: %s' % format_time(total_time))
    if estimated:
        L.append(' | ETA: %s' % format_time(estimated))
    else:
        L.append(' | ETA: Unknown')

    msg = ''.join(L)
    sys.stdout.write(msg)
    for i in range(term_width - int(total_len) - len(msg) - 3):
        sys.stdout.write(' ')

    # go back to the center of the bar
    for i in range(term_width - int(total_len / 2) + 2):
        sys.stdout.write('\b')
    sys.stdout.write(' %d/%d ' % (current + 1, total))

    if current < total - 1:
        sys.stdout.write('\r')
    else:
        sys.stdout.write('\n')

    sys.stdout.flush()

def rename(folder):
    """
    This function renames files in a specified folder.

    :param folder: folder of files that should by renamed
    :return: None
    """

    os.chdir(folder)

    for filename in os.listdir(folder):
        os.rename(filename, filename.replace('.tif', ''))

    # create list of available slides
    files = [f for f in os.listdir(folder) if path.isfile(path.join(folder, f))]
    print(files[0])

def find_nearest(array, value):
    """
    This function finds the array element with the value closest to chosen value and
    returns it index.

    :param array: array to be searched for the closest value
    :param value: value we want to find in array
    :return: index of the element with the closest value
    """

    array = np.asarray(array)
    index = (np.abs(array - value)).argmin()

    return index

def parse_args():
    """
    This function parses passing arguments.

    :return: parsed arguments
    """

    # creating an ArgumentParser object
    parser = argparse.ArgumentParser()

    # process passing arguments and return them
    parser.add_argument('-t', '--type')
    parser.add_argument('-s', '--start')
    parser.add_argument('-e', '--end')
    parser.add_argument('-d', '--data')
    return parser.parse_args()


from typing import Optional
from typing import List
from pathlib import Path
from packaging.version import Version

pkg_version: Optional[str] = None

def get_user_conf():
    from ..system import get_user_conf_path
    from ..util.util import load_yml_conf
    from os.path import exists

    conf_path = get_user_conf_path()
    if exists(conf_path):
        ret = load_yml_conf(conf_path)
        ret["conf_path"] = str(conf_path)
        return ret
    else:
        return None

def get_current_package_version() -> str:
    from pkg_resources import get_distribution

    version = get_distribution("oakvar").version
    return version

def get_package_versions() -> Optional[List[str]]:
    """
    Return available oakvar versions from pypi, sorted asc
    """
    import json
    from requests import get
    from requests.exceptions import ConnectionError
    from ..exceptions import InternetConnectionError

    try:
        r = get("https://pypi.org/pypi/oakvar/json", timeout=(3, None))
    except ConnectionError:
        raise InternetConnectionError()
    if r.status_code == 200:
        d = json.loads(r.text)
        all_vers: List[str] = list(d["releases"].keys())
        all_vers.sort(key=Version)
        return all_vers
    else:
        return None

def input_formats():
    import os
    from ..system import get_modules_dir

    modules_dir = get_modules_dir()
    assert modules_dir is not None
    formats = set()
    d = os.path.join(modules_dir, "converters")
    if os.path.exists(d):
        fns = os.listdir(d)
        for fn in fns:
            if fn.endswith("-converter"):
                formats.add(fn.split("-")[0])
    return formats

def fn_new_exampleinput(d: str) -> Path:
    from pathlib import Path
    import shutil

    fn = "exampleinput"
    ifn = Path(get_packagedir()) / "lib" / "assets" / fn
    ofn = Path(d) / fn
    shutil.copyfile(ifn, ofn)
    return ofn

def create_new_module(module_name: str, module_type: str, outer=None) -> bool:
    from shutil import copytree
    from pathlib import Path
    from ..system import get_modules_dir
    from ..module.cache import get_module_cache

    modules_dir = get_modules_dir()
    if not modules_dir:
        if outer:
            outer.error("modules_dir does not exist. Run `ov system setup`?")
        return False
    module_dir = Path(modules_dir) / (module_type + "s") / module_name
    template_dir = (
        Path(get_packagedir()) / "lib" / "assets" / "module_templates" / module_type
    )
    if not template_dir.exists():
        e = ValueError(
            f"{template_dir} does not exist. Maybe a wrong module_type {module_type}?"
        )
        if outer:
            outer.error(e)
        raise e
    copytree(template_dir, module_dir)
    for fn in module_dir.iterdir():
        new_fn = str(fn).replace("template", module_name)
        fn.rename(new_fn)
    customize_module_template(module_name, module_dir / f"{module_name}.md")
    customize_module_template(module_name, module_dir / f"{module_name}.yml")
    get_module_cache().update_local()
    return True

def customize_module_template(module_name: str, fpath: Path):
    with open(fpath) as f:
        lines = f.readlines()
    with open(fpath, "w") as wf:
        for line in lines:
            new_line = line.replace("MODULE_TITLE", module_name)
            wf.write(new_line)

def recursive_update(d1, d2):
    """
    Recursively merge two dictionaries and return a copy.
    d1 is merged into d2. Keys in d1 that are not present in d2 are preserved
    at all levels. The default Dict.update() only preserved keys at the top
    level.
    """
    import copy

    d3 = copy.deepcopy(d1)  # Copy perhaps not needed. Test.
    for k, v in d2.items():
        if k in d3:
            orig_v = d3[k]
            if isinstance(v, dict):
                if isinstance(orig_v, dict) is False:
                    d3[k] = v
                else:
                    t = recursive_update(d3.get(k, {}), v)
                    d3[k] = t
            else:
                d3[k] = d2[k]
        else:
            d3[k] = v
    return d3

def report_issue():
    import webbrowser

    webbrowser.open("http://github.com/rkimoakbioinformatics/oakvar/issues")

def set_user_conf_prop(key, val):
    import oyaml as yaml
    from pathlib import Path
    from ..system import get_user_conf_path

    if isinstance(val, Path):
        val = str(val)
    user_conf = get_user_conf()
    if not user_conf:
        user_conf = {key: val}
    else:
        user_conf[key] = val
    user_conf_path = get_user_conf_path()
    wf = open(user_conf_path, "w")
    yaml.dump(user_conf, wf, default_flow_style=False)
    wf.close()

def oakvar_version():
    global pkg_version
    if not pkg_version:
        pkg_version = get_current_package_version()
    return pkg_version

def get_packagedir():
    from pathlib import Path

    return Path(__file__).parent.parent.parent.absolute()

def get_platform():
    from platform import platform

    pl = platform()
    if pl.startswith("Windows"):
        pl = "windows"
    elif pl.startswith("Darwin") or pl.startswith("macOS"):
        pl = "macos"
    elif pl.startswith("Linux"):
        pl = "linux"
    else:
        pl = "linux"
    return pl

def get_max_version_supported_for_migration():
    return Version("1.7.0")

def get_latest_package_version() -> Optional[Version]:
    vers = get_package_versions()
    if vers:
        latest_ver = max([Version(v) for v in vers])
    else:
        latest_ver = None
    return latest_ver


import cv2

srcImage = cv2.imread("G:/pythondealPic/originalPic1.jpg")

logo = cv2.imread("G:/pythondealPic/mask2x.png")
# 使用numpy数组切片，获取ROI(region of interest)，也就是感兴趣区域
imgeROI = srcImage[350:350+logo.shape[0], 800:800+logo.shape[1]]

cv2.addWeighted(imgeROI, 0.8, logo, 0.2, 0, imgeROI)
# 竟然等价以下
# srcImage[350:350+logo.shape[1], 800:800+logo.shape[0]] = cv2.addWeighted(imgeROI, 0.5, logo, 0.3, 0)
#将混合后的文件，重新保存为dota-logo.jpg
cv2.imwrite("G:/pythondealPic/qqw.jpg", srcImage)

# cv2.imshow("AddWeight Image", srcImage)
cv2.waitKey(0)


import tensorflow as tf
from keras.models import Sequential
from keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
import pandas as pd
import numpy as np
import chardet
from datetime import datetime
# Charger les données à partir du fichier CSV
filename = '/Users/paulbergeon/Documents/mémoireESSCA/Classeur10.csv'

with open(filename, 'rb') as f:
    result = chardet.detect(f.read())

df = pd.read_csv(filename, encoding=result['encoding'], delimiter=';')
df['Date'] = pd.to_datetime(df['Date'], format='%d/%m/%Y')

# Normaliser le score de sentiment sur une plage de valeurs de 0 à 1
df['Sentiment'] = df['Sentiment'] / 100.0

# Prétraiter les données et les normaliser
scaler = MinMaxScaler(feature_range=(0, 1))
df['Close normalisé'] = scaler.fit_transform(df['Close'].values.reshape(-1, 1))
df['Variation normalisée'] = scaler.fit_transform(df['Variation'].values.reshape(-1, 1))
df['Sentiment normalisé'] = scaler.fit_transform(df['Sentiment'].values.reshape(-1, 1))

# Diviser les données en ensembles d'entraînement et de test
train_size = int(len(df) * 0.7)
train_data, test_data = df[:train_size], df[train_size:]

# Créer les ensembles d'entrée et de sortie pour le modèle LSTM
def create_dataset(dataset, look_back=1):
    X, Y = [], []
    for i in range(len(dataset) - look_back - 1):
        a = dataset[i:(i + look_back), :]
        X.append(a)
        Y.append(dataset[i + look_back, 2])
    return np.array(X), np.array(Y)

look_back = 20
train_X, train_Y = create_dataset(train_data[['Sentiment normalisé', 'Close normalisé', 'Variation normalisée']].values, look_back)
test_X, test_Y = create_dataset(test_data[['Sentiment normalisé', 'Close normalisé', 'Variation normalisée']].values, look_back)

# Créer et entraîner le modèle LSTM
model = Sequential()
model.add(LSTM(50, input_shape=(look_back, 3)))
model.add(Dense(1))
model.compile(loss='mean_squared_error', optimizer='adam')

model.fit(train_X, train_Y, epochs=100, batch_size=1, verbose=2)

# Effectuer des prédictions et évaluer les performances du modèle
test_predict = model.predict(test_X)

# Inverser la normalisation des prédictions
test_predict = scaler.inverse_transform(test_predict)

# Créer un DataFrame pour stocker les prédictions avec leurs dates correspondantes
test_dates = test_data['Date'].iloc[look_back+1:].reset_index(drop=True)
test_predict_df = pd.DataFrame(data=test_predict.flatten(), columns=['Prediction'])
test_predict_df['Date'] = test_dates
test_predict_df = test_predict_df.set_index('Date')

# Inverser la normalisation des vraies valeurs de prix de clôture
test_Y = scaler.inverse_transform([test_Y])[0]

# Créer un DataFrame pour stocker les vraies valeurs avec leurs dates correspondantes
test_y_df = pd.DataFrame(data=test_Y, columns=['Close'])
test_y_df['Date'] = test_dates
test_y_df = test_y_df.set_index('Date')

# Afficher les prédictions avec les vraies valeurs de prix de clôture
print(test_y_df.join(test_predict_df))

# Effectuer des prédictions et évaluer les performances du modèle
test_predict = model.predict(test_X)

test_predict = scaler.inverse_transform(test_predict)
test_Y = scaler.inverse_transform([test_Y])

mae = np.mean(np.abs(test_predict - test_Y))
rmse = np.sqrt(np.mean(np.square(test_predict - test_Y)))

print("MAE:", mae)
print("RMSE:", rmse)
print(test_Y)
print(test_X)
import matplotlib.pyplot as plt

# Tracer les prédictions et les données réelles
plt.plot(test_Y.flatten(), label='Données réelles')
plt.plot(test_predict.flatten(), label='Prédictions')
plt.xlabel('Période')
plt.ylabel('Prix de clôture')
plt.title('Prédictions du modèle LSTM')
plt.legend()
plt.show()


#!/usr/bin/env python
# coding=utf-8
#此文件实现一些检测当中用到的预处理
from torchvision import transforms
import random

class RandomHorizontalFlipWithBBox(object):
    #以一定概率进行水平翻转 同时翻转ground truth box的位置
    def __init__(self,p=0.5):
        self.p = p
        self.transform = transforms.RandomHorizontalFlip(p)
    def __call__(self,img,bbox): #bbox [[xmin,ymin,xmax,ymax]]
        if random.random() < self.p:
            img_w,img_h = img.size
            #print bbox
            for i in bbox:
                temp = i[0]
                i[0] = img_w - i[2] #new_x_min = img_w - x_max
                i[2] = img_w - temp
            #print bbox
            transform = transforms.RandomHorizontalFlip(1)
            img = transform(img)
            return img,bbox
        return img,bbox

class ResizeWithBBox(object):
    #改变图像大小，同时也改变ground truth box的大小
    #如果是元组类型，输出的图片尺寸与output_size匹配，按照(height,width)顺序
    #如果是整数类型，输出的图片按等比例缩放
    def __init__(self,output_size):
        assert isinstance(output_size,(tuple,int))
        self.output_size = output_size

    def __call__(self,img,bbox):
        img_w,img_h = img.size
        if isinstance(self.output_size,int):
            if h>w:
                new_h,new_w = self.output_size*h/w,self.output_size
            else:
                new_h,new_w = self.output_size,self.output_size*w/h
        else:
            new_h,new_w = self.output_size

        new_h,new_w = int(new_h),int(new_w)
        for i in bbox:
            i[0] = int(i[0]*new_w/img_w)
            i[1] = int(i[1]*new_h/img_h)
            i[2] = int(i[2]*new_w/img_w)
            i[3] = int(i[3]*new_h/img_h)
      #  print bbox
        transform = transforms.Resize(self.output_size)
        img = transform(img)
        return img,bbox


import logging
import sys

from pydcop import infrastructure
from pydcop.algorithms import amaxsum
from pydcop.dcop import relations
from pydcop.dcop.objects import Variable
from pydcop.dcop.relations import UnaryFunctionRelation
from pydcop.infrastructure.run import synchronous_single_run
from pydcop.infrastructure.communication import InProcessCommunicationLayer

logging.basicConfig(level=logging.DEBUG)
logging.info('MaxSum smart coloring test')

"""
Sample : Graph coloring with costs :

2 variables:
* x1 with 2 colors [0,1]
* x2 with 3 colors [0, 1, 2]

1 'all_diff' hard constraint: x1 and x2 must choose different colors

A cost is associated with each color:
* for x1:  {0: 0, 1: -3}
* for x2:   {0: 0, 1: -2, 2: -1}
Costs as modelled as unary constraints on x1 and x2.

"""

def distribue_agent_for_all(variables, factors):
    comm = infrastructure.communication.InProcessCommunicationLayer()

    node_agents = []
    for v in variables:
        f_for_variable = [f.name for f in factors if v.name in
                          [i.name for i in f.dimensions]]

        a = infrastructure.Agent('Var_' + v.name, comm)
        a.add_computation(amaxsum.VariableAlgo(v, f_for_variable))
        node_agents.append(a)

    for f in factors:
        a = infrastructure.Agent('Fact_' + f.name, comm)
        a.add_computation(amaxsum.FactorAlgo(f))
        node_agents.append(a)

    return node_agents

def graph_coloring_pb():

    # Variables and domain
    d1 = [0, 1]
    d2 = [0, 1, 2]
    x1 = Variable('x1', d1)
    x2 = Variable('x2', d2)

    # Cost functions for x1 and x2
    x1_cost = UnaryFunctionRelation('x1_cost', x1,
                                    lambda v: {0: 0, 1: -3}[v])
    x2_cost = UnaryFunctionRelation('x2_cost', x2,
                                    lambda v: {0: 0, 1: -2, 2: -1}[v])

    # Constraint x1 != x2
    # Without any cost
    @relations.AsNAryFunctionRelation(x1, x2)
    def all_diff(x1_val, x2_val):
        if x1_val == x2_val:
            return 10000
        return 0

    # Map the factor graph to agents
    variables = [x1, x2]
    factors = [x1_cost, x2_cost, all_diff]
    node_agents= distribue_agent_for_all(variables, factors)

    # and solve it
    results, _, _ = synchronous_single_run(node_agents)

    print(results)
    if results['x1'] == 1 and results['x2'] == 2:
        logging.info('SUCCESS !! ')
        return 0
    else:
        logging.info('invalid result found, needs some debugging ...' + str(
            results))
        return 1

def run_test():
    return graph_coloring_pb()

if __name__ == "__main__":
    sys.exit(run_test())


#https://www.acmicpc.net/problem/16932

import sys
from collections import deque

n, m = map(int,sys.stdin.readline().split())
a = [list(map(int,sys.stdin.readline().split())) for _ in range(n)]
MIN = 0
answer = MIN

dx = [0,0,1,-1]
dy = [1,-1,0,0]

group = [[0] * m for _ in range(n)]
group_number = 0
group_size = [0]

def bfs(i,j):
    global group_number
    group_number += 1
    q = deque()
    q.append((i,j))
    group[i][j] = group_number
    cnt = 1
    while q:
        x, y = q.popleft()
        for k in range(4):
            nx = x + dx[k]
            ny = y + dy[k]
            if 0 <= nx < n and 0 <= ny < m and group[nx][ny] == 0 and a[nx][ny] == 1:
                group[nx][ny] = group_number
                q.append((nx,ny))
                cnt += 1

    group_size.append(cnt)

for i in range(n):
    for j in range(m):
        if a[i][j] == 1 and group[i][j] == 0:
            bfs(i,j)

for i in range(n):
    for j in range(m):
        if a[i][j] == 0:
            d = set()
            for k in range(4):
                nx = i + dx[k]
                ny = j + dy[k]
                if 0 <= nx < n and 0 <= ny < m and a[nx][ny] == 1:
                    d.add(group[nx][ny])
            total = 1
            for near in d:
                total += group_size[near]
            answer = max(total,answer)

print(answer)

BOJ(Baekjoon Online Judge)/Graph/16932_모양 만들기(shaping).py


import sys
import numpy as np
import midi
from PIL import Image
import matplotlib.pyplot as plt

input_path = str(sys.argv[1])
output_path = str(sys.argv[2])

piano_roll = np.array(Image.open(input_path))
piano_roll = np.transpose(np.flip(piano_roll, 0), (1,0))
piano_roll = np.lib.pad(piano_roll, ((0,1),(0,1)), 'constant', constant_values=((0,0),(0,0)))
# plt.imshow(piano_roll, cmap='gray')
piano_roll[piano_roll<130] = 0
piano_roll[piano_roll>0] = 255

pattern = midi.Pattern(resolution=4)
track = midi.Track()
pattern.append(track)
lastEvent = 0
sweep = np.zeros_like(piano_roll[0], dtype=np.bool)
for tick, line in enumerate(piano_roll):
    for i, note in enumerate(line):
        if note>0 and not sweep[i]: # a note on event
            track.append(midi.NoteOnEvent(tick=tick-lastEvent, data=[i+34 ,127]))
            lastEvent = tick
            sweep[i]=True
        if note==0 and sweep[i]: # a note off event
            track.append(midi.NoteOffEvent(tick=tick-lastEvent, data=[i+34 ,0]))
            lastEvent = tick
            sweep[i]=False
track.append(midi.EndOfTrackEvent(tick=1))
midi.write_midifile(output_path, pattern)


import argparse
import collections
import logging
import math
import os
import re
import time
from concurrent.futures import ProcessPoolExecutor

import numpy as np
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
import torch
import torch.nn as nn
import torch.utils.data
import transformers
from data_pipeline import create_pretraining_dataloader
from learning_rates import AnnealingLR
from memory_tracker import memory_status, memory_status_cpu
from sharded_data_parallel_checkpoint import get_buffer_names, get_param_shapes
from smdistributed.modelparallel.torch.nn import FusedLayerNorm as LayerNorm
from smdistributed.modelparallel.torch.nn.huggingface.gptj import (
    translate_hf_gptj_state_dict_to_smdistributed,
    translate_state_dict_to_hf_gptj,
)
from torch import optim
from torch.nn.parallel.distributed import DistributedDataParallel
from transformers import (
    CONFIG_MAPPING,
    MODEL_FOR_CAUSAL_LM_MAPPING,
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GPTJConfig,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import is_main_process

logging.getLogger("torch.distributed.distributed_c10d").setLevel(logging.ERROR)
logger = logging.getLogger(__name__)

def get_learning_rate_scheduler(optimizer, args):
    # Add linear learning rate scheduler.
    if args.lr_decay_iters is not None:
        num_iters = args.lr_decay_iters
    else:
        num_iters = args.max_steps
    num_iters = max(1, num_iters)
    init_step = 0
    warmup_iter = args.warmup * num_iters
    plateau_iter = warmup_iter + args.plateau * num_iters
    lr_scheduler = AnnealingLR(
        optimizer,
        start_lr=args.lr,
        warmup_iter=warmup_iter,
        plateau_iter=plateau_iter,
        total_iters=num_iters,
        decay_style=args.lr_decay_style,
        last_iter=init_step,
        min_lr=args.min_lr,
        use_checkpoint_lr_scheduler=args.load_partial or args.load_full,
        override_lr_scheduler=False,
    )

    return lr_scheduler

def get_param_groups_by_weight_decay(module):
    weight_decay_params = {"params": []}
    no_weight_decay_params = {"params": [], "weight_decay": 0.0}
    param_ids = set()
    for module_ in module.modules():
        if isinstance(module_, LayerNorm):
            for p in list(module_._parameters.values()):
                if p is not None and id(p) not in param_ids:
                    no_weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
        else:
            for n, p in list(module_._parameters.items()):
                if p is not None and n != "bias" and id(p) not in param_ids:
                    weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
            for n, p in list(module_._parameters.items()):
                if p is not None and n == "bias" and id(p) not in param_ids:
                    no_weight_decay_params["params"].append(p)
                    param_ids.add(id(p))
    return weight_decay_params, no_weight_decay_params

# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    if args.logits_output:
        output = model(
            input_ids=input_ids, attention_mask=attention_mask, labels=input_ids
        )
        loss = output["loss"]
    else:
        loss = model(
            input_ids=input_ids, attention_mask=attention_mask, labels=input_ids
        )["loss"]

    model.backward(loss)

    if args.logits_output:
        return output

    return loss

# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)[
        "loss"
    ]
    return loss

def eval_model(model, dataloader, num_batches, use_bert_data):
    model = model.eval()
    n_batches = 0
    loss = 0.0

    with torch.no_grad():
        for batch_idx, input_data in enumerate(dataloader):
            if use_bert_data:
                input_ids, _, attention_mask, _, _ = input_data
            else:
                input_ids, attention_mask = input_data
            if batch_idx >= num_batches:
                break

            loss += test_step(model, input_ids, attention_mask).reduce_mean()
            n_batches += 1

    if n_batches > 0:
        torch.distributed.all_reduce(loss, group=smp.get_dp_process_group())
        loss /= smp.dp_size()
        loss /= n_batches
        loss = loss.item()
        ppl = math.exp(loss)
    else:
        loss = -1.0
        ppl = -1.0

    return loss, ppl

def train(
    model,
    optimizer,
    lr_scheduler,
    model_config,
    start_train_path_index,
    start_batch_index,
    num_params,
    total_steps,
    args,
):
    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before train step")
    model.train()
    if args.parallel_proc_data_processing:
        pool = ProcessPoolExecutor(1)

    dp_rank = smp.dp_rank() if not args.prescaled_batch else smp.rdp_rank()
    dp_size = smp.dp_size() if not args.prescaled_batch else smp.rdp_size()
    data_type = "BERT" if args.use_bert_data else "GPT"

    if args.use_bert_data:
        train_paths = sorted(
            [
                os.path.join(args.training_dir, p)
                for p in os.listdir(args.training_dir)
                if os.path.isfile(os.path.join(args.training_dir, p))
                and "training" in p
            ]
        )
    else:
        if args.zipped_data > 0:
            file_extension = ".json.gz"
        else:
            file_extension = ".json"
        train_paths = sorted(
            [
                os.path.join(args.training_dir, p)
                for p in os.listdir(args.training_dir)
                if p.endswith(file_extension)
            ]
        )

    train_dataloader = create_pretraining_dataloader(
        [train_paths[start_train_path_index]],
        args.train_batch_size,
        args.max_context_width,
        seed=args.seed,
        dp_rank=dp_rank,
        dp_size=dp_size,
        shuffle=args.same_seed < 1,
        zipped=args.zipped_data > 0,
        use_last_file_only=args.fast_validation > 0,
        data_type=data_type,
    )

    if args.validation_freq is not None:
        # load all validation examples
        if smp.rank() == 0:
            print("Creating val dataloader")
        if args.use_bert_data:
            val_paths = sorted(
                [
                    os.path.join(args.test_dir, p)
                    for p in os.listdir(args.test_dir)
                    if os.path.isfile(os.path.join(args.test_dir, p)) and "testing" in p
                ]
            )

        else:
            if args.zipped_data > 0:
                file_extension = ".json.gz"
            else:
                file_extension = ".json"
            val_paths = sorted(
                [
                    os.path.join(args.test_dir, p)
                    for p in os.listdir(args.test_dir)
                    if p.endswith(file_extension)
                ]
            )
        val_dataloader = create_pretraining_dataloader(
            val_paths,
            args.val_batch_size,
            args.max_context_width,
            seed=args.seed,
            dp_rank=dp_rank,
            dp_size=dp_size,
            shuffle=True,
            zipped=args.zipped_data > 0,
            use_last_file_only=args.fast_validation > 0,
            data_type=data_type,
        )
        if smp.rank() == 0:
            print("Created val dataloader")

    start = time.time()
    throughput = None
    to_save = {"loss": [], "val_loss": []}
    loss_metric = 0

    def grad_accumulation_boundary(batch_idx):
        return batch_idx % args.gradient_accumulation == args.gradient_accumulation - 1

    def should_record():
        # only record the ranks that in the tp group that contains global rank 0
        if smp.tp_size() > 1:
            tp_group = smp.get_tp_group()
            return 0 in tp_group
        else:
            return smp.rank() == 0

    # Set the same seed for computation
    set_seed(args.seed)

    for index in range(start_train_path_index, args.epochs * len(train_paths)):
        next_train_path_index = (index + 1) % len(train_paths)
        curr_train_path_index = index % len(train_paths)

        if total_steps >= args.max_steps:
            break

        if args.parallel_proc_data_processing:
            dataset_future = pool.submit(
                create_pretraining_dataloader,
                [train_paths[next_train_path_index]],
                args.train_batch_size,
                args.max_context_width,
                seed=args.seed,
                dp_rank=dp_rank,
                dp_size=dp_size,
                shuffle=args.same_seed < 1,
                zipped=args.zipped_data > 0,
                use_last_file_only=args.fast_validation > 0,
                data_type=data_type,
            )

        if smp.rank() == 0:
            if args.use_bert_data:
                print(
                    f"Reading data from training path {train_dataloader.dataset.input_file}"
                )
            else:
                print(
                    f"Reading data from training path {train_dataloader.dataset.input_paths}"
                )

        for batch_idx, input_data in enumerate(train_dataloader):
            if batch_idx < start_batch_index:
                if smp.rank() == 0:
                    print(
                        f"Resuming from saved batch index {start_batch_index}, skipping batch {batch_idx}..."
                    )
                if start_batch_index == len(train_dataloader):
                    # If saving at the last batch of the file, read from the next file
                    start_batch_index = 0
                    break
                continue
            else:
                start_batch_index = 0

            if args.use_bert_data:
                input_ids, _, attention_mask, _, _ = input_data
            else:
                input_ids, attention_mask = input_data

            if total_steps >= args.max_steps:
                break

            step_start = time.time()

            if grad_accumulation_boundary(batch_idx - 1):
                optimizer.zero_grad(set_to_none=True)

            if args.logits_output:
                train_output = train_step(
                    model, optimizer, input_ids, attention_mask, args
                )
                loss_mb = train_output["loss"]
                logits_mb = train_output["logits"]
                if smp.tp_size() > 1:
                    logits = torch.cat(tuple(logits_mb.outputs), dim=1)
                else:
                    logits = torch.cat(tuple(logits_mb.outputs), dim=0)
            else:
                # Return value, loss_mb is a StepOutput object
                loss_mb = train_step(model, optimizer, input_ids, attention_mask, args)

            # smdistributed: Average the loss across microbatches.
            loss = loss_mb.reduce_mean()
            if not args.validation_freq:
                loss_metric = loss.item()

            if args.enable_memory_profiling > 0:
                memory_status_cpu("After_train_step_cpu")
                memory_status(msg="After_train_step")

            if args.clean_cache > 0:
                # empty the cache to avoid OOM
                torch.cuda.empty_cache()

            if grad_accumulation_boundary(batch_idx):
                if args.fp16:
                    optimizer.clip_master_grads(args.grad_clip)

                optimizer.step()
                if not (args.fp16 and optimizer.overflow):
                    lr_scheduler.step()

                if args.enable_memory_profiling > 0:
                    memory_status(msg="After_opt_step")

            total_steps += 1
            time_elapsed = time.time() - start
            step_time = time.time() - step_start
            sample_processed = input_ids.shape[0] * dp_size
            throughput = sample_processed / step_time
            tokens_per_gpu = input_ids.shape[0] * input_ids.shape[1]

            # Based on the formula in https://developer.nvidia.com/blog/scaling-language-model-training-to-a-trillion-parameters-using-megatron/
            tflops_per_gpu = 8 * num_params * tokens_per_gpu / step_time / 1e12
            if smp.rank() == 0 and not total_steps % args.logging_freq:
                print(
                    f"({int(time_elapsed)}s), Batch {total_steps - 1} Loss: {loss.item()}, Speed: {throughput} samples/sec, TFLOPS/GPU: {tflops_per_gpu}"
                )

            # evaluate on validation
            if args.validation_freq and not (total_steps % args.validation_freq):
                cur_state = np.random.get_state()
                model = model.eval()
                val_loss, val_ppl = eval_model(
                    model, val_dataloader, args.validation_batches, args.use_bert_data
                )
                if is_main_process(smp.rank()):
                    print(
                        f"({int(time.time()-start)}s) Batch {total_steps - 1} Validation loss: {val_loss}"
                    )
                    print(
                        f"({int(time.time()-start)}s) Batch {total_steps - 1} Validation perplexity: {val_ppl}"
                    )
                loss_metric = val_loss
                if args.logits_output:
                    to_save["val_loss"].append(val_loss)
                model = model.train()
                if args.preserve_np_state > 0:
                    np.random.set_state(cur_state)

            # checkpoint
            if not (total_steps % args.checkpoint_freq):
                user_content = {
                    "cli_args": args.__dict__,
                    "num_params": num_params,
                    "total_steps": total_steps,
                    "start_train_path_index": curr_train_path_index,
                    "model_config": model_config,
                    "start_batch_index": batch_idx + 1,
                }
                # to reconstruct the full model
                if args.sharded_data_parallel_degree > 1:
                    user_content["buffer_names"] = get_buffer_names(model)
                    user_content["param_shapes"] = get_param_shapes(model, optimizer)
                user_content["lr_scheduler"] = lr_scheduler.state_dict()
                smp.save_checkpoint(
                    args.checkpoint_dir,
                    tag=f"total_steps{total_steps}",
                    partial=True,
                    model=model,
                    optimizer=optimizer,
                    user_content=user_content,
                    num_kept_partial_checkpoints=args.num_kept_checkpoints,
                )

            if args.logits_output:
                to_save["loss"].append(loss.item())

        if total_steps >= args.max_steps:
            if should_record() and args.logits_output:
                to_save["logits"] = logits.detach().cpu()
                output_file = f"rank_{smp.rank()}_" + args.logits_output
                torch.save(to_save, os.path.join(args.model_dir, output_file))
                print(
                    f"logits and loss saved at {os.path.join(args.model_dir, output_file)}"
                )
            break

        del train_dataloader

        if args.parallel_proc_data_processing:
            s = time.time()
            train_dataloader = dataset_future.result(timeout=None)
            wait_time = time.time() - s
            if wait_time > 1:
                # TODO if this happens, we should try num_workers>1 in dataloader
                print(
                    f"[{smp.rank()}] Waited {wait_time} for data loader to be ready. Please check if dataloader performance can be improved to avoid these waits."
                )
        else:
            train_dataloader = create_pretraining_dataloader(
                [train_paths[next_train_path_index]],
                args.train_batch_size,
                args.max_context_width,
                seed=args.seed,
                dp_rank=dp_rank,
                dp_size=dp_size,
                shuffle=args.same_seed < 1,
                zipped=args.zipped_data > 0,
                use_last_file_only=args.fast_validation > 0,
                data_type=data_type,
            )

    return total_steps, throughput, loss_metric

def parse_args():
    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.

    opt_grp = parser.add_argument_group(
        title="optimization", description="arguments for optimization"
    )
    opt_grp.add_argument(
        "--train_batch_size",
        type=int,
        default=4,
        help="batch size per dp rank, for tensor parallelism degree 8 with pipeline parallel degree 1 this means 8*this batch size per node",
    )
    opt_grp.add_argument("--val_batch_size", type=int, default=4)
    opt_grp.add_argument("--max_steps", type=int, default=100)
    opt_grp.add_argument("--seed", type=int, default=12345)
    opt_grp.add_argument("--same_seed", type=int, default=0)
    opt_grp.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    opt_grp.add_argument(
        "--fp16", default=0, type=int, help="automatic mixed precision training"
    )
    opt_grp.add_argument(
        "--bf16", default=0, type=int, help="automatic mixed precision training"
    )
    opt_grp.add_argument("--sharded_data_parallel_degree", default=1, type=int)
    opt_grp.add_argument(
        "--grad_clip", default=1.0, type=float, help="gradient clipping"
    )
    opt_grp.add_argument(
        "--weight_decay", default=0.01, type=float, help="weight decay"
    )
    opt_grp.add_argument(
        "--beta1", default=0.9, type=float, help="beta1 parameter for Adam optimizer"
    )
    opt_grp.add_argument(
        "--beta2", default=0.95, type=float, help="beta2 parameter for Adam optimizer"
    )
    opt_grp.add_argument(
        "--activation_checkpointing",
        type=int,
        default=1,
        help="enable gradient checkpointing to reduce memory consumption",
    )
    parser.add_argument(
        "--logging_freq",
        type=int,
        default=1,
        help="number of iterations between logging",
    )

    # I/O
    io_grp = parser.add_argument_group(
        title="io", description="location for input and output"
    )
    io_grp.add_argument(
        "--use_bert_data", type=int, default=0, help="use wiki corpus data for training"
    )
    io_grp.add_argument(
        "--zipped_data", type=int, default=0, help="input data is zipped files"
    )
    io_grp.add_argument(
        "--epochs",
        type=int,
        default=1,
        help="times of iterating over the training dataset",
    )
    io_grp.add_argument(
        "--output-data-dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"]
    )
    io_grp.add_argument(
        "--checkpoint-dir",
        type=str,
        default="/opt/ml/checkpoints",
        help="Saves partial checkpoints (model, optimizer) to this dir, and loads latest checkpoint from this if load_partial is specified.",
    )
    io_grp.add_argument(
        "--model-dir",
        type=str,
        default=os.environ["SM_MODEL_DIR"],
        help="Saves full model for inference to this dir. Also used if load_full is given to load the model. Note the lack of optimizer state here.",
    )
    io_grp.add_argument(
        "--training-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"]
    )
    io_grp.add_argument("--test-dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
    io_grp.add_argument(
        "--parallel_proc_data_processing",
        type=int,
        default=0,
        help="Load data in parallel with a different process. At any point a process can have two files in memory. With tensor parallelism, each of the 8 processes on an instance will then have 2 files in memory. Depending on file sizes this may or may not be feasible. With pipeline parallelism this was not a problem as only 1 rank on an instance loaded data.",
    )
    io_grp.add_argument(
        "--save_final_full_model",
        type=int,
        default=0,
        help="Enabling this will save a combined model only at the end",
    )
    io_grp.add_argument(
        "--load_partial", type=int, default=0, help="Load from partial checkpoints"
    )
    io_grp.add_argument(
        "--load_full", type=int, default=0, help="Load from full checkpoints"
    )
    io_grp.add_argument(
        "--logits_output", type=str, default="", help="Path to save logits and loss"
    )
    io_grp.add_argument(
        "--prescaled_batch", type=int, default=1, help="use prescaled batch"
    )

    # configure model size
    model_grp = parser.add_argument_group(
        title="model", description="arguments to describe model configuration"
    )
    model_grp.add_argument("--max_context_width", type=int, default=1024)
    model_grp.add_argument("--vocab_size", type=int, default=50400)
    model_grp.add_argument("--hidden_width", type=int, default=768)
    model_grp.add_argument("--num_layers", type=int, default=12)
    model_grp.add_argument("--num_heads", type=int, default=12)
    model_grp.add_argument("--resid_pdrop", type=float, default=0.1)
    model_grp.add_argument("--embd_pdrop", type=float, default=0.1)
    model_grp.add_argument("--attn_pdrop", type=float, default=0.1)
    model_grp.add_argument("--summary_first_pdrop", type=float, default=0.1)
    model_grp.add_argument(
        "--use_adamw", type=int, default=0, help="Use adamw optimizer"
    )
    model_grp.add_argument(
        "--finetune_6b",
        type=int,
        default=0,
        help="Flag to enable finetune 6B GPTJ model",
    )
    model_grp.add_argument(
        "--use_distributed_transformer",
        type=int,
        default=1,
        help="Use distributed transformer",
    )
    model_grp.add_argument(
        "--checkpoint_sublayers",
        type=int,
        default=0,
        help="Apply activation checkpointing to submodules of each transformer layer",
    )

    smp_grp = parser.add_argument_group(title="smp", description="smp")
    smp_grp.add_argument("--tensor_parallel_degree", type=int, default=1)
    smp_grp.add_argument("--pipeline_parallel_degree", type=int, default=1)
    smp_grp.add_argument("--microbatches", type=int, default=1)
    smp_grp.add_argument("--active_microbatches", type=int, default=None)
    smp_grp.add_argument("--optimize", type=str, default="speed")
    smp_grp.add_argument("--activation_strategy", type=str, default="each")
    smp_grp.add_argument("--shard_optimizer_state", type=int, default=0)
    smp_grp.add_argument("--offload_activations", type=int, default=0)
    smp_grp.add_argument("--fast_mode", type=int, default=0)
    smp_grp.add_argument("--static_mode", type=int, default=0)
    smp_grp.add_argument("--delayed_param", type=int, default=0)
    smp_grp.add_argument("--same_partition_load", type=int, default=0)
    smp_grp.add_argument("--attention_in_fp32", type=int, default=0)
    smp_grp.add_argument("--placement_strategy", type=str, default="cluster")
    smp_grp.add_argument("--activation_loading_horizon", type=int, default=4)
    smp_grp.add_argument("--skip_tracing", type=int, default=0)
    smp_grp.add_argument("--query_key_layer_scaling", type=int, default=1)
    smp_grp.add_argument("--fused_softmax", type=int, default=1)
    smp_grp.add_argument("--fused_dropout", type=int, default=0)
    smp_grp.add_argument("--fused_bias_gelu", type=int, default=1)
    smp_grp.add_argument("--gradient_accumulation", type=int, default=1)

    parser.add_argument(
        "--num_kept_checkpoints",
        type=int,
        default=5,
        help="how many checkpoints to keep before deleting",
    )
    parser.add_argument(
        "--checkpoint_freq",
        type=int,
        default=10000,
        help="number of iterations between checkpointing",
    )
    parser.add_argument(
        "--validation_freq",
        type=int,
        default=None,
        help="number of iterations to print validation loss",
    )
    parser.add_argument(
        "--validation_batches",
        type=int,
        default=10,
        help="number of batches to estimate validation loss",
    )
    parser.add_argument(
        "--manual_partition",
        type=int,
        default=0,
        help="evenly distribute layers across the partitions",
    )
    parser.add_argument(
        "--partition_assignment",
        type=str,
        default="",
        help="number of transformer layers assigned to each partition",
    )
    parser.add_argument(
        "--preserve_np_state",
        type=int,
        default=0,
        help="Perserve the numpy random state between validation",
    )
    parser.add_argument(
        "--fast_validation",
        type=int,
        default=1,
        help="Running validation only with the last data file for faster speed",
    )
    parser.add_argument(
        "--gather_if_shard",
        type=int,
        default=1,
        help="When sharding opt states is enabled, gather the opt checkpoint to rdp rank 0 during saving",
    )
    parser.add_argument(
        "--clean_cache",
        type=int,
        default=0,
        help="Clean torch reserved memory at he end of every step",
    )
    parser.add_argument(
        "--use_fsx", type=int, default=0, help="Using FSx for checkpointing"
    )
    parser.add_argument(
        "--enable_memory_profiling", type=int, default=0, help="Enable memory profile"
    )

    # learning rate
    lr_grp = parser.add_argument_group(
        title="lr", description="arguments for learning rate schedule"
    )
    lr_grp.add_argument("--lr", type=float, default=None, help="Initial learning rate.")
    lr_grp.add_argument(
        "--lr_decay_style",
        type=str,
        default="linear",
        choices=["constant", "linear", "cosine", "exponential", "plateau"],
        help="Learning rate decay function.",
    )
    lr_grp.add_argument(
        "--lr_decay_iters",
        type=int,
        default=None,
        help="number of iterations to decay learning rate over,"
        " If None defaults to train iters",
    )
    lr_grp.add_argument(
        "--min_lr",
        type=float,
        default=0.0,
        help="Minumum value for learning rate. The scheduler"
        "clip values below this threshold.",
    )
    lr_grp.add_argument(
        "--warmup",
        type=float,
        default=0.01,
        help="Percentage of total iterations to warmup on "
        "(.01 = 1 percent of all training iters).",
    )
    lr_grp.add_argument(
        "--plateau",
        type=float,
        default=0.4,
        help="Percentage of total iterations to keep at max if using plateau lr",
    )

    ci_grp = parser.add_argument_group(title="ci", description="ci related settings")
    ci_grp.add_argument(
        "--ci", default=False, action="store_true", help="Whether enable ci"
    )
    ci_grp.add_argument("--time_to_train", type=int, help="time to train threshold")
    ci_grp.add_argument("--throughput", type=float, help="throughput threshold")
    ci_grp.add_argument("--loss", type=float, help="loss threshold")
    args, _ = parser.parse_known_args()
    return args

def compute_num_params(model):
    num_params = 0
    seen = set()
    for p in model.parameters():
        if p not in seen:
            seen.add(p)
            if hasattr(p, "ds_shape"):
                num_params += np.prod(p.ds_shape)
            else:
                num_params += np.prod(p.size())

    return num_params

def main():
    args = parse_args()

    if args.partition_assignment != "" and args.manual_partition == 0:
        print("[Warning] partition_assignment is set, enable manual_partition")
        args.manual_partition = 1

    # any value here is overriden by the config set in notebook when launching the sagemaker job
    smp_config = {
        "ddp": True,
        "tensor_parallel_degree": args.tensor_parallel_degree,
        "pipeline_parallel_degree": args.pipeline_parallel_degree,
        "microbatches": args.microbatches,
        "shard_optimizer_state": args.shard_optimizer_state > 0,
        "prescaled_batch": args.prescaled_batch > 0,
        "fp16": args.fp16 > 0,
        "bf16": args.bf16 > 0,
        "offload_activations": args.offload_activations > 0,
        "delayed_parameter_initialization": args.delayed_param > 0,
        "optimize": args.optimize,
        "placement_strategy": args.placement_strategy,
        "activation_loading_horizon": args.activation_loading_horizon,
        "skip_tracing": args.skip_tracing > 0,
        "auto_partition": False if args.manual_partition else True,
        "default_partition": 0,
        "static_mode": args.static_mode > 0,
        "fast_mode": args.fast_mode > 0,
        "sharded_data_parallel_degree": args.sharded_data_parallel_degree,
    }
    if args.active_microbatches is not None:
        smp_config["active_microbatches"] = args.active_microbatches

    smp.init(smp_config)

    if smp.rank() == 0:
        print("Arguments:", args.__dict__)
        print(f"Transformers version: {transformers.__version__}")
        print(
            f"smdistributed.modelparallel version: {smdistributed.modelparallel.__version__}"
        )
        print(f"smdistributed config: {smp_config}")

    if args.save_final_full_model and smp.rank() == 0:
        print(
            f"[Warning] Note that save_final_full_model only saves the final model at the end of all steps. It does not save optimizer state. Optimizer state is only saved with partial models which are saved at checkpointing_freq during training. If you want to restart training you need partial checkpoints."
        )

    if args.partition_assignment != "":
        partition_assignment = args.partition_assignment.split(",")
        assert (
            len(partition_assignment) == smp.pp_size()
        ), f"partition_assignment must have the same size as pipeline parallel degree, but getting {len(partition_assignment)} vs {smp.pp_size()}"

    model_config = GPTJConfig(
        vocab_size=args.vocab_size,
        n_positions=args.max_context_width,
        n_embd=args.hidden_width,
        n_layer=args.num_layers,
        n_head=args.num_heads,
        n_inner=None,
        activation_function="gelu_new",
        resid_pdrop=args.resid_pdrop,
        embd_pdrop=args.embd_pdrop,
        attn_pdrop=args.attn_pdrop,
        layer_norm_epsilon=1e-05,
        initializer_range=0.02,
        summary_type="cls_index",
        summary_use_proj=True,
        summary_activation=None,
        summary_proj_to_labels=True,
        summary_first_dropout=args.summary_first_pdrop,
        # gradient_checkpointing=args.gradient_checkpointing > 0,
        use_cache=False,
        bos_token_id=50256,
        eos_token_id=50256,
        return_dict=True,
    )

    # the following improves start-up time by skipping proper initialization
    # of weights in the original model. this is not a problem because DistributedModel
    # will override those weights anyway when we use distributed transformer.
    if args.use_distributed_transformer > 0:
        from transformers.modeling_utils import PreTrainedModel

        PreTrainedModel.init_weights = lambda x: None

    set_seed(args.seed)

    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before model creation")

    if args.fp16 and args.bf16:
        raise ValueError("FP16 and BF16 cannot be simultaneously enabled.")
    elif args.fp16:
        dtype = torch.float16
    elif args.bf16:
        dtype = torch.bfloat16
    else:
        dtype = torch.get_default_dtype()

    if args.finetune_6b:
        with smp.model_creation(
            tensor_parallelism=smp.tp_size() > 1
            or args.use_distributed_transformer > 0,
            dtype=dtype,
            attention_in_fp32=args.attention_in_fp32 > 0,
            query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
            fused_softmax=args.fused_softmax > 0,
            fused_dropout=args.fused_dropout > 0,
            fused_bias_gelu=args.fused_bias_gelu > 0,
        ):
            model = AutoModelForCausalLM.from_pretrained(
                "EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16
            )
            model_config = model.config
    else:
        with smp.model_creation(
            tensor_parallelism=smp.tp_size() > 1
            or args.use_distributed_transformer > 0,
            dtype=dtype,
            attention_in_fp32=args.attention_in_fp32 > 0,
            query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
            fused_softmax=args.fused_softmax > 0,
            fused_dropout=args.fused_dropout > 0,
            fused_bias_gelu=args.fused_bias_gelu > 0,
        ):
            model = AutoModelForCausalLM.from_config(
                model_config, torch_dtype=torch.float16
            )
            model_config = model.config

    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="after model creation")

    num_params = compute_num_params(model)
    if smp.rank() == 0:
        print(f"# total parameters: {num_params}")

    # smdistributed: Set the device to the GPU ID used by the current process.
    # Input tensors should be transferred to this device.
    torch.cuda.set_device(smp.local_rank())
    device = torch.device("cuda")

    if not args.same_seed:
        # Set seed by tp_rank to prevent weights from being the same on different tp_ranks
        set_seed(args.seed + smp.tp_rank())

    # smdistributed: Use the DistributedModel container to provide the model
    # to be partitioned across different ranks. For the rest of the script,
    # the returned DistributedModel object should be used in place of
    # the model provided for DistributedModel class instantiation.
    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="before dist model creation")
    model = smp.DistributedModel(
        model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation
    )

    if args.enable_memory_profiling > 0:
        memory_status_cpu(msg="after dist model creation")

    m = model.get_module()
    if args.use_distributed_transformer > 0:
        transformer_layers = m.transformer.seq_layers
    else:
        transformer_layers = m.transformer.h

    if args.manual_partition:
        print(f"Manual partition enabled")
        if args.partition_assignment != "":
            get_num_layers = lambda x: int(partition_assignment[x])
            total_layers = sum(
                [get_num_layers(pp_rank) for pp_rank in range(smp.pp_size())]
            )
            assert (
                total_layers == args.num_layers
            ), f"partition_assignment must have the same total transformer layers as model, but getting {total_layers} vs {args.num_layers}"
        else:
            # evenly distribute layers across all partitions
            div, rem = divmod(args.num_layers, smp.pp_size())
            get_num_layers = lambda x: (div + 1 if x >= smp.pp_size() - rem else div)
        assignments = []
        for pp_rank in range(smp.pp_size()):
            nl = get_num_layers(pp_rank)
            print(f"{nl} layers assigned to partition {pp_rank}")
            assignments += [pp_rank for _ in range(nl)]

        for i, c in enumerate(transformer_layers.children()):
            smp.set_partition(c, assignments[i])

    param_groups = get_param_groups_by_weight_decay(m)

    if args.use_adamw > 0:
        optimizer = optim.AdamW(
            param_groups,
            betas=(args.beta1, args.beta2),
            lr=args.lr,
            weight_decay=args.weight_decay,
        )
    else:
        optimizer = optim.Adam(
            param_groups,
            betas=(args.beta1, args.beta2),
            lr=args.lr,
            weight_decay=args.weight_decay,
        )

    if args.activation_checkpointing:
        if args.use_distributed_transformer or smp.tp_size() > 1:
            if args.checkpoint_sublayers:
                for c in transformer_layers.children():
                    smp.set_activation_checkpointing(c.attention)
                    smp.set_activation_checkpointing(c.output)
            else:
                smp.set_activation_checkpointing(
                    transformer_layers, strategy=args.activation_strategy
                )
        else:
            for c in transformer_layers.children():
                if args.checkpoint_sublayers:
                    smp.set_activation_checkpointing(c.attn)
                    smp.set_activation_checkpointing(c.mlp)
                else:
                    smp.set_activation_checkpointing(c)

    optimizer = smp.DistributedOptimizer(
        optimizer,
        static_loss_scale=None,
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
    )
    lr_scheduler = get_learning_rate_scheduler(optimizer, args)

    if args.enable_memory_profiling > 0:
        model.register_post_partition_hook(
            lambda model, optimizer: memory_status(msg="After_partition")
        )

    # load after wrapping model and optimizer with smp Distributed...
    if args.load_full or args.load_partial:
        if args.load_partial and args.load_full:
            print(
                "Since both --load_partial and --load_full set, will try to load from full checkpoint."
                "If the intention is to load from partial checkpoint, please don't set --load_full"
            )
        partial = not args.load_full
        path = args.checkpoint_dir if partial else args.model_dir
        tag = None if partial else "fullmodel.pt"
        user_content = smp.resume_from_checkpoint(path, tag=tag, partial=partial)
        total_steps = user_content["total_steps"] if partial else 0
        start_train_path_index = user_content.get("start_train_path_index", 0)
        start_batch_index = user_content.get("start_batch_index", 0)
        if "lr_scheduler" in user_content:
            lr_scheduler.load_state_dict(user_content["lr_scheduler"])
    else:
        total_steps = 0
        start_train_path_index = 0
        start_batch_index = 0

    start = time.time()
    total_steps, throughput, loss = train(
        model,
        optimizer,
        lr_scheduler,
        model_config,
        start_train_path_index,
        start_batch_index,
        num_params,
        total_steps,
        args,
    )
    time_to_train = time.time() - start
    if args.ci:
        print(f"[SMP_METRIC]__GPTJ__Time_to_train__{time_to_train}")
        print(f"[SMP_METRIC]__GPTJ__samples/second__{throughput}")
        print(f"[SMP_METRIC]__GPTJ__Loss__{loss}")
        if not args.load_partial and not args.load_full:
            assert time_to_train < args.time_to_train
            assert throughput > args.throughput
            if args.loss:
                assert loss < args.loss

    if args.save_final_full_model:
        # saves full model at the end
        user_content = {
            "cli_args": args.__dict__,
            "num_params": num_params,
            "total_steps": total_steps,
            "model_config": model_config,
        }
        if args.sharded_data_parallel_degree > 1:
            # When sharded_data_parallel_degree > 1, saving full model is not supported, saving partial instead
            # To get the full model, one can use the following API
            # > from sharded_data_parallel_checkpoint import get_full_state_dict_from_sharded_data_parallel_checkpoint
            # > full_model = get_full_state_dict_from_sharded_data_parallel_checkpoint(args.model_dir, tag=f"sharded_data_parallel_final_full_{num_params}", dtype=torch.float32)
            # > if args.use_distributed_transformer > 0: # translate the state_dict to hf format if distributed transformer is used
            # >     full_model = smp.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(full_model, max_seq_len=args.max_context_width)
            # Note: the shared parameter will not be reflected so during loading you might need to load with strict=False
            user_content["buffer_names"] = get_buffer_names(model)
            user_content["param_shapes"] = get_param_shapes(model, optimizer)
            smp.save_checkpoint(
                args.model_dir,
                tag=f"sharded_data_parallel_final_full_{num_params}",
                partial=True,
                model=model,
                optimizer=optimizer,
                user_content=user_content,
            )
        else:
            smp.save_checkpoint(
                args.model_dir,
                tag="fullmodel.pt",
                partial=False,
                model=model,
                user_content=user_content,
            )

    smp.barrier()
    if smp.rank() == 0:
        print("SMP training finished successfully")

if __name__ == "__main__":
    main()